simplify, add acero engine

Author: TrevorBergeron

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

@@ -15,13 +15,13 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Optional
+import asyncio
+from typing import TYPE_CHECKING, Optional, cast
 
-from bigframes.core import bigframe_node
+import bigframes.core.compile.substrait.compiler as substrait_compiler
+import bigframes.core.rewrite as rewrite
+from bigframes.core import bigframe_node, nodes
 from bigframes.session import executor, semi_executor
-import bigframes.core.rewrite.slices as slices_rewrite
-from bigframes.core import nodes
-import asyncio
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -66,49 +66,93 @@ def consume(self, plan_proto: bytes, tables: dict[str, pa.Table]) -> pa.Table:
         ctx = datafusion.SessionContext()
 
         for name, table in tables.items():
-             df = ctx.from_arrow_table(table)
-             ctx.register_table(name, df)
-        
+            df = ctx.from_arrow_table(table)
+            ctx.register_table(name, df)
+
         import datafusion.substrait
 
-        datafusion_substrait_plan = datafusion.substrait.Serde.deserialize_bytes(plan_proto)
-        logical_plan = datafusion.substrait.Consumer.from_substrait_plan(ctx, datafusion_substrait_plan)
+        datafusion_substrait_plan = datafusion.substrait.Serde.deserialize_bytes(
+            plan_proto
+        )
+        logical_plan = datafusion.substrait.Consumer.from_substrait_plan(
+            ctx, datafusion_substrait_plan
+        )
         df = ctx.create_dataframe_from_logical_plan(logical_plan)
         return df.to_arrow_table()
 
 
+class AceroSubstraitConsumer(SubstraitConsumer):
+    """
+    Executes Substrait plans using Apache Arrow Acero.
+    """
+
+    def consume(self, plan_proto: bytes, tables: dict[str, pa.Table]) -> pa.Table:
+        import pyarrow.substrait as pa_substrait
+
+        def provide_table(name: list[str], schema: pa.Schema) -> pa.Table:
+            return tables[name[0]]
+
+        batch_reader = pa_substrait.run_query(plan_proto, table_provider=provide_table)
+        return batch_reader.read_all()
+
+
 class SubstraitExecutor(semi_executor.SemiExecutor):
     """
     Executes plans by compiling them to Substrait and running them via a consumer.
     """
 
-    def __init__(self, consumer: SubstraitConsumer):
+    def __init__(
+        self,
+        consumer: SubstraitConsumer,
+        compiler: substrait_compiler.SubstraitCompiler,
+    ):
         self._consumer = consumer
-        # Lazy import to avoid circular dependencies
-        from bigframes.core.compile.substrait.compiler import SubstraitCompiler
-        self._compiler = SubstraitCompiler()
+        self._compiler = compiler
+
+    @classmethod
+    def default_for_engine(cls, engine_name: str) -> SubstraitExecutor:
+        if engine_name == "acero":
+            return cls(
+                AceroSubstraitConsumer(),
+                substrait_compiler.SubstraitCompiler(
+                    duration_type="int", use_precision_types=False
+                ),
+            )
+        elif engine_name == "datafusion":
+            return cls(
+                DataFusionSubstraitConsumer(),
+                substrait_compiler.SubstraitCompiler(duration_type="int"),
+            )
+        else:
+            raise ValueError(f"Unknown engine: {engine_name}")
 
     async def execute(
         self,
         plan: bigframe_node.BigFrameNode,
         ordered: bool,
         peek: Optional[int] = None,
     ) -> Optional[executor.ExecuteResult]:
-        plan = plan.bottom_up(slices_rewrite.rewrite_slice)
+        plan = plan.bottom_up(rewrite.rewrite_slice)
+        # Only needed for acero technically, datafusion can handle timedeltas
+        plan = plan.bottom_up(rewrite.rewrite_timedelta_expressions)
+
+        from bigframes.core import expression
 
-        from bigframes.core import expression, rewrite
         output_cols = tuple((expression.DerefOp(id), id.name) for id in plan.ids)
         result_node = nodes.ResultNode(
             plan,
             output_cols=output_cols,
         )
-        import typing
-        result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
+        result_node = cast(nodes.ResultNode, rewrite.column_pruning(result_node))
         result_node = rewrite.defer_order(result_node, output_hidden_row_keys=False)
 
         rewritten_plan = result_node.child
 
-        if ordered and result_node.order_by and result_node.order_by.all_ordering_columns:
+        if (
+            ordered
+            and result_node.order_by
+            and result_node.order_by.all_ordering_columns
+        ):
             rewritten_plan = nodes.OrderByNode(
                 rewritten_plan,
                 by=tuple(result_node.order_by.all_ordering_columns),
@@ -118,7 +162,9 @@ async def execute(
         if rewritten_plan.ids != original_ids:
             rewritten_plan = nodes.SelectionNode(
                 rewritten_plan,
-                input_output_pairs=tuple(nodes.AliasedRef.identity(id) for id in original_ids)
+                input_output_pairs=tuple(
+                    nodes.AliasedRef.identity(id) for id in original_ids
+                ),
             )
 
         if not self._can_execute(rewritten_plan):
@@ -130,19 +176,24 @@ async def execute(
 
         tables = {}
         for node in rewritten_plan.unique_nodes():
-             if isinstance(node, nodes.ReadLocalNode):
-                  table_name = f"table_{id(node)}"
-                  table = node.local_data_source.data
-                  table = table.select([item.source_id for item in node.scan_list.items])
-                  table = table.rename_columns([item.id.sql for item in node.scan_list.items])
-                  if node.offsets_col is not None:
-                       from bigframes.core import pyarrow_utils
-                       table = pyarrow_utils.append_offsets(table, node.offsets_col.sql)
-                  tables[table_name] = table
-
-        pa_table = await asyncio.to_thread(self._consumer.consume, substrait_plan_proto, tables)
-
-        if peek is not None:    
+            if isinstance(node, nodes.ReadLocalNode):
+                table_name = f"table_{id(node)}"
+                table = node.local_data_source.to_pyarrow_table(duration_type="int")
+                table = table.select([item.source_id for item in node.scan_list.items])
+                table = table.rename_columns(
+                    [item.id.sql for item in node.scan_list.items]
+                )
+                if node.offsets_col is not None:
+                    from bigframes.core import pyarrow_utils
+
+                    table = pyarrow_utils.append_offsets(table, node.offsets_col.sql)
+                tables[table_name] = table
+
+        pa_table = await asyncio.to_thread(
+            self._consumer.consume, substrait_plan_proto, tables
+        )
+
+        if peek is not None:
             pa_table = pa_table.slice(0, peek)
 
         return executor.LocalExecuteResult(

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

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-import dataclasses
+import asyncio
 import weakref
 from typing import TYPE_CHECKING, Union
 
@@ -31,11 +31,13 @@
     import bigframes.core
 
 
-@dataclasses.dataclass
 class SubstraitTestExecutor(bigframes.session.executor.Executor):
-    def __init__(self):
-        from bigframes.session.substrait_executor import DataFusionSubstraitConsumer, SubstraitExecutor
-        self.executor = SubstraitExecutor(DataFusionSubstraitConsumer())
+    def __init__(
+        self, consumer: bigframes.session.substrait_executor.SubstraitConsumer
+    ):
+        from bigframes.session.substrait_executor import SubstraitExecutor
+
+        self.executor = SubstraitExecutor(consumer)
 
     def execute(
         self,
@@ -46,11 +48,15 @@ def execute(
             raise ValueError(
                 f"SubstraitTestExecutor does not support destination spec: {execution_spec.destination_spec}"
             )
-        
-        result = self.executor.execute(array_value.node, ordered=True, peek=execution_spec.peek)
+
+        result = asyncio.run(
+            self.executor.execute(
+                array_value.node, ordered=True, peek=execution_spec.peek
+            )
+        )
         if result is None:
             raise NotImplementedError("SubstraitExecutor cannot execute this plan")
-            
+
         return result
 
     def cached(
@@ -63,7 +69,7 @@ def cached(
 
 
 class TestSession(bigframes.session.Session):
-    def __init__(self):
+    def __init__(self, executor: SubstraitTestExecutor):
         self._location = None  # type: ignore
         self._bq_kms_key_name = None  # type: ignore
         self._clients_provider = None  # type: ignore
@@ -85,7 +91,7 @@ def __init__(self):
         self._metrics = bigframes.session.metrics.ExecutionMetrics()
         self._function_session = None  # type: ignore
         self._temp_storage_manager = None  # type: ignore
-        self._executor = SubstraitTestExecutor()
+        self._executor = executor
         self._loader = None  # type: ignore
 
     def read_pandas(self, pandas_dataframe, write_engine="default"):

packages/bigframes/bigframes/session/substrait_executor.py

@@ -82,16 +82,34 @@ def sqlglot_engine(
 
 
 @pytest.fixture(scope="session")
-def substrait_datafusion_engine(
-) -> semi_executor.SemiExecutor:
-    return substrait_executor.SubstraitExecutor(
-        consumer = substrait_executor.DataFusionSubstraitConsumer()
-    )
+def substrait_datafusion_engine() -> semi_executor.SemiExecutor:
+    return substrait_executor.SubstraitExecutor.default_for_engine("datafusion")
 
 
-@pytest.fixture(scope="session", params=["pyarrow", "polars", "bq", "bq-sqlglot", "substrait-datafusion"])
+@pytest.fixture(scope="session")
+def substrait_acero_engine() -> semi_executor.SemiExecutor:
+    return substrait_executor.SubstraitExecutor.default_for_engine("acero")
+
+
+@pytest.fixture(
+    scope="session",
+    params=[
+        "pyarrow",
+        "polars",
+        "bq",
+        "bq-sqlglot",
+        "substrait-datafusion",
+        "substrait-acero",
+    ],
+)
 def engine(
-    request, pyarrow_engine, polars_engine, bq_engine, sqlglot_engine, substrait_datafusion_engine
+    request,
+    pyarrow_engine,
+    polars_engine,
+    bq_engine,
+    sqlglot_engine,
+    substrait_datafusion_engine,
+    substrait_acero_engine,
 ) -> semi_executor.SemiExecutor:
     if request.param == "pyarrow":
         return pyarrow_engine
@@ -103,6 +121,8 @@ def engine(
         return sqlglot_engine
     if request.param == "substrait-datafusion":
         return substrait_datafusion_engine
+    if request.param == "substrait-acero":
+        return substrait_acero_engine
     raise ValueError(f"Unrecognized param: {request.param}")
 
 

packages/bigframes/bigframes/testing/substrait_session.py

@@ -54,7 +54,9 @@ def apply_agg_to_all_valid(
     return new_arr
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot", "substrait-datafusion"], indirect=True)
+@pytest.mark.parametrize(
+    "engine", ["polars", "bq", "bq-sqlglot", "substrait-datafusion"], indirect=True
+)
 def test_engines_aggregate_post_filter_size(
     scalars_array_value: array_value.ArrayValue,
     engine,