diff --git a/bigframes/core/compile/sqlglot/expressions/datetime_ops.py b/bigframes/core/compile/sqlglot/expressions/datetime_ops.py index 949b122a1d..4a62f970fc 100644 --- a/bigframes/core/compile/sqlglot/expressions/datetime_ops.py +++ b/bigframes/core/compile/sqlglot/expressions/datetime_ops.py @@ -21,6 +21,7 @@ import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op +register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op @register_unary_op(ops.FloorDtOp, pass_op=True) @@ -79,6 +80,13 @@ def _(expr: TypedExpr) -> sge.Expression: return sge.Extract(this=sge.Identifier(this="SECOND"), expression=expr.expr) +@register_binary_op(ops.timestamp_diff_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + return sge.TimestampDiff( + this=left.expr, expression=right.expr, unit=sge.Var(this="MICROSECOND") + ) + + @register_unary_op(ops.StrftimeOp, pass_op=True) def _(expr: TypedExpr, op: ops.StrftimeOp) -> sge.Expression: return sge.func("FORMAT_TIMESTAMP", sge.convert(op.date_format), expr.expr) diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_diff/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_diff/out.sql new file mode 100644 index 0000000000..671fe3a952 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_diff/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + TIMESTAMP_DIFF(`bfcol_0`, `bfcol_0`, MICROSECOND) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py index 6384dc79a9..c6be21d214 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py @@ -270,3 +270,11 @@ def test_sub_timedelta(scalar_types_df: bpd.DataFrame, snapshot): bf_df["timedelta_sub_timedelta"] = bf_df["duration_col"] - bf_df["duration_col"] snapshot.assert_match(bf_df.sql, "out.sql") + + +def test_timestamp_diff(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col", "date_col"]] + sql = utils._apply_binary_op( + bf_df, ops.timestamp_diff_op, "timestamp_col", "timestamp_col" + ) + snapshot.assert_match(sql, "out.sql")