ruff

Author: TrevorBergeron

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

@@ -313,7 +313,6 @@ def _compile_concat(self, node: nodes.ConcatNode) -> algebra_pb2.Rel:
     def _compile_aggregate(self, node: nodes.AggregateNode) -> algebra_pb2.Rel:
         input_rel = self._compile_node(node.child)
 
-        import bigframes.dtypes as dtypes
         import bigframes.operations.aggregations as agg_ops
 
         child_ids = list(node.child.ids)
@@ -390,7 +389,7 @@ def _compile_aggregate(self, node: nodes.AggregateNode) -> algebra_pb2.Rel:
                         idx = child_ids.index(col_id)
                         field_expr = algebra_pb2.Expression()
                         field_expr.selection.direct_reference.struct_field.field = idx
-                        
+
                         arg = measure.measure.arguments.add()
                         arg.value.CopyFrom(field_expr)
                     except ValueError:
@@ -406,7 +405,9 @@ def _compile_aggregate(self, node: nodes.AggregateNode) -> algebra_pb2.Rel:
                 not_null_op.scalar_function.function_reference = self._EXTENSIONS[
                     "is_not_null"
                 ]
-                json_format.ParseDict({"bool": {}}, not_null_op.scalar_function.output_type)
+                json_format.ParseDict(
+                    {"bool": {}}, not_null_op.scalar_function.output_type
+                )
                 not_null_op.scalar_function.arguments.add().value.CopyFrom(key_expr)
                 not_null_exprs.append(not_null_op)
 
@@ -417,7 +418,9 @@ def _compile_aggregate(self, node: nodes.AggregateNode) -> algebra_pb2.Rel:
                     and_expr.scalar_function.function_reference = self._EXTENSIONS[
                         "and"
                     ]
-                    json_format.ParseDict({"bool": {}}, and_expr.scalar_function.output_type)
+                    json_format.ParseDict(
+                        {"bool": {}}, and_expr.scalar_function.output_type
+                    )
                     and_expr.scalar_function.arguments.add().value.CopyFrom(expr)
                     and_expr.scalar_function.arguments.add().value.CopyFrom(e)
                     expr = and_expr
@@ -450,57 +453,89 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
             expr.selection.direct_reference.struct_field.field = idx
 
         # 2. Map window frame bounds (RowsWindowBounds / RangeWindowBounds / None)
-        bounds_type = algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_UNSPECIFIED
+        bounds_type = (
+            algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_UNSPECIFIED
+        )
         lower_bound = algebra_pb2.Expression.WindowFunction.Bound()
         upper_bound = algebra_pb2.Expression.WindowFunction.Bound()
 
         if node.window_spec.bounds is not None:
             if isinstance(node.window_spec.bounds, window_spec.RowsWindowBounds):
-                bounds_type = algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_ROWS
-                
+                bounds_type = (
+                    algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_ROWS
+                )
+
                 # Lower bound mapping
                 start = node.window_spec.bounds.start
                 if start is None:
-                    lower_bound.unbounded.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.Unbounded())
+                    lower_bound.unbounded.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
+                    )
                 elif start == 0:
-                    lower_bound.current_row.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.CurrentRow())
+                    lower_bound.current_row.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.CurrentRow()
+                    )
                 elif start < 0:
                     lower_bound.preceding.offset = -start
                 else:
                     lower_bound.following.offset = start
-                    
+
                 # Upper bound mapping
                 end = node.window_spec.bounds.end
                 if end is None:
-                    upper_bound.unbounded.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.Unbounded())
+                    upper_bound.unbounded.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
+                    )
                 elif end == 0:
-                    upper_bound.current_row.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.CurrentRow())
+                    upper_bound.current_row.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.CurrentRow()
+                    )
                 elif end < 0:
                     upper_bound.preceding.offset = -end
                 else:
                     upper_bound.following.offset = end
 
             elif isinstance(node.window_spec.bounds, window_spec.RangeWindowBounds):
-                bounds_type = algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_RANGE
+                bounds_type = (
+                    algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_RANGE
+                )
                 start = node.window_spec.bounds.start
                 if start is None:
-                    lower_bound.unbounded.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.Unbounded())
+                    lower_bound.unbounded.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
+                    )
                 elif start == pd.Timedelta(0):
-                    lower_bound.current_row.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.CurrentRow())
+                    lower_bound.current_row.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.CurrentRow()
+                    )
                 else:
-                    raise NotImplementedError("Range window bounds with non-zero offsets are not supported yet")
+                    raise NotImplementedError(
+                        "Range window bounds with non-zero offsets are not supported yet"
+                    )
 
                 end = node.window_spec.bounds.end
                 if end is None:
-                    upper_bound.unbounded.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.Unbounded())
+                    upper_bound.unbounded.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
+                    )
                 elif end == pd.Timedelta(0):
-                    upper_bound.current_row.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.CurrentRow())
+                    upper_bound.current_row.CopyFrom(
+                        algebra_pb2.Expression.WindowFunction.Bound.CurrentRow()
+                    )
                 else:
-                    raise NotImplementedError("Range window bounds with non-zero offsets are not supported yet")
+                    raise NotImplementedError(
+                        "Range window bounds with non-zero offsets are not supported yet"
+                    )
         else:
-            bounds_type = algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_ROWS
-            lower_bound.unbounded.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.Unbounded())
-            upper_bound.unbounded.CopyFrom(algebra_pb2.Expression.WindowFunction.Bound.Unbounded())
+            bounds_type = (
+                algebra_pb2.Expression.WindowFunction.BoundsType.BOUNDS_TYPE_ROWS
+            )
+            lower_bound.unbounded.CopyFrom(
+                algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
+            )
+            upper_bound.unbounded.CopyFrom(
+                algebra_pb2.Expression.WindowFunction.Bound.Unbounded()
+            )
 
         # 3. Project each window aggregation expression as a WindowFunction expression
         for agg_idx, col_def in enumerate(node.agg_exprs):
@@ -553,7 +588,9 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
             win_func.phase = algebra_pb2.AGGREGATION_PHASE_INITIAL_TO_RESULT
 
             bound_expr = ex.bind_schema_fields(agg, node.child.field_by_id)
-            type_dict = self._convert_type(dtypes.dtype_for_etype(bound_expr.output_type))
+            type_dict = self._convert_type(
+                dtypes.dtype_for_etype(bound_expr.output_type)
+            )
             json_format.ParseDict(type_dict, win_func.output_type)
 
             if distinct or isinstance(agg.op, agg_ops.NuniqueOp):
@@ -574,7 +611,9 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
             # Set sorting keys (sorts)
             for ord_expr in node.window_spec.ordering:
                 sort_field = win_func.sorts.add()
-                sort_pb = self._compile_expression(ord_expr.scalar_expression, node.child)
+                sort_pb = self._compile_expression(
+                    ord_expr.scalar_expression, node.child
+                )
                 sort_field.expr.CopyFrom(sort_pb)
 
                 is_asc = ord_expr.direction.is_ascending
@@ -596,7 +635,7 @@ def _compile_window(self, node: nodes.WindowOpNode) -> algebra_pb2.Rel:
                         idx = child_ids.index(col_id)
                         field_expr = algebra_pb2.Expression()
                         field_expr.selection.direct_reference.struct_field.field = idx
-                        
+
                         arg = win_func.arguments.add()
                         arg.value.CopyFrom(field_expr)
                     except ValueError:

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

@@ -15,17 +15,16 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import cast, Dict, Tuple
 
+import bigframes.dtypes as dtypes
+import bigframes.operations as ops
+import bigframes.operations.aggregations as agg_ops
 from bigframes.core import (
     agg_expressions,
     expression,
     identifiers,
     nodes,
 )
-import bigframes.dtypes as dtypes
-import bigframes.operations as ops
-import bigframes.operations.aggregations as agg_ops
 
 
 def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
@@ -58,22 +57,21 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
                 is_bool = col_dtype == dtypes.BOOL_DTYPE
                 if isinstance(
                     agg.op, (agg_ops.StdOp, agg_ops.VarOp, agg_ops.PopVarOp)
-                ) or (
-                    isinstance(agg.op, agg_ops.MeanOp)
-                    and is_bool
-                ):
+                ) or (isinstance(agg.op, agg_ops.MeanOp) and is_bool):
                     cast_aggs.append((agg_idx, col_id, dtypes.FLOAT_DTYPE))
                 elif isinstance(agg.op, agg_ops.SumOp) and is_bool:
                     cast_aggs.append((agg_idx, col_id, dtypes.INT_DTYPE))
 
     # If we need pre-projection (casts or size constants)
     if size_aggs or cast_aggs:
         assignments = []
-        
+
         cast_agg_to_col_id = {}
         for agg_idx, col_id, target_dtype in cast_aggs:
             new_id = identifiers.ColumnId(f"bf_cast_{col_id.name}_{agg_idx}")
-            cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(expression.deref(col_id.name))
+            cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(
+                expression.deref(col_id.name)
+            )
             assignments.append((cast_expr, new_id))
             cast_agg_to_col_id[(agg_idx, col_id)] = new_id
 
@@ -103,10 +101,12 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
                 new_exprs = []
                 for col_id in agg.column_references:
                     if (agg_idx, col_id) in cast_agg_to_col_id:
-                        new_exprs.append(expression.deref(cast_agg_to_col_id[(agg_idx, col_id)].name))
+                        new_exprs.append(
+                            expression.deref(cast_agg_to_col_id[(agg_idx, col_id)].name)
+                        )
                     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]
@@ -132,16 +132,20 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
     output_ids = group_ids + agg_ids
 
     expected_types = [item.dtype for item in node.schema.items]
-    
+
     assignments = []
     selection_pairs = []
     for idx, (out_id, target_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(expression.deref(out_id.name))
+        cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(
+            expression.deref(out_id.name)
+        )
         assignments.append((cast_expr, cast_id))
-        
-        selection_pairs.append((nodes.AliasedRef(expression.deref(cast_id.name), out_id), out_id))
-        
+
+        selection_pairs.append(
+            (nodes.AliasedRef(expression.deref(cast_id.name), out_id), out_id)
+        )
+
     post_project = nodes.ProjectionNode(
         node,
         assignments=tuple(assignments),
@@ -150,7 +154,7 @@ def rewrite_substrait_aggregations(node: nodes.BigFrameNode) -> nodes.BigFrameNo
         post_project,
         input_output_pairs=tuple(ref for ref, _ in selection_pairs),
     )
-    
+
     return post_selection
 
 
@@ -170,7 +174,7 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
     # Collect size and cast requirements for the window agg expressions
     size_aggs = []
     cast_aggs = []
-    
+
     for agg_idx, col_def in enumerate(node.agg_exprs):
         agg = col_def.expression
         if isinstance(agg.op, (agg_ops.SizeOp, agg_ops.SizeUnaryOp)):
@@ -182,22 +186,21 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
                 is_bool = col_dtype == dtypes.BOOL_DTYPE
                 if isinstance(
                     agg.op, (agg_ops.StdOp, agg_ops.VarOp, agg_ops.PopVarOp)
-                ) or (
-                    isinstance(agg.op, agg_ops.MeanOp)
-                    and is_bool
-                ):
+                ) or (isinstance(agg.op, agg_ops.MeanOp) and is_bool):
                     cast_aggs.append((agg_idx, col_id, dtypes.FLOAT_DTYPE))
                 elif isinstance(agg.op, agg_ops.SumOp) and is_bool:
                     cast_aggs.append((agg_idx, col_id, dtypes.INT_DTYPE))
 
     # If we need pre-projection (casts or size constants)
     if size_aggs or cast_aggs:
         assignments = []
-        
+
         cast_agg_to_col_id = {}
         for agg_idx, col_id, target_dtype in cast_aggs:
             new_id = identifiers.ColumnId(f"bf_window_cast_{col_id.name}_{agg_idx}")
-            cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(expression.deref(col_id.name))
+            cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(
+                expression.deref(col_id.name)
+            )
             assignments.append((cast_expr, new_id))
             cast_agg_to_col_id[(agg_idx, col_id)] = new_id
 
@@ -220,7 +223,7 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         for agg_idx, col_def in enumerate(node.agg_exprs):
             agg = col_def.expression
             out_col_id = col_def.id
-            
+
             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(
@@ -230,10 +233,12 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
                 new_exprs = []
                 for col_id in agg.column_references:
                     if (agg_idx, col_id) in cast_agg_to_col_id:
-                        new_exprs.append(expression.deref(cast_agg_to_col_id[(agg_idx, col_id)].name))
+                        new_exprs.append(
+                            expression.deref(cast_agg_to_col_id[(agg_idx, col_id)].name)
+                        )
                     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]
@@ -245,9 +250,7 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
             else:
                 rewritten_agg = agg
 
-            rewritten_agg_exprs.append(
-                nodes.ColumnDef(rewritten_agg, out_col_id)
-            )
+            rewritten_agg_exprs.append(nodes.ColumnDef(rewritten_agg, out_col_id))
 
         node = dataclasses.replace(
             node,
@@ -256,24 +259,32 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         )
 
     # Post-projection to enforce output schema types for newly introduced window columns:
-    child_output_ids = list(node.child.ids) if not isinstance(child, nodes.ProjectionNode) else list(child.child.ids)
-    
+    child_output_ids = (
+        list(node.child.ids)
+        if not isinstance(child, nodes.ProjectionNode)
+        else list(child.child.ids)
+    )
+
     assignments = []
     selection_pairs = []
-    
+
     for child_id in child_output_ids:
         selection_pairs.append((nodes.AliasedRef.identity(child_id), child_id))
-        
+
     for col_def, field in zip(node.agg_exprs, node.added_fields):
         out_id = col_def.id
         target_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(expression.deref(out_id.name))
+        cast_expr = ops.AsTypeOp(to_type=target_dtype).as_expr(
+            expression.deref(out_id.name)
+        )
         assignments.append((cast_expr, cast_id))
-        
-        selection_pairs.append((nodes.AliasedRef(expression.deref(cast_id.name), out_id), out_id))
-        
+
+        selection_pairs.append(
+            (nodes.AliasedRef(expression.deref(cast_id.name), out_id), out_id)
+        )
+
     post_project = nodes.ProjectionNode(
         node,
         assignments=tuple(assignments),
@@ -282,5 +293,5 @@ def rewrite_substrait_windows(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         post_project,
         input_output_pairs=tuple(ref for ref, _ in selection_pairs),
     )
-    
+
     return post_selection