fixes

Author: TrevorBergeron

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, rewrite
+from bigframes.core import agg_expressions, bigframe_node, nodes, rewrite
 from bigframes.core.compile import lowering
 
 
@@ -499,12 +499,12 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
                 bounds_type = (
                     algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_RANGE
                 )
-                start = node.window_spec.bounds.start
-                if start is None:
+                range_start = node.window_spec.bounds.start
+                if range_start is None:
                     lower_bound.unbounded.CopyFrom(
                         algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
                     )
-                elif start == pd.Timedelta(0):
+                elif range_start == pd.Timedelta(0):
                     lower_bound.current_row.CopyFrom(
                         algebra_pb2.Expression.WindowFunction.Bound.CurrentRow()
                     )
@@ -513,12 +513,12 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
                         "Range window bounds with non-zero offsets are not supported yet"
                     )
 
-                end = node.window_spec.bounds.end
-                if end is None:
+                range_end = node.window_spec.bounds.end
+                if range_end is None:
                     upper_bound.unbounded.CopyFrom(
                         algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
                     )
-                elif end == pd.Timedelta(0):
+                elif range_end == pd.Timedelta(0):
                     upper_bound.current_row.CopyFrom(
                         algebra_pb2.Expression.WindowFunction.Bound.CurrentRow()
                     )
@@ -540,6 +540,7 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
         # 3. Project each window aggregation expression as a WindowFunction expression
         for agg_idx, col_def in enumerate(node.agg_exprs):
             agg = col_def.expression
+            assert isinstance(agg, agg_expressions.Aggregation)
             distinct = False
 
             if isinstance(agg.op, agg_ops.SumOp):
@@ -798,10 +799,10 @@ def _compile_scalar_constant(
                 pb_expr.literal.precision_timestamp.precision = 6
                 pb_expr.literal.precision_timestamp.value = us
         elif isinstance(val, datetime.date):
-            epoch = datetime.date(1970, 1, 1)
-            days = (val - epoch).days
+            date_epoch = datetime.date(1970, 1, 1)
+            days = (val - date_epoch).days
             pb_expr.literal.date = days
-        elif pd.isna(val):
+        elif pd.isna(val):  # type: ignore[call-overload]
             pb_expr.literal.null.varchar.length = 0
         else:
             pb_expr.literal.string = str(val)
@@ -911,7 +912,7 @@ def _get_expression_dtype(
         import bigframes.dtypes as dtypes
 
         if isinstance(expr, ex.ScalarConstantExpression):
-            if expr.value is None or pd.isna(expr.value):
+            if expr.value is None or pd.isna(expr.value):  # type: ignore[call-overload]
                 return None
             return expr.dtype or dtypes.infer_literal_type(expr.value)
         elif isinstance(expr, ex.DerefOp):
@@ -1110,7 +1111,7 @@ def _compile_isin(
     @_compile_op.register(generic_ops.FillNaOp)
     def _compile_fillna_op(
         self,
-        op: generic_ops.FillNaOp,
+        op: ops.BinaryOp,
         inputs: Sequence[ex.Expression],
         child: nodes.BigFrameNode,
     ) -> algebra_pb2.Expression:
@@ -1237,7 +1238,7 @@ def _compile_standard_unaryops(
     @_compile_op.register(numeric_ops.PosOp)
     def _compile_pos_op(
         self,
-        op: numeric_ops.PosOp,
+        op: ops.UnaryOp,
         inputs: Sequence[ex.Expression],
         child: nodes.BigFrameNode,
     ) -> algebra_pb2.Expression:
@@ -1247,7 +1248,7 @@ def _compile_pos_op(
     @_compile_op.register(numeric_ops.NegOp)
     def _compile_neg_op(
         self,
-        op: numeric_ops.NegOp,
+        op: ops.UnaryOp,
         inputs: Sequence[ex.Expression],
         child: nodes.BigFrameNode,
     ) -> algebra_pb2.Expression:
@@ -1270,7 +1271,7 @@ def _compile_neg_op(
     @_compile_op.register(generic_ops.InvertOp)
     def _compile_invert_op(
         self,
-        op: generic_ops.InvertOp,
+        op: ops.UnaryOp,
         inputs: Sequence[ex.Expression],
         child: nodes.BigFrameNode,
     ) -> algebra_pb2.Expression:

packages/bigframes/bigframes/core/rewrite/substrait_agg.py

@@ -48,7 +48,7 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
     ]
 
     # Collect cast aggregations (bool->float for mean/stddev/var, bool->int for sum)
-    cast_aggs = []
+    cast_aggs: list[tuple[int, identifiers.ColumnId, dtypes.Dtype]] = []
     for agg_idx, (agg, _) in enumerate(node.aggregations):
         if hasattr(agg, "column_references"):
             for col_id in agg.column_references:
@@ -92,6 +92,7 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
         # Rewrite aggregations to use the projected columns
         rewritten_aggs = []
         for agg_idx, (agg, out_col_id) in enumerate(node.aggregations):
+            rewritten_agg: agg_expressions.Aggregation
             if isinstance(agg.op, (agg_ops.SizeOp, agg_ops.SizeUnaryOp)):
                 new_col_id = size_agg_to_col_id[agg_idx]
                 rewritten_agg = agg_expressions.UnaryAggregation(
@@ -107,14 +108,7 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
                     else:
                         new_exprs.append(expression.deref(col_id.name))
 
-                if isinstance(agg, agg_expressions.UnaryAggregation):
-                    rewritten_agg = agg_expressions.UnaryAggregation(
-                        agg.op, new_exprs[0]
-                    )
-                else:
-                    rewritten_agg = agg_expressions.UnaryAggregation(
-                        agg.op, new_exprs[0]
-                    )
+                rewritten_agg = agg.replace_args(*new_exprs)
             else:
                 rewritten_agg = agg
 
@@ -135,9 +129,9 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
 
     assignments = []
     selection_pairs = []
-    for idx, (out_id, target_dtype) in enumerate(zip(output_ids, expected_types)):
+    for idx, (out_id, out_dtype) in enumerate(zip(output_ids, expected_types)):
         cast_id = identifiers.ColumnId(f"bf_out_cast_{out_id.name}")
-        cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(
+        cast_expr = ops.AsTypeOp(to_type=out_dtype).as_expr(
             expression.deref(out_id.name)
         )
         assignments.append((cast_expr, cast_id))
@@ -173,10 +167,11 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
 
     # Collect size and cast requirements for the window agg expressions
     size_aggs = []
-    cast_aggs = []
+    cast_aggs: list[tuple[int, identifiers.ColumnId, dtypes.Dtype]] = []
 
     for agg_idx, col_def in enumerate(node.agg_exprs):
         agg = col_def.expression
+        assert isinstance(agg, agg_expressions.Aggregation)
         if isinstance(agg.op, (agg_ops.SizeOp, agg_ops.SizeUnaryOp)):
             size_aggs.append((agg_idx, agg))
         elif hasattr(agg, "column_references"):
@@ -222,7 +217,9 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         rewritten_agg_exprs = []
         for agg_idx, col_def in enumerate(node.agg_exprs):
             agg = col_def.expression
+            assert isinstance(agg, agg_expressions.Aggregation)
             out_col_id = col_def.id
+            rewritten_agg: agg_expressions.Aggregation
 
             if isinstance(agg.op, (agg_ops.SizeOp, agg_ops.SizeUnaryOp)):
                 new_col_id = size_agg_to_col_id[agg_idx]
@@ -239,14 +236,7 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
                     else:
                         new_exprs.append(expression.deref(col_id.name))
 
-                if isinstance(agg, agg_expressions.UnaryAggregation):
-                    rewritten_agg = agg_expressions.UnaryAggregation(
-                        agg.op, new_exprs[0]
-                    )
-                else:
-                    rewritten_agg = agg_expressions.UnaryAggregation(
-                        agg.op, new_exprs[0]
-                    )
+                rewritten_agg = agg.replace_args(*new_exprs)
             else:
                 rewritten_agg = agg
 
@@ -273,10 +263,10 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
 
     for col_def, field in zip(node.agg_exprs, node.added_fields):
         out_id = col_def.id
-        target_dtype = field.dtype
+        out_dtype = field.dtype
 
         cast_id = identifiers.ColumnId(f"bf_window_out_cast_{out_id.name}")
-        cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(
+        cast_expr = ops.AsTypeOp(to_type=out_dtype).as_expr(
             expression.deref(out_id.name)
         )
         assignments.append((cast_expr, cast_id))

packages/bigframes/bigframes/session/substrait_executor.py

@@ -21,7 +21,7 @@
 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
+from bigframes.session import execution_spec, executor, semi_executor
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -129,9 +129,10 @@ def default_for_engine(cls, engine_name: str) -> SubstraitExecutor:
     async def execute(
         self,
         plan: bigframe_node.BigFrameNode,
-        ordered: bool,
-        peek: Optional[int] = None,
+        execution_spec: execution_spec.ExecutionSpec,
     ) -> Optional[executor.ExecuteResult]:
+        ordered = execution_spec.ordered
+        peek = execution_spec.peek
         plan = plan.bottom_up(rewrite.rewrite_slice)
         # Only needed for acero technically, datafusion can handle timedeltas
         plan = plan.bottom_up(rewrite.rewrite_timedelta_expressions)

packages/bigframes/bigframes/testing/substrait_session.py

@@ -35,9 +35,18 @@ class SubstraitTestExecutor(bigframes.session.executor.Executor):
     def __init__(
         self, consumer: bigframes.session.substrait_executor.SubstraitConsumer
     ):
-        from bigframes.session.substrait_executor import SubstraitExecutor
+        from bigframes.core.compile.substrait.compiler import SubstraitCompiler
+        from bigframes.session.substrait_executor import (
+            AceroSubstraitConsumer,
+            SubstraitExecutor,
+        )
+
+        if isinstance(consumer, AceroSubstraitConsumer):
+            compiler = SubstraitCompiler(duration_type="int", use_precision_types=False)
+        else:
+            compiler = SubstraitCompiler(duration_type="int")
 
-        self.executor = SubstraitExecutor(consumer)
+        self.executor = SubstraitExecutor(consumer, compiler)
 
     def execute(
         self,
@@ -51,7 +60,10 @@ def execute(
 
         result = asyncio.run(
             self.executor.execute(
-                array_value.node, ordered=True, peek=execution_spec.peek
+                array_value.node,
+                bigframes.session.execution_spec.ExecutionSpec(
+                    ordered=True, peek=execution_spec.peek
+                ),
             )
         )
         if result is None:

packages/bigframes/mypy.ini

@@ -45,5 +45,11 @@ ignore_missing_imports = True
 [mypy-anywidget]
 ignore_missing_imports = True
 
+[mypy-substrait.*]
+ignore_missing_imports = True
+
+[mypy-datafusion.*]
+ignore_missing_imports = True
+
 [mypy-bigframes_vendored.*]
 ignore_errors = True

packages/bigframes/testing/constraints-3.10.txt

@@ -18,7 +18,7 @@ pandas==1.5.3
 pandas-gbq==0.26.1
 polars==1.21.0
 substrait==0.29.0
-datafusion==45.0.0
+datafusion==45.2.0
 pyarrow==23.0.1
 pydata-google-auth==1.8.2
 pyiceberg==0.7.1