From 5287c007b3371c971a42494d81209c82773ffd06 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 7 Jul 2026 23:59:44 +0000 Subject: [PATCH] fix(bigframes): Fix sqlglot backend regressions --- .../compile/sqlglot/expression_compiler.py | 4 ++ .../compile/sqlglot/expressions/array_ops.py | 2 +- .../sqlglot/expressions/datetime_ops.py | 59 +++++++++++++++++-- .../small/engines/test_comparison_ops.py | 16 ++++- .../system/small/engines/test_temporal_ops.py | 49 ++++++++++++++- 5 files changed, 122 insertions(+), 8 deletions(-) diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/expression_compiler.py b/packages/bigframes/bigframes/core/compile/sqlglot/expression_compiler.py index d07cb2186fbc..b412249a39f4 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/expression_compiler.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/expression_compiler.py @@ -48,6 +48,10 @@ class ExpressionCompiler: sge.LT, sge.EQ, sge.NEQ, + sge.Like, + sge.RegexpLike, + sge.In, + sge.Between, # Logical operations sge.And, sge.Or, diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/array_ops.py b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/array_ops.py index 8831b922e4ce..52a80b2a5c1c 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/array_ops.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/array_ops.py @@ -105,7 +105,7 @@ def _(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression: @register_unary_op(ops.ArrayToStringOp, pass_op=True) def _(expr: TypedExpr, op: ops.ArrayToStringOp) -> sge.Expression: - return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'") + return sge.ArrayToString(this=expr.expr, expression=sge.convert(op.delimiter)) @register_nary_op(ops.ToArrayOp) diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/datetime_ops.py b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/datetime_ops.py index 1f145f75879c..399b30622738 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/datetime_ops.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/datetime_ops.py @@ -373,9 +373,11 @@ def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression: ) return sge.Cast(this=result, to="DATETIME") + if expr.dtype == dtypes.TIMESTAMP_DTYPE: + return sge.func("DATETIME", expr.expr, sge.convert("UTC")) + if expr.dtype in ( dtypes.STRING_DTYPE, - dtypes.TIMESTAMP_DTYPE, dtypes.DATETIME_DTYPE, dtypes.DATE_DTYPE, ): @@ -387,9 +389,10 @@ def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression: if factor != 1: value = sge.Mul(this=value, expression=sge.convert(factor)) value = sge.func("TRUNC", value) - return sge.Cast( - this=sge.func("TIMESTAMP_MICROS", sge.Cast(this=value, to="INT64")), - to="DATETIME", + return sge.func( + "DATETIME", + sge.func("TIMESTAMP_MICROS", sge.Cast(this=value, to="INT64")), + sge.convert("UTC"), ) @@ -400,7 +403,7 @@ def _(expr: TypedExpr, op: ops.ToTimestampOp) -> sge.Expression: if expr.dtype != dtypes.STRING_DTYPE: result = sge.Cast(this=result, to="STRING") return sge.func( - "PARSE_TIMESTAMP", sge.convert(op.format), expr.expr, sge.convert("UTC") + "PARSE_TIMESTAMP", sge.convert(op.format), result, sge.convert("UTC") ) if expr.dtype in ( @@ -683,3 +686,49 @@ def _integer_label_to_datetime_op_yearly_freq( this=next_month_date, expression=sge.Interval(this=one, unit="DAY") ) return sge.Cast(this=x_label, to=sqlglot_types.from_bigframes_dtype(y.dtype)) + + +@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_binary_op(ops.timestamp_sub_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + return sge.TimestampSub( + this=left.expr, expression=right.expr, unit=sge.Var(this="MICROSECOND") + ) + + +@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.date_add_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + left_expr = sge.Cast(this=left.expr, to="TIMESTAMP") + return sge.TimestampAdd( + this=left_expr, expression=right.expr, unit=sge.Var(this="MICROSECOND") + ) + + +@register_binary_op(ops.date_sub_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + left_expr = sge.Cast(this=left.expr, to="TIMESTAMP") + return sge.TimestampSub( + this=left_expr, expression=right.expr, unit=sge.Var(this="MICROSECOND") + ) + + +@register_binary_op(ops.date_diff_op) +def _(left: TypedExpr, right: TypedExpr) -> sge.Expression: + diff = sge.DateDiff(this=left.expr, expression=right.expr, unit=sge.Var(this="DAY")) + return sge.Mul( + this=diff, + expression=sge.convert(int(UNIT_TO_US_CONVERSION_FACTORS["d"])), + ) diff --git a/packages/bigframes/tests/system/small/engines/test_comparison_ops.py b/packages/bigframes/tests/system/small/engines/test_comparison_ops.py index de0f110fa9b1..cd6ece558636 100644 --- a/packages/bigframes/tests/system/small/engines/test_comparison_ops.py +++ b/packages/bigframes/tests/system/small/engines/test_comparison_ops.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import itertools import pytest @@ -19,7 +20,7 @@ import bigframes.operations as ops from bigframes.core import array_value from bigframes.session import polars_executor -from bigframes.testing.engine_utils import assert_equivalence_execution +from bigframes.testing.engine_utils import SPEC, assert_equivalence_execution pytest.importorskip("polars") @@ -68,3 +69,16 @@ def test_engines_project_comparison_op( # bool col actually doesn't work properly for bq engine arr = apply_op_pairwise(scalars_array_value, op, excluded_cols=["string_col"]) assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine) + + +@pytest.mark.parametrize("engine", ["bq-sqlglot"], indirect=True) +def test_engines_precedence_like_and_in( + scalars_array_value: array_value.ArrayValue, engine +): + exprs = [ + ops.eq_op.as_expr("bool_col", ops.StrContainsOp("a").as_expr("string_col")), + ] + arr, _ = scalars_array_value.compute_values(exprs) + res = asyncio.run(engine.execute(arr.node, SPEC)) + assert res is not None + assert len(res.batches().to_pandas()) > 0 diff --git a/packages/bigframes/tests/system/small/engines/test_temporal_ops.py b/packages/bigframes/tests/system/small/engines/test_temporal_ops.py index cf22ce709752..61b1b06b1f29 100644 --- a/packages/bigframes/tests/system/small/engines/test_temporal_ops.py +++ b/packages/bigframes/tests/system/small/engines/test_temporal_ops.py @@ -12,12 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio + +import pandas as pd import pytest import bigframes.operations as ops +from bigframes import dtypes from bigframes.core import array_value +from bigframes.core import expression as ex from bigframes.session import polars_executor -from bigframes.testing.engine_utils import assert_equivalence_execution +from bigframes.testing.engine_utils import SPEC, assert_equivalence_execution pytest.importorskip("polars") @@ -64,3 +69,45 @@ def test_engines_date_accessors(scalars_array_value: array_value.ArrayValue, eng arr, _ = scalars_array_value.compute_values(exprs) assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine) + + +@pytest.mark.parametrize("engine", ["bq", "bq-sqlglot"], indirect=True) +def test_engines_temporal_arithmetic( + scalars_array_value: array_value.ArrayValue, engine +): + exprs = [ + ops.timestamp_add_op.as_expr( + "timestamp_col", ex.const(pd.Timedelta(seconds=1), dtypes.TIMEDELTA_DTYPE) + ), + ops.timestamp_sub_op.as_expr( + "timestamp_col", ex.const(pd.Timedelta(seconds=1), dtypes.TIMEDELTA_DTYPE) + ), + ops.date_add_op.as_expr( + "date_col", ex.const(pd.Timedelta(days=1), dtypes.TIMEDELTA_DTYPE) + ), + ops.date_sub_op.as_expr( + "date_col", ex.const(pd.Timedelta(days=1), dtypes.TIMEDELTA_DTYPE) + ), + ops.timestamp_diff_op.as_expr("timestamp_col", "timestamp_col"), + ops.date_diff_op.as_expr("date_col", "date_col"), + ] + + arr, _ = scalars_array_value.compute_values(exprs) + res = asyncio.run(engine.execute(arr.node, SPEC)) + assert res is not None + assert len(res.batches().to_pandas()) > 0 + + +@pytest.mark.parametrize("engine", ["bq", "bq-sqlglot"], indirect=True) +def test_engines_to_datetime(scalars_array_value: array_value.ArrayValue, engine): + exprs = [ + ops.ToDatetimeOp().as_expr("timestamp_col"), + ] + arr, _ = scalars_array_value.compute_values(exprs) + res = asyncio.run(engine.execute(arr.node, SPEC)) + assert res is not None + df = res.batches().to_pandas() + # The input timestamp was: TIMESTAMP('2021-07-21T17:43:43.945289+00:00') + # The output should be naive DATETIME('2021-07-21T17:43:43.945289') + val = df.iloc[0, -1] + assert pd.Timestamp(val) == pd.Timestamp("2021-07-21T17:43:43.945289")