shift to engine tests, fix issues

Author: TrevorBergeron

…igframes/core/compile/polars/lowering.py …rames/bigframes/core/compile/lowering.pypackages/bigframes/bigframes/core/compile/polars/lowering.py renamed to packages/bigframes/bigframes/core/compile/lowering.py packages/bigframes/bigframes/core/compile/polars/lowering.py renamed to packages/bigframes/bigframes/core/compile/lowering.py

@@ -492,11 +492,19 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
     LowerFloorOp(),
 )
 
+SUBSTRAIT_LOWERING_RULES = (*LOWER_COMPARISONS,)
+
 
 def lower_ops_to_polars(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
     return op_lowering.lower_ops(root, rules=POLARS_LOWERING_RULES)
 
 
+def lower_ops_to_substrait(
+    root: bigframe_node.BigFrameNode,
+) -> bigframe_node.BigFrameNode:
+    return op_lowering.lower_ops(root, rules=SUBSTRAIT_LOWERING_RULES)
+
+
 def _numeric_to_timedelta(expr: expression.Expression) -> expression.Expression:
     """rounding logic used for emulating timedelta ops"""
     rounded_value = ops.where_op.as_expr(

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

@@ -31,6 +31,7 @@
 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.compile import lowering
 
 
 class SubstraitCompiler:
@@ -53,6 +54,7 @@ def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
         if not self.can_compile(plan):
             return None
 
+        plan = lowering.lower_ops_to_substrait(plan)
         pb_rel = self._compile_node(plan)
 
         pb_plan = plan_pb2.Plan()

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

@@ -46,7 +46,17 @@ def apply_op_pairwise(
     return new_arr
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+@pytest.mark.parametrize(
+    "engine",
+    [
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+    indirect=True,
+)
 @pytest.mark.parametrize(
     "op",
     [

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

@@ -48,7 +48,17 @@ def apply_op_pairwise(
     return new_arr
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
+@pytest.mark.parametrize(
+    "engine",
+    [
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+    indirect=True,
+)
 @pytest.mark.parametrize(
     "op",
     [

packages/bigframes/tests/system/small/engines/test_slicing.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,
+)
 @pytest.mark.parametrize(
     ("start", "stop", "step"),
     [