googleapis · sycai · Feb 19, 2025 · Feb 15, 2025 · Feb 18, 2025 · Feb 18, 2025
@@ -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,

@@ -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
@@ -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):

@@ -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)