more tests, casting fix

Author: TrevorBergeron

packages/bigframes/bigframes/core/compile/polars/compiler.py

@@ -39,7 +39,6 @@
 import bigframes.operations.numeric_ops as num_ops
 import bigframes.operations.string_ops as string_ops
 from bigframes.core import agg_expressions, identifiers, nodes, ordering, window_spec
-from bigframes.core.compile.polars import lowering
 
 polars_installed = True
 if TYPE_CHECKING:
@@ -672,6 +671,8 @@ def compile(self, plan: nodes.BigFrameNode) -> pl.LazyFrame:
             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)
+            from bigframes.core.compile import lowering
+
             node = lowering.lower_ops_to_polars(node)
             return self.compile_node(node)
 
@@ -763,6 +764,8 @@ def compile_join(self, node: nodes.JoinNode):
             left_on = []
             right_on = []
             for left_ex, right_ex in node.conditions:
+                from bigframes.core.compile import lowering
+
                 left_ex, right_ex = lowering._coerce_comparables(left_ex, right_ex)
                 left_on.append(self.expr_compiler.compile_expression(left_ex))
                 right_on.append(self.expr_compiler.compile_expression(right_ex))
@@ -782,6 +785,8 @@ def compile_isin(self, node: nodes.InNode):
             right = right.with_columns(pl.lit(True).alias(node.indicator_col.sql))
 
             right_col = ex.ResolvedDerefOp.from_field(node.right_child.fields[0])
+            from bigframes.core.compile import lowering
+
             left_ex, right_ex = lowering._coerce_comparables(node.left_col, right_col)
 
             left_pl_ex = self.expr_compiler.compile_expression(left_ex)

packages/bigframes/bigframes/core/compile/substrait/compiler.py

@@ -30,7 +30,7 @@
 import bigframes.operations.generic_ops as generic_ops
 import bigframes.operations.numeric_ops as numeric_ops
 import bigframes.operations.struct_ops as struct_ops
-from bigframes.core import bigframe_node, nodes
+from bigframes.core import bigframe_node, nodes, rewrite
 from bigframes.core.compile import lowering
 
 
@@ -54,6 +54,8 @@ def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
         if not self.can_compile(plan):
             return None
 
+        # Need to bind types in before lowering
+        plan = rewrite.bind_schema_to_tree(plan)
         plan = lowering.lower_ops_to_substrait(plan)
         pb_rel = self._compile_node(plan)
 
@@ -674,7 +676,8 @@ def _compile_cast(
         type_dict = self._convert_type(target_dtype)
         json_format.ParseDict(type_dict, cast.type)
 
-        cast.failure_behavior = algebra_pb2.Expression.Cast.FAILURE_BEHAVIOR_RETURN_NULL
+        # alternative: FAILURE_BEHAVIOR_RETURN_NULL not supported by acero
+        cast.failure_behavior = algebra_pb2.Expression.Cast.FAILURE_BEHAVIOR_THROW_EXCEPTION
         return pb_expr
 
     def _get_expression_dtype(

packages/bigframes/bigframes/core/rewrite/__init__.py

@@ -25,6 +25,7 @@
     try_reduce_to_local_scan,
     try_reduce_to_table_scan,
 )
+from bigframes.core.rewrite.schema_binding import bind_schema_to_tree
 from bigframes.core.rewrite.select_pullup import defer_selection
 from bigframes.core.rewrite.slices import pull_out_limit, pull_up_limits, rewrite_slice
 from bigframes.core.rewrite.timedeltas import rewrite_timedelta_expressions
@@ -37,6 +38,7 @@
 
 __all__ = [
     "as_sql_nodes",
+    "bind_schema_to_tree",
     "extract_ctes",
     "legacy_join_as_projection",
     "try_row_join",

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

@@ -24,7 +24,17 @@
 REFERENCE_ENGINE = polars_executor.PolarsExecutor()
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+@pytest.mark.parametrize(
+    "engine",
+    [
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+    indirect=True,
+)
 def test_engines_filter_bool_col(
     scalars_array_value: array_value.ArrayValue,
     engine,
@@ -35,7 +45,17 @@ def test_engines_filter_bool_col(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+@pytest.mark.parametrize(
+    "engine",
+    [
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+    indirect=True,
+)
 def test_engines_filter_expr_cond(
     scalars_array_value: array_value.ArrayValue,
     engine,
@@ -47,7 +67,17 @@ def test_engines_filter_expr_cond(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+@pytest.mark.parametrize(
+    "engine",
+    [
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+    indirect=True,
+)
 def test_engines_filter_true(
     scalars_array_value: array_value.ArrayValue,
     engine,
@@ -57,7 +87,17 @@ def test_engines_filter_true(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+@pytest.mark.parametrize(
+    "engine",
+    [
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+    indirect=True,
+)
 def test_engines_filter_false(
     scalars_array_value: array_value.ArrayValue,
     engine,