From 05a5963954236e6ad38305afc71a1ff4b4a25428 Mon Sep 17 00:00:00 2001 From: jialuo Date: Tue, 4 Nov 2025 02:05:06 +0000 Subject: [PATCH] chore: Migrate timestamp_add_op operator to SQLGlot Migrated the timestamp_add_op operator to SQLGlot. --- .../compile/sqlglot/expressions/datetime_ops.py | 15 +++++++++++++++ .../test_datetime_ops/test_timestamp_add/out.sql | 16 ++++++++++++++++ .../test_timestamp_diff/out.sql | 13 +++++++++++++ .../sqlglot/expressions/test_datetime_ops.py | 16 ++++++++++------ 4 files changed, 54 insertions(+), 6 deletions(-) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_add/out.sql create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_diff/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/datetime_ops.py b/bigframes/core/compile/sqlglot/expressions/datetime_ops.py index 949b122a1d..ffc2fde522 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,20 @@ 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_binary_op(ops.timestamp_add_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + return sge.TimestampAdd( + 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_add/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_add/out.sql new file mode 100644 index 0000000000..7d3b8ed206 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_timestamp_add/out.sql @@ -0,0 +1,16 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `timestamp_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + `bfcol_0` AS `bfcol_4`, + TIMESTAMP_ADD(`bfcol_1`, INTERVAL 1 MICROSECOND) AS `bfcol_5` + FROM `bfcte_0` +) +SELECT + `bfcol_4` AS `rowindex`, + `bfcol_5` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file 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..f41869c6a3 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py @@ -142,16 +142,20 @@ def test_second(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") -def test_strftime(scalar_types_df: bpd.DataFrame, snapshot): - col_name = "timestamp_col" - bf_df = scalar_types_df[[col_name]] - sql = utils._apply_ops_to_sql( - bf_df, [ops.StrftimeOp("%Y-%m-%d").as_expr(col_name)], [col_name] +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") +def test_timestamp_add(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + bf_df["timestamp_col"] = bf_df["timestamp_col"] + pd.Timedelta(1, unit="us") + snapshot.assert_match(bf_df.sql, "out.sql") + + def test_time(scalar_types_df: bpd.DataFrame, snapshot): col_name = "timestamp_col" bf_df = scalar_types_df[[col_name]]