chore: add array operators to SQLGlot compiler (#1852)

* [WIP] Add array operators. Need to finish tests

* add tests

* fix lint

* fix typos

* Use sge.Bracket() for safe_offset

Author: sycai

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

@@ -14,26 +14,22 @@
 
 from __future__ import annotations
 
-import typing
-
 import sqlglot.expressions as sge
 
 from bigframes import dtypes
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot.expressions.op_registration import OpRegistration
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
-BinaryOpCompiler = typing.Callable[[ops.BinaryOp, TypedExpr, TypedExpr], sge.Expression]
-
-BINARY_OP_REIGSTRATION = OpRegistration[BinaryOpCompiler]()
+BINARY_OP_REGISTRATION = OpRegistration()
 
 
 def compile(op: ops.BinaryOp, left: TypedExpr, right: TypedExpr) -> sge.Expression:
-    return BINARY_OP_REIGSTRATION[op](op, left, right)
+    return BINARY_OP_REGISTRATION[op](op, left, right)
 
 
 # TODO: add parenthesize for operators
-@BINARY_OP_REIGSTRATION.register(ops.add_op)
+@BINARY_OP_REGISTRATION.register(ops.add_op)
 def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
     if left.dtype == dtypes.STRING_DTYPE and right.dtype == dtypes.STRING_DTYPE:
         # String addition
@@ -43,7 +39,6 @@ def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.Add(this=left.expr, expression=right.expr)
 
 
-@BINARY_OP_REIGSTRATION.register(ops.ge_op)
-def compile_ge(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
-
+@BINARY_OP_REGISTRATION.register(ops.ge_op)
+def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.GTE(this=left.expr, expression=right.expr)

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

@@ -14,19 +14,14 @@
 
 from __future__ import annotations
 
-import typing
-
 import sqlglot.expressions as sge
 
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot.expressions.op_registration import OpRegistration
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
-# No simpler way to specify that the compilation function expects varargs.
-NaryOpCompiler = typing.Callable[..., sge.Expression]
-
-NARY_OP_REIGSTRATION = OpRegistration[NaryOpCompiler]()
+NARY_OP_REGISTRATION = OpRegistration()
 
 
 def compile(op: ops.NaryOp, *args: TypedExpr) -> sge.Expression:
-    return NARY_OP_REIGSTRATION[op](op, *args)
+    return NARY_OP_REGISTRATION[op](op, *args)

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

@@ -15,30 +15,40 @@
 from __future__ import annotations
 
 import typing
-from typing import Generic, TypeVar
+
+from sqlglot import expressions as sge
 
 from bigframes import operations as ops
 
-T = TypeVar("T")
+# We should've been more specific about input types. Unfortunately,
+# MyPy doesn't support more rigorous checks.
+CompilationFunc = typing.Callable[..., sge.Expression]
 
 
-class OpRegistration(Generic[T]):
-    _registered_ops: dict[str, T] = {}
+class OpRegistration:
+    def __init__(self) -> None:
+        self._registered_ops: dict[str, CompilationFunc] = {}
 
     def register(
         self, op: ops.ScalarOp | type[ops.ScalarOp]
-    ) -> typing.Callable[[T], T]:
-        key = typing.cast(str, op.name)
-
-        def decorator(item: T):
+    ) -> typing.Callable[[CompilationFunc], CompilationFunc]:
+        def decorator(item: CompilationFunc):
+            def arg_checker(*args, **kwargs):
+                if not isinstance(args[0], ops.ScalarOp):
+                    raise ValueError(
+                        f"The first parameter must be an operator. Got {type(args[0])}"
+                    )
+                return item(*args, **kwargs)
+
+            key = typing.cast(str, op.name)
             if key in self._registered_ops:
                 raise ValueError(f"{key} is already registered")
             self._registered_ops[key] = item
-            return item
+            return arg_checker
 
         return decorator
 
-    def __getitem__(self, key: str | ops.ScalarOp) -> T:
+    def __getitem__(self, key: str | ops.ScalarOp) -> CompilationFunc:
         if isinstance(key, ops.ScalarOp):
             return self._registered_ops[key.name]
         return self._registered_ops[key]

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

@@ -14,22 +14,16 @@
 
 from __future__ import annotations
 
-import typing
-
 import sqlglot.expressions as sge
 
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot.expressions.op_registration import OpRegistration
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
-TernaryOpCompiler = typing.Callable[
-    [ops.TernaryOp, TypedExpr, TypedExpr, TypedExpr], sge.Expression
-]
-
-TERNATRY_OP_REIGSTRATION = OpRegistration[TernaryOpCompiler]()
+TERNATRY_OP_REGISTRATION = OpRegistration()
 
 
 def compile(
     op: ops.TernaryOp, expr1: TypedExpr, expr2: TypedExpr, expr3: TypedExpr
 ) -> sge.Expression:
-    return TERNATRY_OP_REIGSTRATION[op](op, expr1, expr2, expr3)
+    return TERNATRY_OP_REGISTRATION[op](op, expr1, expr2, expr3)

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

@@ -16,16 +16,57 @@
 
 import typing
 
+import sqlglot
 import sqlglot.expressions as sge
 
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot.expressions.op_registration import OpRegistration
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
-UnaryOpCompiler = typing.Callable[[ops.UnaryOp, TypedExpr], sge.Expression]
-
-UNARY_OP_REIGSTRATION = OpRegistration[UnaryOpCompiler]()
+UNARY_OP_REGISTRATION = OpRegistration()
 
 
 def compile(op: ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
-    return UNARY_OP_REIGSTRATION[op](op, expr)
+    return UNARY_OP_REGISTRATION[op](op, expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.ArrayToStringOp)
+def _(op: ops.ArrayToStringOp, expr: TypedExpr) -> sge.Expression:
+    return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
+
+
+@UNARY_OP_REGISTRATION.register(ops.ArrayIndexOp)
+def _(op: ops.ArrayIndexOp, expr: TypedExpr) -> sge.Expression:
+    return sge.Bracket(
+        this=expr.expr,
+        expressions=[sge.Literal.number(op.index)],
+        safe=True,
+        offset=False,
+    )
+
+
+@UNARY_OP_REGISTRATION.register(ops.ArraySliceOp)
+def _(op: ops.ArraySliceOp, expr: TypedExpr) -> sge.Expression:
+    slice_idx = sqlglot.to_identifier("slice_idx")
+
+    conditions: typing.List[sge.Predicate] = [slice_idx >= op.start]
+
+    if op.stop is not None:
+        conditions.append(slice_idx < op.stop)
+
+    # local name for each element in the array
+    el = sqlglot.to_identifier("el")
+
+    selected_elements = (
+        sge.select(el)
+        .from_(
+            sge.Unnest(
+                expressions=[expr.expr],
+                alias=sge.TableAlias(columns=[el]),
+                offset=slice_idx,
+            )
+        )
+        .where(*conditions)
+    )
+
+    return sge.array(selected_elements)

tests/unit/core/compile/sqlglot/expressions/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unit/core/compile/sqlglot/snapshots/test_compile_scalar_expr/test_compile_numerical_add/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql

@@ -0,0 +1,15 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `rowindex` AS `bfcol_0`,
+    `string_list_col` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`repeated_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_1`[SAFE_OFFSET(1)] AS `bfcol_4`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_0` AS `rowindex`,
+  `bfcol_4` AS `string_list_col`
+FROM `bfcte_1`