refactor: fix test_date_series_diff_agg on agg_ops.DiffOp

Author: chelsea-lin

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -23,6 +23,7 @@
 from bigframes.core import window_spec
 import bigframes.core.compile.sqlglot.aggregations.op_registration as reg
 from bigframes.core.compile.sqlglot.aggregations.windows import apply_window_if_present
+from bigframes.core.compile.sqlglot.expressions import constants
 import bigframes.core.compile.sqlglot.expressions.typed_expr as typed_expr
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 from bigframes.operations import aggregations as agg_ops
@@ -326,6 +327,15 @@ def _(
             unit=sge.Identifier(this="MICROSECOND"),
         )
 
+    if column.dtype == dtypes.DATE_DTYPE:
+        date_diff = sge.DateDiff(
+            this=column.expr, expression=shifted, unit=sge.Identifier(this="DAY")
+        )
+        return sge.Cast(
+            this=sge.Floor(this=date_diff * constants._DAY_TO_MICROSECONDS),
+            to="INT64",
+        )
+
     raise TypeError(f"Cannot perform diff on type {column.dtype}")
 
 

bigframes/core/compile/sqlglot/expressions/constants.py

@@ -20,6 +20,7 @@
 _NAN = sge.Cast(this=sge.convert("NaN"), to="FLOAT64")
 _INF = sge.Cast(this=sge.convert("Infinity"), to="FLOAT64")
 _NEG_INF = sge.Cast(this=sge.convert("-Infinity"), to="FLOAT64")
+_DAY_TO_MICROSECONDS = sge.convert(86400000000)
 
 # Approx Highest number you can pass in to EXP function and get a valid FLOAT64 result
 # FLOAT64 has 11 exponent bits, so max values is about 2**(2**10)

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_diff_w_date/out.sql

@@ -0,0 +1,15 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `date_col`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    CAST(FLOOR(
+      DATE_DIFF(`date_col`, LAG(`date_col`, 1) OVER (ORDER BY `date_col` ASC NULLS LAST), DAY) * 86400000000
+    ) AS INT64) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `diff_date`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/test_unary_compiler.py

@@ -247,6 +247,17 @@ def test_diff_w_datetime(scalar_types_df: bpd.DataFrame, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
+def test_diff_w_date(scalar_types_df: bpd.DataFrame, snapshot):
+    col_name = "date_col"
+    bf_df_date = scalar_types_df[[col_name]]
+    window = window_spec.WindowSpec(ordering=(ordering.ascending_over(col_name),))
+    op = agg_exprs.UnaryAggregation(
+        agg_ops.DiffOp(periods=1), expression.deref(col_name)
+    )
+    sql = _apply_unary_window_op(bf_df_date, op, window, "diff_date")
+    snapshot.assert_match(sql, "out.sql")
+
+
 def test_diff_w_timestamp(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df_timestamp = scalar_types_df[[col_name]]