add more ops to substrait compiler

Author: TrevorBergeron

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

@@ -53,12 +53,15 @@ def compile(self, plan: bigframe_node.BigFrameNode) -> Optional[bytes]:
             ("sub", 2),
             ("mul", 3),
             ("div", 4),
-            ("eq", 5),
+            ("equal", 5),
             ("ne", 6),
             ("lt", 7),
             ("gt", 8),
             ("le", 9),
             ("ge", 10),
+            ("sum", 11),
+            ("max", 12),
+            ("and", 13),
         ]
         for name, anchor in extensions:
              ext = pb_plan.extensions.add()
@@ -90,8 +93,6 @@ def _compile_node(self, node: bigframe_node.BigFrameNode) -> Dict[str, Any]:
             return self._compile_selection(node)
         elif isinstance(node, nodes.FilterNode):
             return self._compile_filter(node)
-        elif isinstance(node, nodes.SliceNode):
-            return self._compile_slice(node)
         elif isinstance(node, nodes.ProjectionNode):
             return self._compile_projection(node)
         elif isinstance(node, nodes.JoinNode):
@@ -114,43 +115,35 @@ def _compile_read(self, node: nodes.ReadLocalNode) -> Dict[str, Any]:
         return json_format.MessageToDict(rel, preserving_proto_field_name=True)
 
     def _compile_selection(self, node: nodes.SelectionNode) -> Dict[str, Any]:
-        # Selection usually maps to ProjectRel or FilterRel depending on if it filters or just selects columns.
-        # If it's just column selection (Projection), it's a ProjectRel.
-        # Let's assume it's a ProjectRel for now.
         input_rel = self._compile_node(node.child)
+        expressions = []
+        child_ids = list(node.child.ids)
+        for aliased_ref in node.input_output_pairs:
+             source_id = aliased_ref.ref.id
+             idx = child_ids.index(source_id)
+             expressions.append({"selection": {"direct_reference": {"struct_field": {"field": idx}}}})
         return {
             "project": {
+                "common": {
+                    "emit": {
+                        "outputMapping": [len(child_ids) + i for i in range(len(expressions))]
+                    }
+                },
                 "input": input_rel,
-                "expressions": [
-                    # Skeletal expression mapping
-                    {"selection": {"direct_reference": {"struct_field": {"field": i}}}} 
-                    for i in range(len(node.schema))
-                ]
+                "expressions": expressions
             }
         }
 
     def _compile_filter(self, node: nodes.FilterNode) -> Dict[str, Any]:
         input_rel = self._compile_node(node.child)
-        condition_rel = self._compile_expression(node.condition, node.child)
+        condition_rel = self._compile_expression(node.predicate, node.child)
         return {
             "filter": {
                 "input": input_rel,
                 "condition": condition_rel
             }
         }
 
-    def _compile_slice(self, node: nodes.SliceNode) -> Dict[str, Any]:
-        input_rel = self._compile_node(node.child)
-        count = node.stop if node.stop is not None else -1
-        offset = node.start if node.start is not None else 0
-        
-        return {
-            "fetch": {
-                "input": input_rel,
-                "offset": offset,
-                "count": count
-            }
-        }
 
     def _compile_projection(self, node: nodes.ProjectionNode) -> Dict[str, Any]:
         input_rel_dict = self._compile_node(node.child)
@@ -193,7 +186,7 @@ def _compile_join(self, node: nodes.JoinNode) -> Dict[str, Any]:
 
              eq_expressions.append({
                  "scalar_function": {
-                     "function_reference": 0,
+                     "function_reference": 5, # eq
                      "arguments": [
                          {"value": {"selection": {"direct_reference": {"struct_field": {"field": left_idx}}}}},
                          {"value": {"selection": {"direct_reference": {"struct_field": {"field": right_idx}}}}}
@@ -203,6 +196,13 @@ def _compile_join(self, node: nodes.JoinNode) -> Dict[str, Any]:
 
         if len(eq_expressions) > 1:
              expr = eq_expressions[0]
+             for e in eq_expressions[1:]:
+                  expr = {
+                      "scalar_function": {
+                          "function_reference": 13, # and
+                          "arguments": [{"value": expr}, {"value": e}]
+                      }
+                  }
         elif len(eq_expressions) == 1:
              expr = eq_expressions[0]
         else:
@@ -229,8 +229,14 @@ def _compile_aggregate(self, node: nodes.AggregateNode) -> Dict[str, Any]:
              groupings.append({"grouping_expressions": grouping_expressions})
 
         measures = []
+        import bigframes.operations.aggregations as agg_ops
         for agg, _ in node.aggregations:
-             func_ref = 1 if "Sum" in type(agg).__name__ else 2
+             if isinstance(agg.op, agg_ops.SumOp):
+                  func_ref = 11
+             elif isinstance(agg.op, agg_ops.MaxOp):
+                  func_ref = 12
+             else:
+                  raise NotImplementedError(f"Aggregation {type(agg.op)} not supported in Substrait compiler yet")
              args = []
              if hasattr(agg, "column_references"):
                   for col_id in agg.column_references:

packages/bigframes/tests/unit/session/test_substrait_executor.py

@@ -133,3 +133,184 @@ def test_execute_projection_add_with_datafusion():
     assert "b" in result_table.column_names
     assert "a" in result_table.column_names
     assert result_table.column("b").to_pylist() == [43, 44, 45]
+
+
+def test_execute_filter_with_datafusion():
+    from bigframes.session.substrait_executor import DataFusionSubstraitConsumer
+    from bigframes.operations.comparison_ops import gt_op
+    
+    consumer = DataFusionSubstraitConsumer()
+    executor = substrait_executor.SubstraitExecutor(consumer)
+    
+    read_node = create_read_local_node()
+    
+    # a > 1
+    filter_expr = ex.OpExpression(
+        op=gt_op,
+        inputs=(
+            ex.DerefOp(identifiers.ColumnId("a")),
+            ex.ScalarConstantExpression(1),
+        ),
+    )
+    plan = nodes.FilterNode(
+        child=read_node,
+        predicate=filter_expr,
+    )
+    
+    result = executor.execute(plan, ordered=True)
+    assert result is not None
+    
+    result_table = pa.Table.from_batches(result.batches().arrow_batches)
+    assert result_table.num_rows == 2
+    assert "a" in result_table.column_names
+    assert result_table.column("a").to_pylist() == [2, 3]
+
+
+def test_execute_aggregate_sum_with_datafusion():
+    from bigframes.session.substrait_executor import DataFusionSubstraitConsumer
+    from bigframes.operations.aggregations import sum_op
+    from bigframes.core.agg_expressions import UnaryAggregation
+    
+    consumer = DataFusionSubstraitConsumer()
+    executor = substrait_executor.SubstraitExecutor(consumer)
+    
+    read_node = create_read_local_node()
+    
+    # sum(a)
+    sum_agg = UnaryAggregation(
+        op=sum_op,
+        arg=ex.DerefOp(identifiers.ColumnId("a")),
+    )
+    
+    plan = nodes.AggregateNode(
+        child=read_node,
+        aggregations=((sum_agg, identifiers.ColumnId("sum_a")),),
+        by_column_ids=(),
+    )
+    
+    result = executor.execute(plan, ordered=True)
+    assert result is not None
+    
+    result_table = pa.Table.from_batches(result.batches().arrow_batches)
+    assert result_table.num_rows == 1
+    assert "sum_a" in result_table.column_names
+    assert result_table.column("sum_a").to_pylist() == [6]
+
+
+def test_execute_aggregate_max_with_datafusion():
+    from bigframes.session.substrait_executor import DataFusionSubstraitConsumer
+    from bigframes.operations.aggregations import max_op
+    from bigframes.core.agg_expressions import UnaryAggregation
+    
+    consumer = DataFusionSubstraitConsumer()
+    executor = substrait_executor.SubstraitExecutor(consumer)
+    
+    read_node = create_read_local_node()
+    
+    # max(a)
+    max_agg = UnaryAggregation(
+        op=max_op,
+        arg=ex.DerefOp(identifiers.ColumnId("a")),
+    )
+    
+    plan = nodes.AggregateNode(
+        child=read_node,
+        aggregations=((max_agg, identifiers.ColumnId("max_a")),),
+        by_column_ids=(),
+    )
+    
+    result = executor.execute(plan, ordered=True)
+    assert result is not None
+    
+    result_table = pa.Table.from_batches(result.batches().arrow_batches)
+    assert result_table.num_rows == 1
+    assert "max_a" in result_table.column_names
+    assert result_table.column("max_a").to_pylist() == [3]
+
+
+def test_execute_join_with_datafusion():
+    from bigframes.session.substrait_executor import DataFusionSubstraitConsumer
+    
+    consumer = DataFusionSubstraitConsumer()
+    executor = substrait_executor.SubstraitExecutor(consumer)
+    
+    # Table 1: a
+    session1 = mocks.create_bigquery_session()
+    table1 = pa.Table.from_pydict({"a": [1, 2, 3]})
+    source1 = local_data.ManagedArrowTable.from_pyarrow(table1)
+    col_id_a = identifiers.ColumnId("a")
+    read_node1 = nodes.ReadLocalNode(
+        local_data_source=source1,
+        session=session1,
+        scan_list=nodes.ScanList(items=(nodes.ScanItem(id=col_id_a, source_id="a"),)),
+    )
+    
+    # Table 2: b
+    session2 = mocks.create_bigquery_session()
+    table2 = pa.Table.from_pydict({"b": [2, 3, 4]})
+    source2 = local_data.ManagedArrowTable.from_pyarrow(table2)
+    col_id_b = identifiers.ColumnId("b")
+    read_node2 = nodes.ReadLocalNode(
+        local_data_source=source2,
+        session=session2,
+        scan_list=nodes.ScanList(items=(nodes.ScanItem(id=col_id_b, source_id="b"),)),
+    )
+    
+    # Join on a = b
+    join_node = nodes.JoinNode(
+        left_child=read_node1,
+        right_child=read_node2,
+        conditions=((ex.DerefOp(col_id_a), ex.DerefOp(col_id_b)),),
+        type="inner",
+        propogate_order=False,
+    )
+    
+    result = executor.execute(join_node, ordered=True)
+    assert result is not None
+    
+    result_table = pa.Table.from_batches(result.batches().arrow_batches)
+    assert result_table.num_rows == 2
+    assert "a" in result_table.column_names
+    assert "b" in result_table.column_names
+    assert result_table.column("a").to_pylist() == [2, 3]
+    assert result_table.column("b").to_pylist() == [2, 3]
+
+
+def test_execute_selection_with_datafusion():
+    from bigframes.session.substrait_executor import DataFusionSubstraitConsumer
+    from bigframes.core.nodes import AliasedRef
+    
+    consumer = DataFusionSubstraitConsumer()
+    executor = substrait_executor.SubstraitExecutor(consumer)
+    
+    # Table with a and b
+    session = mocks.create_bigquery_session()
+    table = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
+    source = local_data.ManagedArrowTable.from_pyarrow(table)
+    col_id_a = identifiers.ColumnId("a")
+    col_id_b = identifiers.ColumnId("b")
+    read_node = nodes.ReadLocalNode(
+        local_data_source=source,
+        session=session,
+        scan_list=nodes.ScanList(
+            items=(
+                nodes.ScanItem(id=col_id_a, source_id="a"),
+                nodes.ScanItem(id=col_id_b, source_id="b"),
+            )
+        ),
+    )
+    
+    # Select only a, and rename it to c
+    col_id_c = identifiers.ColumnId("c")
+    selection_node = nodes.SelectionNode(
+        child=read_node,
+        input_output_pairs=(AliasedRef(ex.DerefOp(col_id_a), col_id_c),),
+    )
+    
+    result = executor.execute(selection_node, ordered=True)
+    assert result is not None
+    
+    result_table = pa.Table.from_batches(result.batches().arrow_batches)
+    assert result_table.num_rows == 3
+    assert result_table.column_names == ["c"]
+    assert result_table.column("c").to_pylist() == [1, 2, 3]