feat: add on parameter in dataframe.rolling() and dataframe.groupby.rolling() (#1556)

* feat: add  parameter in dataframe/groupby.rolling()

* fix test

* add docs

* update doc

* update docs

* remove on param from windowspec

Author: sycai

bigframes/core/blocks.py

@@ -987,7 +987,7 @@ def apply_nary_op(
     def multi_apply_window_op(
         self,
         columns: typing.Sequence[str],
-        op: agg_ops.WindowOp,
+        op: agg_ops.UnaryWindowOp,
         window_spec: windows.WindowSpec,
         *,
         skip_null_groups: bool = False,
@@ -1058,7 +1058,7 @@ def project_exprs(
     def apply_window_op(
         self,
         column: str,
-        op: agg_ops.WindowOp,
+        op: agg_ops.UnaryWindowOp,
         window_spec: windows.WindowSpec,
         *,
         result_label: Label = None,

bigframes/core/groupby/dataframe_group_by.py

@@ -310,6 +310,7 @@ def rolling(
         self,
         window: int,
         min_periods=None,
+        on: str | None = None,
         closed: Literal["right", "left", "both", "neither"] = "right",
     ) -> windows.Window:
         window_spec = window_specs.WindowSpec(
@@ -320,8 +321,15 @@ def rolling(
         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
+            block,
+            window_spec,
+            self._selected_cols,
+            drop_null_groups=self._dropna,
+            skip_agg_column_id=skip_agg_col_id,
         )
 
     @validations.requires_ordering()
@@ -511,7 +519,7 @@ def _aggregate_all(
 
     def _apply_window_op(
         self,
-        op: agg_ops.WindowOp,
+        op: agg_ops.UnaryWindowOp,
         window: typing.Optional[window_specs.WindowSpec] = None,
         numeric_only: bool = False,
     ):

bigframes/core/groupby/series_group_by.py

@@ -294,7 +294,7 @@ def _aggregate(self, aggregate_op: agg_ops.UnaryAggregateOp) -> series.Series:
 
     def _apply_window_op(
         self,
-        op: agg_ops.WindowOp,
+        op: agg_ops.UnaryWindowOp,
         discard_name=False,
         window: typing.Optional[window_specs.WindowSpec] = None,
         never_skip_nulls: bool = False,

bigframes/core/window/rolling.py

@@ -34,12 +34,16 @@ def __init__(
         value_column_ids: typing.Sequence[str],
         drop_null_groups: bool = True,
         is_series: bool = False,
+        skip_agg_column_id: str | None = None,
     ):
         self._block = block
         self._window_spec = window_spec
         self._value_column_ids = value_column_ids
         self._drop_null_groups = drop_null_groups
         self._is_series = is_series
+        # The column ID that won't be aggregated on.
+        # This is equivalent to pandas `on` parameter in rolling()
+        self._skip_agg_column_id = skip_agg_column_id
 
     def count(self):
         return self._apply_aggregate(agg_ops.count_op)
@@ -66,10 +70,37 @@ def _apply_aggregate(
         self,
         op: agg_ops.UnaryAggregateOp,
     ):
-        block = self._block
-        labels = [block.col_id_to_label[col] for col in self._value_column_ids]
-        block, result_ids = block.multi_apply_window_op(
-            self._value_column_ids,
+        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"
+            )
+
+        if self._is_series:
+            from bigframes.series import Series
+
+            return Series(agg_block)
+        else:
+            from bigframes.dataframe import DataFrame
+
+            # Preserve column order.
+            column_labels = [
+                self._block.col_id_to_label[col_id] for col_id in self._value_column_ids
+            ]
+            return DataFrame(agg_block)._reindex_columns(column_labels)
+
+    def _aggregate_block(
+        self, op: agg_ops.UnaryAggregateOp, agg_col_ids: typing.List[str]
+    ) -> blocks.Block:
+        block, result_ids = self._block.multi_apply_window_op(
+            agg_col_ids,
             op,
             self._window_spec,
             skip_null_groups=self._drop_null_groups,
@@ -85,13 +116,5 @@ def _apply_aggregate(
             )
             block = block.set_index(col_ids=index_ids)
 
-        if self._is_series:
-            from bigframes.series import Series
-
-            return Series(block.select_columns(result_ids).with_column_labels(labels))
-        else:
-            from bigframes.dataframe import DataFrame
-
-            return DataFrame(
-                block.select_columns(result_ids).with_column_labels(labels)
-            )
+        labels = [self._block.col_id_to_label[col] for col in agg_col_ids]
+        return block.select_columns(result_ids).with_column_labels(labels)

bigframes/core/window_spec.py

@@ -187,10 +187,12 @@ def __post_init__(self):
 class WindowSpec:
     """
     Specifies a window over which aggregate and analytic function may be applied.
-    grouping_keys: set of column ids to group on
-    preceding: Number of preceding rows in the window
-    following: Number of preceding rows in the window
-    ordering: List of columns ids and ordering direction to override base ordering
+
+    Attributes:
+        grouping_keys: A set of column ids to group on
+        bounds: The window boundaries
+        ordering: A list of columns ids and ordering direction to override base ordering
+        min_periods: The minimum number of observations in window required to have a value
     """
 
     grouping_keys: Tuple[ex.DerefOp, ...] = tuple()

bigframes/dataframe.py

@@ -3312,14 +3312,21 @@ def rolling(
         self,
         window: int,
         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(
-            self._block, window_def, self._block.value_columns
+            self._block,
+            window_def,
+            self._block.value_columns,
+            skip_agg_column_id=skip_agg_col_id,
         )
 
     @validations.requires_ordering()
@@ -3483,7 +3490,7 @@ def pct_change(self, periods: int = 1) -> DataFrame:
 
     def _apply_window_op(
         self,
-        op: agg_ops.WindowOp,
+        op: agg_ops.UnaryWindowOp,
         window_spec: windows.WindowSpec,
     ):
         block, result_ids = self._block.multi_apply_window_op(

bigframes/series.py

@@ -1378,7 +1378,9 @@ def _apply_aggregation(
     ) -> Any:
         return self._block.get_stat(self._value_column, op)
 
-    def _apply_window_op(self, op: agg_ops.WindowOp, window_spec: windows.WindowSpec):
+    def _apply_window_op(
+        self, op: agg_ops.UnaryWindowOp, window_spec: windows.WindowSpec
+    ):
         block = self._block
         block, result_id = block.apply_window_op(
             self._value_column, op, window_spec=window_spec, result_label=self.name

tests/system/small/test_window.py

@@ -20,12 +20,9 @@
 def rolling_dfs(scalars_dfs):
     bf_df, pd_df = scalars_dfs
 
-    target_cols = ["int64_too", "float64_col", "bool_col"]
+    target_cols = ["int64_too", "float64_col", "int64_col"]
 
-    bf_df = bf_df[target_cols].set_index("bool_col")
-    pd_df = pd_df[target_cols].set_index("bool_col")
-
-    return bf_df, pd_df
+    return bf_df[target_cols], pd_df[target_cols]
 
 
 @pytest.fixture(scope="module")
@@ -49,31 +46,65 @@ def test_dataframe_rolling_closed_param(rolling_dfs, closed):
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 def test_dataframe_groupby_rolling_closed_param(rolling_dfs, closed):
     bf_df, pd_df = rolling_dfs
+    # Need to specify column subset for comparison due to b/406841327
+    check_columns = ["float64_col", "int64_col"]
 
     actual_result = (
-        bf_df.groupby(level=0).rolling(window=3, closed=closed).sum().to_pandas()
+        bf_df.groupby(bf_df["int64_too"] % 2)
+        .rolling(window=3, closed=closed)
+        .sum()
+        .to_pandas()
     )
 
-    expected_result = pd_df.groupby(level=0).rolling(window=3, closed=closed).sum()
-    pd.testing.assert_frame_equal(actual_result, expected_result, check_dtype=False)
+    expected_result = (
+        pd_df.groupby(pd_df["int64_too"] % 2).rolling(window=3, closed=closed).sum()
+    )
+    pd.testing.assert_frame_equal(
+        actual_result[check_columns], expected_result, check_dtype=False
+    )
 
 
-def test_dataframe_rolling_default_closed_param(rolling_dfs):
+def test_dataframe_rolling_on(rolling_dfs):
     bf_df, pd_df = rolling_dfs
 
-    actual_result = bf_df.rolling(window=3).sum().to_pandas()
+    actual_result = bf_df.rolling(window=3, on="int64_too").sum().to_pandas()
 
-    expected_result = pd_df.rolling(window=3).sum()
+    expected_result = pd_df.rolling(window=3, on="int64_too").sum()
     pd.testing.assert_frame_equal(actual_result, expected_result, check_dtype=False)
 
 
-def test_dataframe_groupby_rolling_default_closed_param(rolling_dfs):
+def test_dataframe_rolling_on_invalid_column_raise_error(rolling_dfs):
+    bf_df, _ = rolling_dfs
+
+    with pytest.raises(ValueError):
+        bf_df.rolling(window=3, on="whatever").sum()
+
+
+def test_dataframe_groupby_rolling_on(rolling_dfs):
     bf_df, pd_df = rolling_dfs
+    # Need to specify column subset for comparison due to b/406841327
+    check_columns = ["float64_col", "int64_col"]
 
-    actual_result = bf_df.groupby(level=0).rolling(window=3).sum().to_pandas()
+    actual_result = (
+        bf_df.groupby(bf_df["int64_too"] % 2)
+        .rolling(window=3, on="float64_col")
+        .sum()
+        .to_pandas()
+    )
 
-    expected_result = pd_df.groupby(level=0).rolling(window=3).sum()
-    pd.testing.assert_frame_equal(actual_result, expected_result, check_dtype=False)
+    expected_result = (
+        pd_df.groupby(pd_df["int64_too"] % 2).rolling(window=3, on="float64_col").sum()
+    )
+    pd.testing.assert_frame_equal(
+        actual_result[check_columns], expected_result, check_dtype=False
+    )
+
+
+def test_dataframe_groupby_rolling_on_invalid_column_raise_error(rolling_dfs):
+    bf_df, _ = rolling_dfs
+
+    with pytest.raises(ValueError):
+        bf_df.groupby(level=0).rolling(window=3, on="whatever").sum()
 
 
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
@@ -103,24 +134,6 @@ def test_series_groupby_rolling_closed_param(rolling_series, closed):
     pd.testing.assert_series_equal(actual_result, expected_result, check_dtype=False)
 
 
-def test_series_rolling_default_closed_param(rolling_series):
-    bf_series, df_series = rolling_series
-
-    actual_result = bf_series.rolling(window=3).sum().to_pandas()
-
-    expected_result = df_series.rolling(window=3).sum()
-    pd.testing.assert_series_equal(actual_result, expected_result, check_dtype=False)
-
-
-def test_series_groupby_rolling_default_closed_param(rolling_series):
-    bf_series, df_series = rolling_series
-
-    actual_result = bf_series.groupby(bf_series % 2).rolling(window=3).sum().to_pandas()
-
-    expected_result = df_series.groupby(df_series % 2).rolling(window=3).sum()
-    pd.testing.assert_series_equal(actual_result, expected_result, check_dtype=False)
-
-
 @pytest.mark.parametrize(
     ("windowing"),
     [
@@ -181,8 +194,12 @@ def test_series_window_agg_ops(rolling_series, windowing, agg_op):
         pytest.param(lambda x: x.var(), id="var"),
     ],
 )
-def test_dataframe_window_agg_ops(rolling_dfs, windowing, agg_op):
-    bf_df, pd_df = rolling_dfs
+def test_dataframe_window_agg_ops(scalars_dfs, windowing, agg_op):
+    bf_df, pd_df = scalars_dfs
+    target_columns = ["int64_too", "float64_col", "bool_col"]
+    index_column = "bool_col"
+    bf_df = bf_df[target_columns].set_index(index_column)
+    pd_df = pd_df[target_columns].set_index(index_column)
 
     bf_result = agg_op(windowing(bf_df)).to_pandas()