remove on param from windowspec

Author: sycai

bigframes/core/blocks.py

@@ -1066,11 +1066,6 @@ def apply_window_op(
         skip_reproject_unsafe: bool = False,
         never_skip_nulls: bool = False,
     ) -> typing.Tuple[Block, str]:
-        if column == window_spec.on:
-            # For row-based window, do nothing
-            # TODO(b/388916840) Support range rolling with "on"
-            return self, column
-
         block = self
         if skip_null_groups:
             for key in window_spec.grouping_keys:

bigframes/core/groupby/dataframe_group_by.py

@@ -317,13 +317,19 @@ def rolling(
             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),
-            on=None if on is None else self._block.resolve_label_exact_or_error(on),
         )
         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()

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/dataframe.py

@@ -3318,10 +3318,15 @@ def rolling(
         window_def = windows.WindowSpec(
             bounds=windows.RowsWindowBounds.from_window_size(window, closed),
             min_periods=min_periods if min_periods is not None else window,
-            on=None if on is None else self._block.resolve_label_exact_or_error(on),
+        )
+        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()