refactor: provide infrastructure for SQLGlot scalar compiler (#1850)

* refactor: provide infrastructure for SQLGlot scalar compiler

* remove redundant code

* remove redundant code

* add TODO back

Author: sycai

bigframes/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.

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

@@ -0,0 +1,49 @@
+# 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.
+
+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]()
+
+
+def compile(op: ops.BinaryOp, left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return BINARY_OP_REIGSTRATION[op](op, left, right)
+
+
+# TODO: add parenthesize for operators
+@BINARY_OP_REIGSTRATION.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
+        return sge.Concat(expressions=[left.expr, right.expr])
+
+    # Numerical addition
+    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:
+
+    return sge.GTE(this=left.expr, expression=right.expr)

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

@@ -0,0 +1,32 @@
+# 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.
+
+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]()
+
+
+def compile(op: ops.NaryOp, *args: TypedExpr) -> sge.Expression:
+    return NARY_OP_REIGSTRATION[op](op, *args)

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

@@ -0,0 +1,44 @@
+# 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.
+
+from __future__ import annotations
+
+import typing
+from typing import Generic, TypeVar
+
+from bigframes import operations as ops
+
+T = TypeVar("T")
+
+
+class OpRegistration(Generic[T]):
+    _registered_ops: dict[str, T] = {}
+
+    def register(
+        self, op: ops.ScalarOp | type[ops.ScalarOp]
+    ) -> typing.Callable[[T], T]:
+        key = typing.cast(str, op.name)
+
+        def decorator(item: T):
+            if key in self._registered_ops:
+                raise ValueError(f"{key} is already registered")
+            self._registered_ops[key] = item
+            return item
+
+        return decorator
+
+    def __getitem__(self, key: str | ops.ScalarOp) -> T:
+        if isinstance(key, ops.ScalarOp):
+            return self._registered_ops[key.name]
+        return self._registered_ops[key]

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

@@ -0,0 +1,35 @@
+# 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.
+
+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]()
+
+
+def compile(
+    op: ops.TernaryOp, expr1: TypedExpr, expr2: TypedExpr, expr3: TypedExpr
+) -> sge.Expression:
+    return TERNATRY_OP_REIGSTRATION[op](op, expr1, expr2, expr3)

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

@@ -0,0 +1,27 @@
+# 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.
+
+import dataclasses
+
+import sqlglot.expressions as sge
+
+from bigframes import dtypes
+
+
+@dataclasses.dataclass(frozen=True)
+class TypedExpr:
+    """SQLGlot expression with type."""
+
+    expr: sge.Expression
+    dtype: dtypes.ExpressionType

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

@@ -0,0 +1,31 @@
+# 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.
+
+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
+
+UnaryOpCompiler = typing.Callable[[ops.UnaryOp, TypedExpr], sge.Expression]
+
+UNARY_OP_REIGSTRATION = OpRegistration[UnaryOpCompiler]()
+
+
+def compile(op: ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return UNARY_OP_REIGSTRATION[op](op, expr)

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -13,25 +13,22 @@
 # limitations under the License.
 from __future__ import annotations
 
-import dataclasses
 import functools
 
 import sqlglot.expressions as sge
 
-from bigframes import dtypes
 from bigframes.core import expression
+from bigframes.core.compile.sqlglot.expressions import (
+    binary_compiler,
+    nary_compiler,
+    ternary_compiler,
+    typed_expr,
+    unary_compiler,
+)
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 import bigframes.operations as ops
 
 
-@dataclasses.dataclass(frozen=True)
-class TypedExpr:
-    """SQLGlot expression with type."""
-
-    expr: sge.Expression
-    dtype: dtypes.ExpressionType
-
-
 @functools.singledispatch
 def compile_scalar_expression(
     expression: expression.Expression,
@@ -63,46 +60,21 @@ def compile_constant_expression(
 def compile_op_expression(expr: expression.OpExpression) -> sge.Expression:
     # Non-recursively compiles the children scalar expressions.
     args = tuple(
-        TypedExpr(compile_scalar_expression(input), input.output_type)
+        typed_expr.TypedExpr(compile_scalar_expression(input), input.output_type)
         for input in expr.inputs
     )
 
     op = expr.op
-    op_name = expr.op.__class__.__name__
-    method_name = f"compile_{op_name.lower()}"
-    method = globals().get(method_name, None)
-    if method is None:
-        raise ValueError(
-            f"Compilation method '{method_name}' not found for operator '{op_name}'."
-        )
-
     if isinstance(op, ops.UnaryOp):
-        return method(op, args[0])
+        return unary_compiler.compile(op, args[0])
     elif isinstance(op, ops.BinaryOp):
-        return method(op, args[0], args[1])
+        return binary_compiler.compile(op, args[0], args[1])
     elif isinstance(op, ops.TernaryOp):
-        return method(op, args[0], args[1], args[2])
+        return ternary_compiler.compile(op, args[0], args[1], args[2])
     elif isinstance(op, ops.NaryOp):
-        return method(op, *args)
+        return nary_compiler.compile(op, *args)
     else:
         raise TypeError(
-            f"Operator '{op_name}' has an unrecognized arity or type "
+            f"Operator '{op.name}' has an unrecognized arity or type "
             "and cannot be compiled."
         )
-
-
-# TODO: add parenthesize for operators
-def compile_addop(op: ops.AddOp, left: TypedExpr, right: TypedExpr) -> sge.Expression:
-    if left.dtype == dtypes.STRING_DTYPE and right.dtype == dtypes.STRING_DTYPE:
-        # String addition
-        return sge.Concat(expressions=[left.expr, right.expr])
-
-    # Numerical addition
-    return sge.Add(this=left.expr, expression=right.expr)
-
-
-def compile_ge(
-    op: ops.ge_op, left: TypedExpr, right: TypedExpr  # type: ignore[valid-type]
-) -> sge.Expression:
-
-    return sge.GTE(this=left.expr, expression=right.expr)