test: Enable floordiv local testing (#1856)

Author: TrevorBergeron

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/expression.py

@@ -19,7 +19,7 @@
 import functools
 import itertools
 import typing
-from typing import Generator, Mapping, TypeVar, Union
+from typing import Callable, Generator, Mapping, TypeVar, Union
 
 import pandas as pd
 
@@ -249,6 +249,10 @@ def is_identity(self) -> bool:
         """True for identity operation that does not transform input."""
         return False
 
+    @abc.abstractmethod
+    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
+        ...
+
     def walk(self) -> Generator[Expression, None, None]:
         yield self
         for child in self.children:
@@ -311,6 +315,9 @@ def __eq__(self, other):
 
         return self.value == other.value and self.dtype == other.dtype
 
+    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
+        return self
+
 
 @dataclasses.dataclass(frozen=True)
 class UnboundVariableExpression(Expression):
@@ -362,6 +369,9 @@ def is_bijective(self) -> bool:
     def is_identity(self) -> bool:
         return True
 
+    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
+        return self
+
 
 @dataclasses.dataclass(frozen=True)
 class DerefOp(Expression):
@@ -414,6 +424,9 @@ def is_bijective(self) -> bool:
     def is_identity(self) -> bool:
         return True
 
+    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
+        return self
+
 
 @dataclasses.dataclass(frozen=True)
 class SchemaFieldRefExpression(Expression):
@@ -463,12 +476,15 @@ def is_bijective(self) -> bool:
     def is_identity(self) -> bool:
         return True
 
+    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
+        return self
+
 
 @dataclasses.dataclass(frozen=True)
 class OpExpression(Expression):
     """An expression representing a scalar operation applied to 1 or more argument sub-expressions."""
 
-    op: bigframes.operations.RowOp
+    op: bigframes.operations.ScalarOp
     inputs: typing.Tuple[Expression, ...]
 
     @property
@@ -553,6 +569,12 @@ def deterministic(self) -> bool:
             all(input.deterministic for input in self.inputs) and self.op.deterministic
         )
 
+    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
+        new_inputs = tuple(t(input) for input in self.inputs)
+        if new_inputs != self.inputs:
+            return dataclasses.replace(self, inputs=new_inputs)
+        return self
+
 
 def bind_schema_fields(
     expr: Expression, field_by_id: Mapping[ids.ColumnId, field.Field]

bigframes/core/nodes.py

@@ -1008,6 +1008,14 @@ def referenced_ids(self) -> COLUMN_SET:
     def _node_expressions(self):
         return (self.predicate,)
 
+    def transform_exprs(
+        self, fn: Callable[[ex.Expression], ex.Expression]
+    ) -> FilterNode:
+        return dataclasses.replace(
+            self,
+            predicate=fn(self.predicate),
+        )
+
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> FilterNode:
@@ -1066,6 +1074,20 @@ def referenced_ids(self) -> COLUMN_SET:
     def _node_expressions(self):
         return tuple(map(lambda x: x.scalar_expression, self.by))
 
+    def transform_exprs(
+        self, fn: Callable[[ex.Expression], ex.Expression]
+    ) -> OrderByNode:
+        new_by = cast(
+            tuple[OrderingExpression, ...],
+            tuple(
+                dataclasses.replace(
+                    by_expr, scalar_expression=fn(by_expr.scalar_expression)
+                )
+                for by_expr in self.by
+            ),
+        )
+        return dataclasses.replace(self, by=new_by)
+
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> OrderByNode:
@@ -1078,14 +1100,9 @@ def remap_refs(
             itertools.chain.from_iterable(map(lambda x: x.referenced_columns, self.by))
         )
         ref_mapping = {id: ex.DerefOp(mappings[id]) for id in all_refs}
-        new_by = cast(
-            tuple[OrderingExpression, ...],
-            tuple(
-                by_expr.bind_refs(ref_mapping, allow_partial_bindings=True)
-                for by_expr in self.by
-            ),
+        return self.transform_exprs(
+            lambda ex: ex.bind_refs(ref_mapping, allow_partial_bindings=True)
         )
-        return dataclasses.replace(self, by=new_by)
 
 
 @dataclasses.dataclass(frozen=True, eq=False)
@@ -1293,6 +1310,12 @@ def _node_expressions(self):
     def additive_base(self) -> BigFrameNode:
         return self.child
 
+    def transform_exprs(
+        self, fn: Callable[[ex.Expression], ex.Expression]
+    ) -> ProjectionNode:
+        new_fields = tuple((fn(ex), id) for ex, id in self.assignments)
+        return dataclasses.replace(self, assignments=new_fields)
+
     def replace_additive_base(self, node: BigFrameNode) -> ProjectionNode:
         return dataclasses.replace(self, child=node)
 

bigframes/core/rewrite/op_lowering.py

@@ -0,0 +1,57 @@
+# 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 abc
+from typing import Sequence
+
+from bigframes.core import bigframe_node, expression, nodes
+import bigframes.operations as ops
+
+
+class OpLoweringRule(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def op(self) -> type[ops.ScalarOp]:
+        ...
+
+    @abc.abstractmethod
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        ...
+
+
+def lower_ops(
+    root: bigframe_node.BigFrameNode, rules: Sequence[OpLoweringRule]
+) -> bigframe_node.BigFrameNode:
+    rules_by_op = {rule.op: rule for rule in rules}
+
+    def lower_expr(expr: expression.Expression):
+        def lower_expr_step(expr: expression.Expression) -> expression.Expression:
+            if isinstance(expr, expression.OpExpression):
+                maybe_rule = rules_by_op.get(expr.op.__class__)
+                if maybe_rule:
+                    return maybe_rule.lower(expr)
+            return expr
+
+        return lower_expr_step(expr.transform_children(lower_expr_step))
+
+    def lower_node(node: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
+        if isinstance(
+            node, (nodes.ProjectionNode, nodes.FilterNode, nodes.OrderByNode)
+        ):
+            return node.transform_exprs(lower_expr)
+        else:
+            return node
+
+    return root.bottom_up(lower_node)

tests/unit/test_dataframe_polars.py

@@ -2150,7 +2150,7 @@ def test_df_corrwith_series(scalars_dfs):
         operator.sub,
         operator.mul,
         operator.truediv,
-        # operator.floordiv,
+        operator.floordiv,
         operator.eq,
         operator.ne,
         operator.gt,
@@ -2163,7 +2163,7 @@ def test_df_corrwith_series(scalars_dfs):
         "subtract",
         "multiply",
         "true_divide",
-        # "floor_divide",
+        "floor_divide",
         "eq",
         "ne",
         "gt",
@@ -2217,8 +2217,8 @@ def test_scalar_binop_str_exception(scalars_dfs):
         (lambda x, y: x.rmul(y, axis="index")),
         (lambda x, y: x.truediv(y, axis="index")),
         (lambda x, y: x.rtruediv(y, axis="index")),
-        # (lambda x, y: x.floordiv(y, axis="index")),
-        # (lambda x, y: x.floordiv(y, axis="index")),
+        (lambda x, y: x.floordiv(y, axis="index")),
+        (lambda x, y: x.floordiv(y, axis="index")),
         (lambda x, y: x.gt(y, axis="index")),
         (lambda x, y: x.ge(y, axis="index")),
         (lambda x, y: x.lt(y, axis="index")),
@@ -2233,8 +2233,8 @@ def test_scalar_binop_str_exception(scalars_dfs):
         "rmul",
         "truediv",
         "rtruediv",
-        # "floordiv",
-        # "rfloordiv",
+        "floordiv",
+        "rfloordiv",
         "gt",
         "ge",
         "lt",