feat: support time range rolling on Series. (#1590)

* [WIP] implement range rolling with mods at expression level * implement range rolling for series * fix lint * update rewrite logic * fix mypy * differentiate range bound and unbounded window * fix mypy * relax sorting key search algo * check whether window key is prefix * address comments * fix type hint * add comment on skipping range window during order pull up * fix lint * check order for all AdditiveNode * explicitly dispatch additive nodes

Author: sycai

bigframes/core/array_value.py

@@ -405,7 +405,7 @@ def project_window_op(
  skip_reproject_unsafe: skips the reprojection step, can be used when performing many non-dependent window operations, user responsible for not nesting window expressions, or using outputs as join, filter or aggregation keys before a reprojection
  """
  # TODO: Support non-deterministic windowing
- if window_spec.row_bounded or not op.order_independent:
+ if window_spec.is_row_bounded or not op.order_independent:
  if self.node.order_ambiguous and not self.session._strictly_ordered:
  if not self.session._allows_ambiguity:
  raise ValueError(

bigframes/core/compile/compiled.py

@@ -29,6 +29,7 @@
 from google.cloud import bigquery
 import pyarrow as pa
 
+from bigframes.core import utils
 import bigframes.core.compile.aggregate_compiler as agg_compiler
 import bigframes.core.compile.googlesql
 import bigframes.core.compile.ibis_types
@@ -231,7 +232,7 @@ def aggregate(
  col_out: agg_compiler.compile_aggregate(
  aggregate,
  bindings,
- order_by=_convert_ordering_to_table_values(table, order_by),
+ order_by=_convert_row_ordering_to_table_values(table, order_by),
  )
  for aggregate, col_out in aggregations
  }
@@ -439,7 +440,7 @@ def project_window_op(
  never_skip_nulls=never_skip_nulls,
  )
 
- if expression.op.order_independent and not window_spec.row_bounded:
+ if expression.op.order_independent and window_spec.is_unbounded:
  # notably percentile_cont does not support ordering clause
  window_spec = window_spec.without_order()
  window = self._ibis_window_from_spec(window_spec)
@@ -517,16 +518,30 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec):
  # 1. Order-independent op (aggregation, cut, rank) with unbound window - no ordering clause needed
  # 2. Order-independent op (aggregation, cut, rank) with range window - use ordering clause, ties allowed
  # 3. Order-depedenpent op (navigation functions, array_agg) or rows bounds - use total row order to break ties.
- if window_spec.ordering:
- order_by = _convert_ordering_to_table_values(
+ if window_spec.is_row_bounded:
+ if not window_spec.ordering:
+ # If window spec has following or preceding bounds, we need to apply an unambiguous ordering.
+ raise ValueError("No ordering provided for ordered analytic function")
+ order_by = _convert_row_ordering_to_table_values(
  self._column_names,
  window_spec.ordering,
  )
- elif window_spec.row_bounded:
- # If window spec has following or preceding bounds, we need to apply an unambiguous ordering.
- raise ValueError("No ordering provided for ordered analytic function")
- else:
+
+ elif window_spec.is_range_bounded:
+ order_by = [
+ _convert_range_ordering_to_table_value(
+ self._column_names,
+ window_spec.ordering[0],
+ )
+ ]
+ # The rest if branches are for unbounded windows
+ elif window_spec.ordering:
  # Unbound grouping window. Suitable for aggregations but not for analytic function application.
+ order_by = _convert_row_ordering_to_table_values(
+ self._column_names,
+ window_spec.ordering,
+ )
+ else:
  order_by = None
 
  window = bigframes_vendored.ibis.window(order_by=order_by, group_by=group_by)
@@ -551,7 +566,7 @@ def is_window(column: ibis_types.Value) -> bool:
  return any(isinstance(op, ibis_ops.WindowFunction) for op in matches)
 
 
-def _convert_ordering_to_table_values(
+def _convert_row_ordering_to_table_values(
  value_lookup: typing.Mapping[str, ibis_types.Value],
  ordering_columns: typing.Sequence[OrderingExpression],
 ) -> typing.Sequence[ibis_types.Value]:
@@ -579,6 +594,30 @@ def _convert_ordering_to_table_values(
  return ordering_values
 
 
+def _convert_range_ordering_to_table_value(
+ value_lookup: typing.Mapping[str, ibis_types.Value],
+ ordering_column: OrderingExpression,
+) -> ibis_types.Value:
+ """Converts the ordering for range windows to Ibis references.
+
+ Note that this method is different from `_convert_row_ordering_to_table_values` in
+ that it does not arrange null values. There are two reasons:
+ 1. Manipulating null positions requires more than one ordering key, which is forbidden
+ by SQL window syntax for range rolling.
+ 2. Pandas does not allow range rolling on timeseries with nulls.
+
+ Therefore, we opt for the simplest approach here: generate the simplest SQL and follow
+ the BigQuery engine behavior.
+ """
+ expr = op_compiler.compile_expression(
+ ordering_column.scalar_expression, value_lookup
+ )
+
+ if ordering_column.direction.is_ascending:
+ return bigframes_vendored.ibis.asc(expr) # type: ignore
+ return bigframes_vendored.ibis.desc(expr) # type: ignore
+
+
 def _string_cast_join_cond(
  lvalue: ibis_types.Column, rvalue: ibis_types.Column
 ) -> ibis_types.BooleanColumn:
@@ -668,8 +707,14 @@ def _add_boundary(
 ) -> ibis_expr_builders.LegacyWindowBuilder:
  if isinstance(bounds, RangeWindowBounds):
  return ibis_window.range(
- start=_to_ibis_boundary(bounds.start),
- end=_to_ibis_boundary(bounds.end),
+ start=_to_ibis_boundary(
+ None
+ if bounds.start is None
+ else utils.timedelta_to_micros(bounds.start)
+ ),
+ end=_to_ibis_boundary(
+ None if bounds.end is None else utils.timedelta_to_micros(bounds.end)
+ ),
  )
  if isinstance(bounds, RowsWindowBounds):
  if bounds.start is not None or bounds.end is not None:

bigframes/core/compile/compiler.py

@@ -82,6 +82,7 @@ def _replace_unsupported_ops(node: nodes.BigFrameNode):
  # TODO: Run all replacement rules as single bottom-up pass
  node = nodes.bottom_up(node, rewrites.rewrite_slice)
  node = nodes.bottom_up(node, rewrites.rewrite_timedelta_expressions)
+ node = nodes.bottom_up(node, rewrites.rewrite_range_rolling)
  return node
 
 

bigframes/core/compile/polars/compiler.py

@@ -16,7 +16,7 @@
 import dataclasses
 import functools
 import itertools
-from typing import cast, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import cast, Optional, Sequence, Tuple, TYPE_CHECKING
 
 import bigframes.core
 from bigframes.core import window_spec
@@ -359,6 +359,7 @@ def compile_window(self, node: nodes.WindowOpNode):
  return df.with_columns([agg_expr])
 
  else: # row-bounded window
+ assert isinstance(window.bounds, window_spec.RowsWindowBounds)
  # Polars API semi-bounded, and any grouped rolling window challenging
  # https://github.com/pola-rs/polars/issues/4799
  # https://github.com/pola-rs/polars/issues/8976
@@ -382,9 +383,7 @@ def compile_window(self, node: nodes.WindowOpNode):
  return pl.concat([df, results], how="horizontal")
 
 
-def _get_period(
- bounds: Union[window_spec.RowsWindowBounds, window_spec.RangeWindowBounds]
-) -> Optional[int]:
+def _get_period(bounds: window_spec.RowsWindowBounds) -> Optional[int]:
  """Returns None if the boundary is infinite."""
  if bounds.start is None or bounds.end is None:
  return None

bigframes/core/nodes.py

@@ -1353,7 +1353,7 @@ def _validate(self):
  """Validate the local data in the node."""
  # Since inner order and row bounds are coupled, rank ops can't be row bounded
  assert (
- not self.window_spec.row_bounded
+ not self.window_spec.is_row_bounded
  ) or self.expression.op.implicitly_inherits_order
  assert all(ref in self.child.ids for ref in self.expression.column_references)
 
@@ -1415,7 +1415,9 @@ def inherits_order(self) -> bool:
  op_inherits_order = (
  not self.expression.op.order_independent
  ) and self.expression.op.implicitly_inherits_order
- return op_inherits_order or self.window_spec.row_bounded
+ # range-bounded windows do not inherit orders because their ordering are
+ # already defined before rewrite time.
+ return op_inherits_order or self.window_spec.is_row_bounded
 
  @property
  def additive_base(self) -> BigFrameNode:

bigframes/core/rewrite/__init__.py

@@ -19,6 +19,7 @@
 from bigframes.core.rewrite.pruning import column_pruning
 from bigframes.core.rewrite.slices import pullup_limit_from_slice, rewrite_slice
 from bigframes.core.rewrite.timedeltas import rewrite_timedelta_expressions
+from bigframes.core.rewrite.windows import rewrite_range_rolling
 
 __all__ = [
  "legacy_join_as_projection",
@@ -29,4 +30,5 @@
  "remap_variables",
  "pull_up_order",
  "column_pruning",
+ "rewrite_range_rolling",
 ]

bigframes/core/rewrite/windows.py

@@ -0,0 +1,45 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import dataclasses
+
+from bigframes import operations as ops
+from bigframes.core import nodes
+
+
+def rewrite_range_rolling(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+ if not isinstance(node, nodes.WindowOpNode):
+ return node
+
+ if not node.window_spec.is_range_bounded:
+ return node
+
+ if len(node.window_spec.ordering) != 1:
+ raise ValueError(
+ "Range rolling should only be performed on exactly one column."
+ )
+
+ ordering_expr = node.window_spec.ordering[0]
+
+ new_ordering = dataclasses.replace(
+ ordering_expr,
+ scalar_expression=ops.UnixMicros().as_expr(ordering_expr.scalar_expression),
+ )
+
+ return dataclasses.replace(
+ node,
+ window_spec=dataclasses.replace(node.window_spec, ordering=(new_ordering,)),
+ )

bigframes/core/window/ordering.py

@@ -0,0 +1,81 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from functools import singledispatch
+
+from bigframes.core import expression as ex
+from bigframes.core import nodes, ordering
+
+
+@singledispatch
+def find_order_direction(
+ root: nodes.BigFrameNode, column_id: str
+) -> ordering.OrderingDirection | None:
+ """Returns the order of the given column with tree traversal. If the column cannot be found,
+ or the ordering information is not available, return None.
+ """
+ return None
+
+
+@find_order_direction.register
+def _(root: nodes.OrderByNode, column_id: str):
+ if len(root.by) == 0:
+ # This is a no-op
+ return find_order_direction(root.child, column_id)
+
+ # Make sure the window key is the prefix of sorting keys.
+ order_expr = root.by[0]
+ scalar_expr = order_expr.scalar_expression
+ if isinstance(scalar_expr, ex.DerefOp) and scalar_expr.id.name == column_id:
+ return order_expr.direction
+
+ return None
+
+
+@find_order_direction.register
+def _(root: nodes.ReversedNode, column_id: str):
+ direction = find_order_direction(root.child, column_id)
+
+ if direction is None:
+ return None
+ return direction.reverse()
+
+
+@find_order_direction.register
+def _(root: nodes.SelectionNode, column_id: str):
+ for alias_ref in root.input_output_pairs:
+ if alias_ref.id.name == column_id:
+ return find_order_direction(root.child, alias_ref.ref.id.name)
+
+
+@find_order_direction.register
+def _(root: nodes.FilterNode, column_id: str):
+ return find_order_direction(root.child, column_id)
+
+
+@find_order_direction.register
+def _(root: nodes.InNode, column_id: str):
+ return find_order_direction(root.left_child, column_id)
+
+
+@find_order_direction.register
+def _(root: nodes.WindowOpNode, column_id: str):
+ return find_order_direction(root.child, column_id)
+
+
+@find_order_direction.register
+def _(root: nodes.ProjectionNode, column_id: str):
+ return find_order_direction(root.child, column_id)

bigframes/core/window/rolling.py

@@ -14,12 +14,18 @@
 
 from __future__ import annotations
 
+import datetime
 import typing
 
 import bigframes_vendored.pandas.core.window.rolling as vendored_pandas_rolling
+import numpy
+import pandas
 
-from bigframes.core import log_adapter, window_spec
+from bigframes import dtypes
+from bigframes.core import expression as ex
+from bigframes.core import log_adapter, ordering, window_spec
 import bigframes.core.blocks as blocks
+from bigframes.core.window import ordering as window_ordering
 import bigframes.operations.aggregations as agg_ops
 
 
@@ -118,3 +124,38 @@ def _aggregate_block(
 
  labels = [self._block.col_id_to_label[col] for col in agg_col_ids]
  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,
+ min_periods: int | None,
+ closed: typing.Literal["right", "left", "both", "neither"],
+ is_series: bool,
+) -> 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")
+
+ order_direction = window_ordering.find_order_direction(
+ block.expr.node, block.index_columns[0]
+ )
+ if order_direction is None:
+ raise ValueError(
+ "The index might not be in a monotonic order. Please sort the index 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
+ ),
+ ),
+ )
+ return Window(block, spec, block.value_columns, is_series=is_series)