more work

Author: TrevorBergeron

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

@@ -39,30 +39,13 @@ class SubstraitCompiler:
     """
     Compiles BigFrameNode plans to Substrait schema (JSON representation).
     """
-
-    def _print_node_tree(self, node: bigframe_node.BigFrameNode, indent: int = 0):
-        import sys
-        try:
-            ids = list(node.ids)
-        except Exception as e:
-            ids = f"<error: {e}>"
-        sys.stderr.write("  " * indent + f"- {type(node).__name__}: ids={ids}\n")
-        sys.stderr.flush()
-        for child in node.child_nodes:
-             self._print_node_tree(child, indent + 1)
-
     def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
         """
         Compiles a BigFrameNode to Substrait bytes (JSON encoded via protobuf).
         """
         if not self.can_compile(plan):
             return None
 
-        import sys
-        sys.stderr.write("DEBUG TREE:\n")
-        sys.stderr.flush()
-        self._print_node_tree(plan)
-
         pb_rel = self._compile_node(plan)
 
         pb_plan = plan_pb2.Plan()
@@ -84,7 +67,6 @@ def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
     def can_compile(self, plan: bigframe_node.BigFrameNode) -> bool:
         """
         Checks if the plan can be compiled to Substrait.
-        For the skeleton, we support ReadLocalNode, SelectionNode, and FilterNode.
         """
         supported_nodes = (
             nodes.ReadLocalNode,
@@ -95,7 +77,6 @@ def can_compile(self, plan: bigframe_node.BigFrameNode) -> bool:
             nodes.JoinNode,
             nodes.AggregateNode,
             nodes.OrderByNode,
-            nodes.PromoteOffsetsNode,
             nodes.WindowOpNode,
             nodes.ConcatNode,
         )
@@ -123,8 +104,6 @@ def _compile_node(self, node: bigframe_node.BigFrameNode) -> algebra_pb2.Rel:
              return self._compile_orderby(node)
         elif isinstance(node, nodes.SliceNode):
              return self._compile_slice(node)
-        elif isinstance(node, nodes.PromoteOffsetsNode):
-             return self._compile_promote_offsets(node)
         elif isinstance(node, nodes.WindowOpNode):
              return self._compile_window(node)
         elif isinstance(node, nodes.ConcatNode):
@@ -180,22 +159,6 @@ def _compile_selection(self, node: nodes.SelectionNode) -> algebra_pb2.Rel:
 
         return rel
 
-    def _compile_promote_offsets(self, node: nodes.PromoteOffsetsNode) -> algebra_pb2.Rel:
-        input_rel = self._compile_node(node.child)
-        
-        rel = algebra_pb2.Rel()
-        project_rel = rel.project
-        project_rel.input.CopyFrom(input_rel)
-        
-        # Add a dummy literal i64 = 0 for the offsets column
-        expr = project_rel.expressions.add()
-        expr.literal.i64 = 0
-        
-        child_ids = list(node.child.ids)
-        project_rel.common.emit.output_mapping.extend(range(len(child_ids) + 1))
-        
-        return rel
-
     def _compile_filter(self, node: nodes.FilterNode) -> algebra_pb2.Rel:
         input_rel = self._compile_node(node.child)
 
@@ -684,6 +647,10 @@ def _compile_aggregate(self, node: nodes.AggregateNode) -> algebra_pb2.Rel:
                   func_ref = self._EXTENSIONS["product"]
              elif isinstance(agg.op, agg_ops.MedianOp):
                   func_ref = self._EXTENSIONS["median"]
+             elif isinstance(agg.op, agg_ops.CovOp):
+                  func_ref = self._EXTENSIONS["cov"]
+             elif isinstance(agg.op, agg_ops.CorrOp):
+                  func_ref = self._EXTENSIONS["corr"]
              else:
                   raise NotImplementedError(f"Aggregation {type(agg.op)} not supported in Substrait compiler yet")
 
@@ -846,6 +813,9 @@ def _compile_slice(self, node: nodes.SliceNode) -> algebra_pb2.Rel:
         "lead": 66,
         "struct": 67,
         "get_field": 68,
+        "pow": 69,
+        "cov": 70,
+        "corr": 71,
     }
 
     _OP_TO_EXTENSION = {
@@ -854,6 +824,8 @@ def _compile_slice(self, node: nodes.SliceNode) -> algebra_pb2.Rel:
         numeric_ops.MulOp: "multiply",
         numeric_ops.DivOp: "divide",
         numeric_ops.ModOp: "mod",
+        numeric_ops.PowOp: "pow",
+        numeric_ops.UnsafePowOp: "pow",
         comparison_ops.EqOp: "equal",
         comparison_ops.NeOp: "not_equal",
         comparison_ops.LtOp: "lt",
@@ -1156,6 +1128,8 @@ def _compile_fillna_op(self, op: generic_ops.FillNaOp, inputs: Sequence[ex.Expre
     @_compile_op.register(numeric_ops.AddOp)
     @_compile_op.register(numeric_ops.SubOp)
     @_compile_op.register(numeric_ops.MulOp)
+    @_compile_op.register(numeric_ops.PowOp)
+    @_compile_op.register(numeric_ops.UnsafePowOp)
     @_compile_op.register(comparison_ops.EqOp)
     @_compile_op.register(comparison_ops.NeOp)
     @_compile_op.register(comparison_ops.LtOp)

packages/bigframes/bigframes/session/substrait_executor.py

@@ -21,6 +21,7 @@
 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
@@ -62,25 +63,14 @@ def consume(self, plan_proto: bytes, tables: dict[str, pa.Table]) -> pa.Table:
                 "Install it with `pip install datafusion`."
             )
 
-        # Create a DataFusion context
         ctx = datafusion.SessionContext()
 
         for name, table in tables.items():
              df = ctx.from_arrow_table(table)
              ctx.register_table(name, df)
 
-        # NOTE: The actual API for running Substrait in DataFusion python bindings may vary.
-        # Assuming something like ctx.from_substrait(plan) or ctx.execute_substrait(plan).
-        # We will need to verify this with the actual datafusion python package if available.
-        # For now, we raise NotImplementedError if we cannot find the method, or try a likely one.
-        
         import datafusion.substrait
 
-        import substrait.plan_pb2 as plan_pb2
-        from google.protobuf import json_format
-        plan_obj = plan_pb2.Plan.FromString(plan_proto)
-        print("DEBUG PLAN JSON:")
-        print(json_format.MessageToJson(plan_obj))
         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)
@@ -98,62 +88,14 @@ def __init__(self, consumer: SubstraitConsumer):
         from bigframes.core.compile.substrait.compiler import SubstraitCompiler
         self._compiler = SubstraitCompiler()
 
-    def execute(
+    async def execute(
         self,
         plan: bigframe_node.BigFrameNode,
         ordered: bool,
         peek: Optional[int] = None,
     ) -> Optional[executor.ExecuteResult]:
-        def resolve_promote_offsets(node: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:
-             if isinstance(node, nodes.PromoteOffsetsNode):
-                  res = self.execute(node.child, ordered=ordered)
-                  if res is None:
-                       return node
-                  table = res.batches().to_arrow_table()
-                  import pyarrow as pa
-                  table = table.append_column(node.col_id.name, pa.array(range(len(table)), type=pa.int64()))
-                  
-                  from bigframes.core import local_data, identifiers
-                  from bigframes.core.schema import ArraySchema, SchemaItem
-                  import bigframes.dtypes
-                  
-                  schema_items = []
-                  for col_name in table.column_names:
-                       if col_name == node.col_id.name:
-                            schema_items.append(SchemaItem(col_name, bigframes.dtypes.INT_DTYPE))
-                       else:
-                            schema_items.append(SchemaItem(col_name, node.child.schema.get_type(col_name)))
-                  new_schema = ArraySchema(tuple(schema_items))
-                  
-                  scan_items = []
-                  for col_name in table.column_names:
-                       col_id = identifiers.ColumnId(col_name)
-                       scan_items.append(nodes.ScanItem(col_id, col_name))
-                  scan_list = nodes.ScanList(tuple(scan_items))
-                  
-                  session = None
-                  for child_node in node.child.unique_nodes():
-                       if isinstance(child_node, nodes.ReadLocalNode):
-                            session = child_node.session
-                            break
-                  
-                  managed_table = local_data.ManagedArrowTable.from_pyarrow(table, schema=new_schema)
-                  new_node = nodes.ReadLocalNode(
-                       local_data_source=managed_table,
-                       scan_list=scan_list,
-                       session=session,
-                       offsets_col=None,
-                  )
-                  return new_node
-             return node
-
-        # 1. Rewrite all SliceNodes to standard Selection/Filter/Projection/PromoteOffsetsNodes
         plan = plan.bottom_up(slices_rewrite.rewrite_slice)
 
-        # 2. Resolve all PromoteOffsetsNodes to concrete local tables
-        plan = plan.bottom_up(resolve_promote_offsets)
-
-        # 3. Wrap plan in a ResultNode to apply defer_order
         from bigframes.core import expression, rewrite
         output_cols = tuple((expression.DerefOp(id), id.name) for id in plan.ids)
         result_node = nodes.ResultNode(
@@ -166,14 +108,12 @@ def resolve_promote_offsets(node: bigframe_node.BigFrameNode) -> bigframe_node.B
 
         rewritten_plan = result_node.child
 
-        # 4. Apply outermost sorting if ordered
         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),
             )
 
-        # 5. Project only the original output columns to preserve correct result schema
         original_ids = tuple(id for id in plan.ids)
         if rewritten_plan.ids != original_ids:
             rewritten_plan = nodes.SelectionNode(
@@ -188,17 +128,6 @@ def resolve_promote_offsets(node: bigframe_node.BigFrameNode) -> bigframe_node.B
         if substrait_plan_proto is None:
             return None
 
-        import google.protobuf.json_format as json_format
-        from substrait.plan_pb2 import Plan
-        plan_proto = Plan()
-        plan_proto.ParseFromString(substrait_plan_proto)
-        import os
-        import uuid
-        os.makedirs("/usr/local/google/home/tbergeron/src/google-cloud-python/packages/bigframes/scratch", exist_ok=True)
-        filename = f"/usr/local/google/home/tbergeron/src/google-cloud-python/packages/bigframes/scratch/plan_{rewritten_plan.__class__.__name__}_{uuid.uuid4().hex[:8]}.json"
-        with open(filename, "w") as f:
-             f.write(json_format.MessageToJson(plan_proto))
-
         tables = {}
         for node in rewritten_plan.unique_nodes():
              if isinstance(node, nodes.ReadLocalNode):
@@ -211,52 +140,9 @@ def resolve_promote_offsets(node: bigframe_node.BigFrameNode) -> bigframe_node.B
                        table = pyarrow_utils.append_offsets(table, node.offsets_col.sql)
                   tables[table_name] = table
 
-        pa_table = self._consumer.consume(substrait_plan_proto, tables)
-
-        # Sanitize pa_table: replace inf/nan/is_inf with null for INT_DTYPE columns
-        import pyarrow.compute as pc
-        import bigframes.dtypes as dtypes
-        import pyarrow as pa
-        sanitized_columns = []
-        for col_name in pa_table.column_names:
-             col_data = pa_table.column(col_name)
-             try:
-                  expected_dtype = rewritten_plan.schema.get_type(col_name)
-             except ValueError:
-                  expected_dtype = None
-             
-             if expected_dtype == dtypes.INT_DTYPE and pa.types.is_floating(col_data.type):
-                  is_nan = pc.is_nan(col_data)
-                  is_inf = pc.is_inf(col_data)
-                  is_invalid = pc.or_(is_nan, is_inf)
-                  null_val = pa.scalar(None, type=col_data.type)
-                  col_data = pc.if_else(is_invalid, null_val, col_data)
-             sanitized_columns.append(col_data)
-        pa_table = pa.Table.from_arrays(sanitized_columns, names=pa_table.column_names)
-
-        # Handle SliceNode post-processing
-        for node in rewritten_plan.unique_nodes():
-             if isinstance(node, nodes.SliceNode):
-                  is_simple = (node.start is None or node.start >= 0) and (node.stop is None or node.stop >= 0) and (node.step is None or node.step == 1)
-                  if not is_simple:
-                       df = pa_table.to_pandas()
-                       df = df.iloc[node.start:node.stop:node.step]
-                       pa_table = pa.Table.from_pandas(df, schema=pa_table.schema)
-        offset_cols = set()
-        for node in rewritten_plan.unique_nodes():
-             if isinstance(node, nodes.PromoteOffsetsNode):
-                  offset_cols.add(node.col_id.name)
-                  
-        for col_name in pa_table.column_names:
-             if col_name in offset_cols:
-                  idx = pa_table.column_names.index(col_name)
-                  pa_table = pa_table.set_column(idx, col_name, pa.array(range(len(pa_table)), type=pa.int64()))
-
-        import sys
-        sys.stderr.write(f"PA_TABLE ON EXECUTE:\n{pa_table.to_pandas()}\n")
-        sys.stderr.flush()
+        pa_table = await asyncio.to_thread(self._consumer.consume, substrait_plan_proto, tables)
 
-        if peek is not None:
+        if peek is not None:    
             pa_table = pa_table.slice(0, peek)
 
         return executor.LocalExecuteResult(

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

@@ -26,6 +26,7 @@
     local_scan_executor,
     polars_executor,
     semi_executor,
+    substrait_executor,
 )
 
 CURRENT_DIR = pathlib.Path(__file__).parent
@@ -81,9 +82,17 @@ def sqlglot_engine(
     )
 
 
-@pytest.fixture(scope="session", params=["pyarrow", "polars", "bq", "bq-sqlglot"])
+@pytest.fixture(scope="session")
+def substrait_datafusion_engine(
+) -> semi_executor.SemiExecutor:
+    return substrait_executor.SubstraitExecutor(
+        consumer = substrait_executor.DataFusionSubstraitConsumer()
+    )
+
+
+@pytest.fixture(scope="session", params=["pyarrow", "polars", "bq", "bq-sqlglot", "substrait-datafusion"])
 def engine(
-    request, pyarrow_engine, polars_engine, bq_engine, sqlglot_engine
+    request, pyarrow_engine, polars_engine, bq_engine, sqlglot_engine, substrait_datafusion_engine
 ) -> semi_executor.SemiExecutor:
     if request.param == "pyarrow":
         return pyarrow_engine
@@ -93,6 +102,8 @@ def engine(
         return bq_engine
     if request.param == "bq-sqlglot":
         return sqlglot_engine
+    if request.param == "substrait-datafusion":
+        return substrait_datafusion_engine
     raise ValueError(f"Unrecognized param: {request.param}")
 
 

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

@@ -55,7 +55,7 @@ def apply_agg_to_all_valid(
     return new_arr
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], 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,