Merge branch 'main' into main_chelsealin_compilejoin

Author: chelsea-lin

.pre-commit-config.yaml

@@ -39,6 +39,6 @@ repos:
     rev: v1.15.0
     hooks:
     -   id: mypy
-        additional_dependencies: [types-requests, types-tabulate, pandas-stubs<=2.2.3.241126]
+        additional_dependencies: [types-requests, types-tabulate, types-PyYAML, pandas-stubs<=2.2.3.241126]
         exclude: "^third_party"
         args: ["--check-untyped-defs", "--explicit-package-bases", "--ignore-missing-imports"]

bigframes/clients.py

@@ -24,6 +24,7 @@
 import google.api_core.exceptions
 import google.api_core.retry
 from google.cloud import bigquery_connection_v1, resourcemanager_v3
+from google.iam.v1 import policy_pb2
 
 logger = logging.getLogger(__name__)
 
@@ -172,10 +173,7 @@ def _ensure_iam_binding(
                     return
 
         # Create a new binding
-        new_binding = {
-            "role": role,
-            "members": [service_account],
-        }  # Use a dictionary to avoid problematic google.iam namespace package.
+        new_binding = policy_pb2.Binding(role=role, members=[service_account])
         policy.bindings.append(new_binding)
         request = {
             "resource": project,

bigframes/core/compile/polars/compiler.py

@@ -23,9 +23,11 @@
 
 import bigframes.core
 from bigframes.core import identifiers, nodes, ordering, window_spec
+from bigframes.core.compile.polars import lowering
 import bigframes.core.expression as ex
 import bigframes.core.guid as guid
 import bigframes.core.rewrite
+import bigframes.core.rewrite.schema_binding
 import bigframes.dtypes
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
@@ -403,6 +405,8 @@ def compile(self, array_value: bigframes.core.ArrayValue) -> pl.LazyFrame:
         node = bigframes.core.rewrite.column_pruning(node)
         node = nodes.bottom_up(node, bigframes.core.rewrite.rewrite_slice)
         node = bigframes.core.rewrite.pull_out_window_order(node)
+        node = bigframes.core.rewrite.schema_binding.bind_schema_to_tree(node)
+        node = lowering.lower_ops_to_polars(node)
         return self.compile_node(node)
 
     @functools.singledispatchmethod

bigframes/core/compile/polars/lowering.py

@@ -0,0 +1,47 @@
+# 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 bigframes import dtypes
+from bigframes.core import bigframe_node, expression
+from bigframes.core.rewrite import op_lowering
+from bigframes.operations import numeric_ops
+import bigframes.operations as ops
+
+# TODO: Would be more precise to actually have separate op set for polars ops (where they diverge from the original ops)
+
+
+class LowerFloorDivRule(op_lowering.OpLoweringRule):
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return numeric_ops.FloorDivOp
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        dividend = expr.children[0]
+        divisor = expr.children[1]
+        using_floats = (dividend.output_type == dtypes.FLOAT_DTYPE) or (
+            divisor.output_type == dtypes.FLOAT_DTYPE
+        )
+        inf_or_zero = (
+            expression.const(float("INF")) if using_floats else expression.const(0)
+        )
+        zero_result = ops.mul_op.as_expr(inf_or_zero, dividend)
+        divisor_is_zero = ops.eq_op.as_expr(divisor, expression.const(0))
+        return ops.where_op.as_expr(zero_result, divisor_is_zero, expr)
+
+
+POLARS_LOWERING_RULES = (LowerFloorDivRule(),)
+
+
+def lower_ops_to_polars(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
+    return op_lowering.lower_ops(root, rules=POLARS_LOWERING_RULES)

bigframes/core/compile/sqlglot/compiler.py

@@ -244,6 +244,14 @@ def compile_concat(
             uid_gen=self.uid_gen,
         )
 
+    @_compile_node.register
+    def compile_explode(
+        self, node: nodes.ExplodeNode, child: ir.SQLGlotIR
+    ) -> ir.SQLGlotIR:
+        offsets_col = node.offsets_col.sql if (node.offsets_col is not None) else None
+        columns = tuple(ref.id.sql for ref in node.column_ids)
+        return child.explode(columns, offsets_col)
+
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):
     node = nodes.bottom_up(node, rewrite.rewrite_slice)

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)