fix(bigframes): Fix sqlglot backend regressions (#17655)

Author: TrevorBergeron

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,

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)

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"])),
+    )

packages/bigframes/tests/system/small/engines/test_comparison_ops.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import itertools
 
 import pytest
 
 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

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

packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_datetime/out.sql

@@ -1,7 +1,7 @@
 SELECT
-  CAST(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col` * 0.001) AS INT64)) AS DATETIME) AS `int64_col`,
+  DATETIME(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col` * 0.001) AS INT64)), 'UTC') AS `int64_col`,
   SAFE_CAST(`string_col` AS DATETIME) AS `string_col`,
-  CAST(TIMESTAMP_MICROS(CAST(TRUNC(`float64_col` * 0.001) AS INT64)) AS DATETIME) AS `float64_col`,
-  SAFE_CAST(`timestamp_col` AS DATETIME) AS `timestamp_col`,
+  DATETIME(TIMESTAMP_MICROS(CAST(TRUNC(`float64_col` * 0.001) AS INT64)), 'UTC') AS `float64_col`,
+  DATETIME(`timestamp_col`, 'UTC') AS `timestamp_col`,
   SAFE_CAST(`string_col` AS DATETIME) AS `string_col_fmt`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`