feat: Enable time range rolling for DataFrame, DataFrameGroupBy and SeriesGroupBy (#1605)

* feat: Enable range rolling for DataFrame, DataFrameGroupBy and SeriesGroupBy * update docs * fix lint * resolve merge conflict

Author: sycai

bigframes/core/groupby/dataframe_group_by.py

@@ -14,11 +14,13 @@
 
 from __future__ import annotations
 
+import datetime
 import typing
 from typing import Literal, Sequence, Tuple, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.groupby as vendored_pandas_groupby
+import numpy
 import pandas as pd
 
 from bigframes import session
@@ -30,6 +32,7 @@
 import bigframes.core.ordering as order
 import bigframes.core.utils as utils
 import bigframes.core.validations as validations
+from bigframes.core.window import rolling
 import bigframes.core.window as windows
 import bigframes.core.window_spec as window_specs
 import bigframes.dataframe as df
@@ -309,28 +312,41 @@ def diff(self, periods=1) -> series.Series:
  @validations.requires_ordering()
  def rolling(
  self,
- window: int,
+ window: int | pd.Timedelta | numpy.timedelta64 | datetime.timedelta | str,
  min_periods=None,
  on: str | None = None,
  closed: Literal["right", "left", "both", "neither"] = "right",
  ) -> windows.Window:
- window_spec = window_specs.WindowSpec(
- bounds=window_specs.RowsWindowBounds.from_window_size(window, closed),
- min_periods=min_periods if min_periods is not None else window,
- grouping_keys=tuple(ex.deref(col) for col in self._by_col_ids),
- )
- block = self._block.order_by(
- [order.ascending_over(col) for col in self._by_col_ids],
- )
- skip_agg_col_id = (
- None if on is None else self._block.resolve_label_exact_or_error(on)
- )
- return windows.Window(
- block,
- window_spec,
- self._selected_cols,
+ if isinstance(window, int):
+ window_spec = window_specs.WindowSpec(
+ bounds=window_specs.RowsWindowBounds.from_window_size(window, closed),
+ min_periods=min_periods if min_periods is not None else window,
+ grouping_keys=tuple(ex.deref(col) for col in self._by_col_ids),
+ )
+ block = self._block.order_by(
+ [order.ascending_over(col) for col in self._by_col_ids],
+ )
+ skip_agg_col_id = (
+ None if on is None else self._block.resolve_label_exact_or_error(on)
+ )
+ return windows.Window(
+ block,
+ window_spec,
+ self._selected_cols,
+ drop_null_groups=self._dropna,
+ skip_agg_column_id=skip_agg_col_id,
+ )
+
+ return rolling.create_range_window(
+ self._block,
+ window,
+ min_periods=min_periods,
+ value_column_ids=self._selected_cols,
+ on=on,
+ closed=closed,
+ is_series=False,
+ grouping_keys=self._by_col_ids,
  drop_null_groups=self._dropna,
- skip_agg_column_id=skip_agg_col_id,
  )
 
  @validations.requires_ordering()

bigframes/core/groupby/series_group_by.py

@@ -14,11 +14,14 @@
 
 from __future__ import annotations
 
+import datetime
 import typing
 from typing import Literal, Sequence, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.groupby as vendored_pandas_groupby
+import numpy
+import pandas
 
 from bigframes import session
 from bigframes.core import expression as ex
@@ -29,6 +32,7 @@
 import bigframes.core.ordering as order
 import bigframes.core.utils as utils
 import bigframes.core.validations as validations
+from bigframes.core.window import rolling
 import bigframes.core.window as windows
 import bigframes.core.window_spec as window_specs
 import bigframes.dataframe as df
@@ -246,24 +250,36 @@ def diff(self, periods=1) -> series.Series:
  @validations.requires_ordering()
  def rolling(
  self,
- window: int,
+ window: int | pandas.Timedelta | numpy.timedelta64 | datetime.timedelta | str,
  min_periods=None,
  closed: Literal["right", "left", "both", "neither"] = "right",
  ) -> windows.Window:
- window_spec = window_specs.WindowSpec(
- bounds=window_specs.RowsWindowBounds.from_window_size(window, closed),
- min_periods=min_periods if min_periods is not None else window,
- grouping_keys=tuple(ex.deref(col) for col in self._by_col_ids),
- )
- block = self._block.order_by(
- [order.ascending_over(col) for col in self._by_col_ids],
- )
- return windows.Window(
- block,
- window_spec,
- [self._value_column],
- drop_null_groups=self._dropna,
+ if isinstance(window, int):
+ window_spec = window_specs.WindowSpec(
+ bounds=window_specs.RowsWindowBounds.from_window_size(window, closed),
+ min_periods=min_periods if min_periods is not None else window,
+ grouping_keys=tuple(ex.deref(col) for col in self._by_col_ids),
+ )
+ block = self._block.order_by(
+ [order.ascending_over(col) for col in self._by_col_ids],
+ )
+ return windows.Window(
+ block,
+ window_spec,
+ [self._value_column],
+ drop_null_groups=self._dropna,
+ is_series=True,
+ )
+
+ return rolling.create_range_window(
+ self._block,
+ window,
+ min_periods=min_periods,
+ value_column_ids=[self._value_column],
+ closed=closed,
  is_series=True,
+ grouping_keys=self._by_col_ids,
+ drop_null_groups=self._dropna,
  )
 
  @validations.requires_ordering()

bigframes/core/window/ordering.py

@@ -78,4 +78,9 @@ def _(root: nodes.WindowOpNode, column_id: str):
 
 @find_order_direction.register
 def _(root: nodes.ProjectionNode, column_id: str):
+ for expr, ref in root.assignments:
+ if ref.name == column_id and isinstance(expr, ex.DerefOp):
+ # This source column is renamed.
+ return find_order_direction(root.child, expr.id.name)
+
  return find_order_direction(root.child, column_id)

bigframes/core/window/rolling.py

@@ -76,18 +76,7 @@ def _apply_aggregate(
  self,
  op: agg_ops.UnaryAggregateOp,
  ):
- agg_col_ids = [
- col_id
- for col_id in self._value_column_ids
- if col_id != self._skip_agg_column_id
- ]
- agg_block = self._aggregate_block(op, agg_col_ids)
-
- if self._skip_agg_column_id is not None:
- # Concat the skipped column to the result.
- agg_block, _ = agg_block.join(
- self._block.select_column(self._skip_agg_column_id), how="outer"
- )
+ agg_block = self._aggregate_block(op)
 
  if self._is_series:
  from bigframes.series import Series
@@ -102,9 +91,12 @@ def _apply_aggregate(
  ]
  return DataFrame(agg_block)._reindex_columns(column_labels)
 
- def _aggregate_block(
- self, op: agg_ops.UnaryAggregateOp, agg_col_ids: typing.List[str]
- ) -> blocks.Block:
+ def _aggregate_block(self, op: agg_ops.UnaryAggregateOp) -> blocks.Block:
+ agg_col_ids = [
+ col_id
+ for col_id in self._value_column_ids
+ if col_id != self._skip_agg_column_id
+ ]
  block, result_ids = self._block.multi_apply_window_op(
  agg_col_ids,
  op,
@@ -123,39 +115,71 @@ def _aggregate_block(
  block = block.set_index(col_ids=index_ids)
 
  labels = [self._block.col_id_to_label[col] for col in agg_col_ids]
+ if self._skip_agg_column_id is not None:
+ result_ids = [self._skip_agg_column_id, *result_ids]
+ labels.insert(0, self._block.col_id_to_label[self._skip_agg_column_id])
+
  return block.select_columns(result_ids).with_column_labels(labels)
 
 
 def create_range_window(
  block: blocks.Block,
  window: pandas.Timedelta | numpy.timedelta64 | datetime.timedelta | str,
+ *,
+ value_column_ids: typing.Sequence[str] = tuple(),
  min_periods: int | None,
+ on: str | None = None,
  closed: typing.Literal["right", "left", "both", "neither"],
  is_series: bool,
+ grouping_keys: typing.Sequence[str] = tuple(),
+ drop_null_groups: bool = True,
 ) -> Window:
 
- index_dtypes = block.index.dtypes
- if len(index_dtypes) > 1:
- raise ValueError("Range rolling on MultiIndex is not supported")
- if index_dtypes[0] != dtypes.TIMESTAMP_DTYPE:
- raise ValueError("Index type should be timestamps with timezones")
+ if on is None:
+ # Rolling on index
+ index_dtypes = block.index.dtypes
+ if len(index_dtypes) > 1:
+ raise ValueError("Range rolling on MultiIndex is not supported")
+ if index_dtypes[0] != dtypes.TIMESTAMP_DTYPE:
+ raise ValueError("Index type should be timestamps with timezones")
+ rolling_key_col_id = block.index_columns[0]
+ else:
+ # Rolling on a specific column
+ rolling_key_col_id = block.resolve_label_exact_or_error(on)
+ if block.expr.get_column_type(rolling_key_col_id) != dtypes.TIMESTAMP_DTYPE:
+ raise ValueError(f"Column {on} type should be timestamps with timezones")
 
  order_direction = window_ordering.find_order_direction(
- block.expr.node, block.index_columns[0]
+ block.expr.node, rolling_key_col_id
  )
  if order_direction is None:
+ target_str = "index" if on is None else f"column {on}"
  raise ValueError(
- "The index might not be in a monotonic order. Please sort the index before rolling."
+ f"The {target_str} might not be in a monotonic order. Please sort by {target_str} before rolling."
  )
  if isinstance(window, str):
  window = pandas.Timedelta(window)
  spec = window_spec.WindowSpec(
  bounds=window_spec.RangeWindowBounds.from_timedelta_window(window, closed),
  min_periods=1 if min_periods is None else min_periods,
  ordering=(
- ordering.OrderingExpression(
- ex.deref(block.index_columns[0]), order_direction
- ),
+ ordering.OrderingExpression(ex.deref(rolling_key_col_id), order_direction),
  ),
+ grouping_keys=tuple(ex.deref(col) for col in grouping_keys),
+ )
+
+ selected_value_col_ids = (
+ value_column_ids if value_column_ids else block.value_columns
+ )
+ # This step must be done after finding the order direction of the window key.
+ if grouping_keys:
+ block = block.order_by([ordering.ascending_over(col) for col in grouping_keys])
+
+ return Window(
+ block,
+ spec,
+ value_column_ids=selected_value_col_ids,
+ is_series=is_series,
+ skip_agg_column_id=None if on is None else rolling_key_col_id,
+ drop_null_groups=drop_null_groups,
  )
- return Window(block, spec, block.value_columns, is_series=is_series)

bigframes/dataframe.py

@@ -67,12 +67,12 @@
 import bigframes.core.utils as utils
 import bigframes.core.validations as validations
 import bigframes.core.window
+from bigframes.core.window import rolling
 import bigframes.core.window_spec as windows
 import bigframes.dtypes
 import bigframes.exceptions as bfe
 import bigframes.formatting_helpers as formatter
 import bigframes.operations as ops
-import bigframes.operations.aggregations
 import bigframes.operations.aggregations as agg_ops
 import bigframes.operations.ai
 import bigframes.operations.plotting as plotting
@@ -3393,23 +3393,33 @@ def _perform_join_by_index(
  @validations.requires_ordering()
  def rolling(
  self,
- window: int,
+ window: int | pandas.Timedelta | numpy.timedelta64 | datetime.timedelta | str,
  min_periods=None,
  on: str | None = None,
  closed: Literal["right", "left", "both", "neither"] = "right",
  ) -> bigframes.core.window.Window:
- window_def = windows.WindowSpec(
- bounds=windows.RowsWindowBounds.from_window_size(window, closed),
- min_periods=min_periods if min_periods is not None else window,
- )
- skip_agg_col_id = (
- None if on is None else self._block.resolve_label_exact_or_error(on)
- )
- return bigframes.core.window.Window(
+ if isinstance(window, int):
+ window_def = windows.WindowSpec(
+ bounds=windows.RowsWindowBounds.from_window_size(window, closed),
+ min_periods=min_periods if min_periods is not None else window,
+ )
+ skip_agg_col_id = (
+ None if on is None else self._block.resolve_label_exact_or_error(on)
+ )
+ return bigframes.core.window.Window(
+ self._block,
+ window_def,
+ self._block.value_columns,
+ skip_agg_column_id=skip_agg_col_id,
+ )
+
+ return rolling.create_range_window(
  self._block,
- window_def,
- self._block.value_columns,
- skip_agg_column_id=skip_agg_col_id,
+ window,
+ min_periods=min_periods,
+ on=on,
+ closed=closed,
+ is_series=False,
  )
 
  @validations.requires_ordering()

bigframes/series.py

@@ -1590,7 +1590,11 @@ def rolling(
  )
 
  return rolling.create_range_window(
- self._block, window, min_periods, closed, is_series=True
+ block=self._block,
+ window=window,
+ min_periods=min_periods,
+ closed=closed,
+ is_series=True,
  )
 
  @validations.requires_ordering()