feat: (Preview) Support diff aggregation for timestamp series. (#1405)

* [WIP] support time series diff

* add tests

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: sycai

bigframes/core/compile/aggregate_compiler.py

@@ -551,6 +551,24 @@ def _(
         raise TypeError(f"Cannot perform diff on type{column.type()}")
 
 
+@compile_unary_agg.register
+def _(
+    op: agg_ops.TimeSeriesDiffOp,
+    column: ibis_types.Column,
+    window=None,
+) -> ibis_types.Value:
+    if not column.type().is_timestamp():
+        raise TypeError(f"Cannot perform time series diff on type{column.type()}")
+
+    original_column = cast(ibis_types.TimestampColumn, column)
+    shifted_column = cast(
+        ibis_types.TimestampColumn,
+        compile_unary_agg(agg_ops.ShiftOp(op.periods), column, window),
+    )
+
+    return original_column.delta(shifted_column, part="microsecond")
+
+
 @compile_unary_agg.register
 def _(
     op: agg_ops.AllOp,

bigframes/core/rewrite/timedeltas.py

@@ -22,6 +22,7 @@
 from bigframes import operations as ops
 from bigframes.core import expression as ex
 from bigframes.core import nodes, schema, utils
+from bigframes.operations import aggregations as aggs
 
 
 @dataclasses.dataclass
@@ -59,6 +60,16 @@ def rewrite_timedelta_expressions(root: nodes.BigFrameNode) -> nodes.BigFrameNod
         by = tuple(_rewrite_ordering_expr(x, root.schema) for x in root.by)
         return nodes.OrderByNode(root.child, by)
 
+    if isinstance(root, nodes.WindowOpNode):
+        return nodes.WindowOpNode(
+            root.child,
+            _rewrite_aggregation(root.expression, root.schema),
+            root.window_spec,
+            root.output_name,
+            root.never_skip_nulls,
+            root.skip_reproject_unsafe,
+        )
+
     return root
 
 
@@ -166,3 +177,25 @@ def _rewrite_floordiv_op(left: _TypedExpr, right: _TypedExpr) -> _TypedExpr:
         return _TypedExpr.create_op_expr(ops.ToTimedeltaOp("us"), result)
 
     return result
+
+
+@functools.cache
+def _rewrite_aggregation(
+    aggregation: ex.Aggregation, schema: schema.ArraySchema
+) -> ex.Aggregation:
+    if not isinstance(aggregation, ex.UnaryAggregation):
+        return aggregation
+    if not isinstance(aggregation.op, aggs.DiffOp):
+        return aggregation
+
+    if isinstance(aggregation.arg, ex.DerefOp):
+        input_type = schema.get_type(aggregation.arg.id.sql)
+    else:
+        input_type = aggregation.arg.dtype
+
+    if dtypes.is_datetime_like(input_type):
+        return ex.UnaryAggregation(
+            aggs.TimeSeriesDiffOp(aggregation.op.periods), aggregation.arg
+        )
+
+    return aggregation

bigframes/operations/aggregations.py

@@ -484,6 +484,25 @@ class DiffOp(UnaryWindowOp):
     def skips_nulls(self):
         return False
 
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        if dtypes.is_datetime_like(input_types[0]):
+            return dtypes.TIMEDELTA_DTYPE
+        return super().output_type(*input_types)
+
+
+@dataclasses.dataclass(frozen=True)
+class TimeSeriesDiffOp(UnaryWindowOp):
+    periods: int
+
+    @property
+    def skips_nulls(self):
+        return False
+
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        if dtypes.is_datetime_like(input_types[0]):
+            return dtypes.TIMEDELTA_DTYPE
+        raise TypeError(f"expect datetime-like types, but got {input_types[0]}")
+
 
 @dataclasses.dataclass(frozen=True)
 class AllOp(UnaryAggregateOp):

tests/system/small/operations/test_datetimes.py

@@ -448,3 +448,15 @@ def test_timestamp_diff_literal_sub_series(scalars_dfs, column, value):
 
     expected_result = value - pd_series
     assert_series_equal(actual_result, expected_result)
+
+
+@pytest.mark.parametrize("column", ["timestamp_col", "datetime_col"])
+def test_timestamp_series_diff_agg(scalars_dfs, column):
+    bf_df, pd_df = scalars_dfs
+    bf_series = bf_df[column]
+    pd_series = pd_df[column]
+
+    actual_result = bf_series.diff().to_pandas()
+
+    expected_result = pd_series.diff()
+    assert_series_equal(actual_result, expected_result)