refactor: WindowOpNode can create multiple cols (#2284)

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -268,8 +268,7 @@ def compute_values(self, assignments: Sequence[ex.Expression]):
 
  def compute_general_expression(self, assignments: Sequence[ex.Expression]):
  named_exprs = [
- expression_factoring.NamedExpression(expr, ids.ColumnId.unique())
- for expr in assignments
+ nodes.ColumnDef(expr, ids.ColumnId.unique()) for expr in assignments
  ]
  # TODO: Push this to rewrite later to go from block expression to planning form
  # TODO: Jointly fragmentize expressions to more efficiently reuse common sub-expressions
@@ -279,7 +278,7 @@ def compute_general_expression(self, assignments: Sequence[ex.Expression]):
  for expr in named_exprs
  )
  )
- target_ids = tuple(named_expr.name for named_expr in named_exprs)
+ target_ids = tuple(named_expr.id for named_expr in named_exprs)
  new_root = expression_factoring.push_into_tree(self.node, fragments, target_ids)
  return (ArrayValue(new_root), target_ids)
 
@@ -403,22 +402,24 @@ def aggregate(
 
  def project_window_expr(
  self,
- expression: agg_expressions.Aggregation,
+ expressions: Sequence[agg_expressions.Aggregation],
  window: WindowSpec,
- skip_reproject_unsafe: bool = False,
  ):
- output_name = self._gen_namespaced_uid()
+ id_strings = [self._gen_namespaced_uid() for _ in expressions]
+ agg_exprs = tuple(
+ nodes.ColumnDef(expression, ids.ColumnId(id_str))
+ for expression, id_str in zip(expressions, id_strings)
+ )
+
  return (
  ArrayValue(
  nodes.WindowOpNode(
  child=self.node,
- expression=expression,
+ agg_exprs=agg_exprs,
  window_spec=window,
- output_name=ids.ColumnId(output_name),
- skip_reproject_unsafe=skip_reproject_unsafe,
  )
  ),
- output_name,
+ id_strings,
  )
 
  def isin(

bigframes/core/block_transforms.py

@@ -232,13 +232,11 @@ def _interpolate_column(
  masked_offsets,
  agg_ops.LastNonNullOp(),
  backwards_window,
- skip_reproject_unsafe=True,
  )
  block, next_value_offset = block.apply_window_op(
  masked_offsets,
  agg_ops.FirstNonNullOp(),
  forwards_window,
- skip_reproject_unsafe=True,
  )
 
  if interpolate_method == "linear":

bigframes/core/blocks.py

@@ -1091,20 +1091,14 @@ def multi_apply_window_op(
  *,
  skip_null_groups: bool = False,
  ) -> typing.Tuple[Block, typing.Sequence[str]]:
- block = self
- result_ids = []
- for i, col_id in enumerate(columns):
- label = self.col_id_to_label[col_id]
- block, result_id = block.apply_window_op(
- col_id,
- op,
- window_spec=window_spec,
- skip_reproject_unsafe=(i + 1) < len(columns),
- result_label=label,
- skip_null_groups=skip_null_groups,
- )
- result_ids.append(result_id)
- return block, result_ids
+ return self.apply_analytic(
+ agg_exprs=(
+ agg_expressions.UnaryAggregation(op, ex.deref(col)) for col in columns
+ ),
+ window=window_spec,
+ result_labels=self._get_labels_for_columns(columns),
+ skip_null_groups=skip_null_groups,
+ )
 
  def multi_apply_unary_op(
  self,
@@ -1181,44 +1175,39 @@ def apply_window_op(
  *,
  result_label: Label = None,
  skip_null_groups: bool = False,
- skip_reproject_unsafe: bool = False,
  ) -> typing.Tuple[Block, str]:
  agg_expr = agg_expressions.UnaryAggregation(op, ex.deref(column))
- return self.apply_analytic(
- agg_expr,
+ block, ids = self.apply_analytic(
+ [agg_expr],
  window_spec,
- result_label,
- skip_reproject_unsafe=skip_reproject_unsafe,
+ [result_label],
  skip_null_groups=skip_null_groups,
  )
+ return block, ids[0]
 
  def apply_analytic(
  self,
- agg_expr: agg_expressions.Aggregation,
+ agg_exprs: Iterable[agg_expressions.Aggregation],
  window: windows.WindowSpec,
- result_label: Label,
+ result_labels: Iterable[Label],
  *,
- skip_reproject_unsafe: bool = False,
  skip_null_groups: bool = False,
- ) -> typing.Tuple[Block, str]:
+ ) -> typing.Tuple[Block, Sequence[str]]:
  block = self
  if skip_null_groups:
  for key in window.grouping_keys:
  block = block.filter(ops.notnull_op.as_expr(key))
- expr, result_id = block._expr.project_window_expr(
- agg_expr,
+ expr, result_ids = block._expr.project_window_expr(
+ tuple(agg_exprs),
  window,
- skip_reproject_unsafe=skip_reproject_unsafe,
  )
  block = Block(
  expr,
  index_columns=self.index_columns,
- column_labels=self.column_labels.insert(
- len(self.column_labels), result_label
- ),
+ column_labels=self.column_labels.append(pd.Index(result_labels)),
  index_labels=self._index_labels,
  )
- return (block, result_id)
+ return (block, result_ids)
 
  def copy_values(self, source_column_id: str, destination_column_id: str) -> Block:
  expr = self.expr.assign(source_column_id, destination_column_id)

bigframes/core/compile/ibis_compiler/ibis_compiler.py

@@ -265,11 +265,13 @@ def compile_aggregate(node: nodes.AggregateNode, child: compiled.UnorderedIR):
 
 @_compile_node.register
 def compile_window(node: nodes.WindowOpNode, child: compiled.UnorderedIR):
- result = child.project_window_op(
- node.expression,
- node.window_spec,
- node.output_name.sql,
- )
+ result = child
+ for cdef in node.agg_exprs:
+ result = result.project_window_op(
+ cdef.expression, # type: ignore
+ node.window_spec,
+ cdef.id.sql,
+ )
  return result
 
 

bigframes/core/compile/polars/compiler.py

@@ -853,20 +853,23 @@ def compile_window(self, node: nodes.WindowOpNode):
  "min_period not yet supported for polars engine"
  )
 
- if (window.bounds is None) or (window.is_unbounded):
- # polars will automatically broadcast the aggregate to the matching input rows
- agg_pl = self.agg_compiler.compile_agg_expr(node.expression)
- if window.grouping_keys:
- agg_pl = agg_pl.over(
- self.expr_compiler.compile_expression(key)
- for key in window.grouping_keys
+ result = df
+ for cdef in node.agg_exprs:
+ assert isinstance(cdef.expression, agg_expressions.Aggregation)
+ if (window.bounds is None) or (window.is_unbounded):
+ # polars will automatically broadcast the aggregate to the matching input rows
+ agg_pl = self.agg_compiler.compile_agg_expr(cdef.expression)
+ if window.grouping_keys:
+ agg_pl = agg_pl.over(
+ self.expr_compiler.compile_expression(key)
+ for key in window.grouping_keys
+ )
+ result = result.with_columns(agg_pl.alias(cdef.id.sql))
+ else: # row-bounded window
+ window_result = self._calc_row_analytic_func(
+ result, cdef.expression, node.window_spec, cdef.id.sql
  )
- result = df.with_columns(agg_pl.alias(node.output_name.sql))
- else: # row-bounded window
- window_result = self._calc_row_analytic_func(
- df, node.expression, node.window_spec, node.output_name.sql
- )
- result = pl.concat([df, window_result], how="horizontal")
+ result = pl.concat([result, window_result], how="horizontal")
  return result
 
  def _calc_row_analytic_func(

bigframes/core/compile/sqlglot/compiler.py

@@ -19,7 +19,15 @@
 
 import sqlglot.expressions as sge
 
-from bigframes.core import expression, guid, identifiers, nodes, pyarrow_utils, rewrite
+from bigframes.core import (
+ agg_expressions,
+ expression,
+ guid,
+ identifiers,
+ nodes,
+ pyarrow_utils,
+ rewrite,
+)
 from bigframes.core.compile import configs
 import bigframes.core.compile.sqlglot.aggregate_compiler as aggregate_compiler
 from bigframes.core.compile.sqlglot.aggregations import windows
@@ -310,67 +318,71 @@ def compile_aggregate(node: nodes.AggregateNode, child: ir.SQLGlotIR) -> ir.SQLG
 @_compile_node.register
 def compile_window(node: nodes.WindowOpNode, child: ir.SQLGlotIR) -> ir.SQLGlotIR:
  window_spec = node.window_spec
- if node.expression.op.order_independent and window_spec.is_unbounded:
- # notably percentile_cont does not support ordering clause
- window_spec = window_spec.without_order()
+ result = child
+ for cdef in node.agg_exprs:
+ assert isinstance(cdef.expression, agg_expressions.Aggregation)
+ if cdef.expression.op.order_independent and window_spec.is_unbounded:
+ # notably percentile_cont does not support ordering clause
+ window_spec = window_spec.without_order()
 
- window_op = aggregate_compiler.compile_analytic(node.expression, window_spec)
+  window_op = aggregate_compiler.compile_analytic(cdef.expression, window_spec)
 
- inputs: tuple[sge.Expression, ...] = tuple(
- scalar_compiler.scalar_op_compiler.compile_expression(
- expression.DerefOp(column)
+ inputs: tuple[sge.Expression, ...] = tuple(
+ scalar_compiler.scalar_op_compiler.compile_expression(
+ expression.DerefOp(column)
+ )
+ for column in cdef.expression.column_references
  )
- for column in node.expression.column_references
- )
 
- clauses: list[tuple[sge.Expression, sge.Expression]] = []
- if window_spec.min_periods and len(inputs) > 0:
- if not node.expression.op.nulls_count_for_min_values:
- # Most operations do not count NULL values towards min_periods
- not_null_columns = [
- sge.Not(this=sge.Is(this=column, expression=sge.Null()))
- for column in inputs
- ]
- # All inputs must be non-null for observation to count
- if not not_null_columns:
- is_observation_expr: sge.Expression = sge.convert(True)
+ clauses: list[tuple[sge.Expression, sge.Expression]] = []
+ if window_spec.min_periods and len(inputs) > 0:
+ if not cdef.expression.op.nulls_count_for_min_values:
+ # Most operations do not count NULL values towards min_periods
+ not_null_columns = [
+ sge.Not(this=sge.Is(this=column, expression=sge.Null()))
+ for column in inputs
+ ]
+ # All inputs must be non-null for observation to count
+ if not not_null_columns:
+ is_observation_expr: sge.Expression = sge.convert(True)
+ else:
+ is_observation_expr = not_null_columns[0]
+ for expr in not_null_columns[1:]:
+ is_observation_expr = sge.And(
+ this=is_observation_expr, expression=expr
+ )
+ is_observation = ir._cast(is_observation_expr, "INT64")
+ observation_count = windows.apply_window_if_present(
+ sge.func("SUM", is_observation), window_spec
+ )
  else:
- is_observation_expr = not_null_columns[0]
- for expr in not_null_columns[1:]:
- is_observation_expr = sge.And(
- this=is_observation_expr, expression=expr
- )
- is_observation = ir._cast(is_observation_expr, "INT64")
- observation_count = windows.apply_window_if_present(
- sge.func("SUM", is_observation), window_spec
- )
- else:
- # Operations like count treat even NULLs as valid observations
- # for the sake of min_periods notnull is just used to convert
- # null values to non-null (FALSE) values to be counted.
- is_observation = ir._cast(
- sge.Not(this=sge.Is(this=inputs[0], expression=sge.Null())),
- "INT64",
- )
- observation_count = windows.apply_window_if_present(
- sge.func("COUNT", is_observation), window_spec
- )
-
- clauses.append(
- (
- observation_count < sge.convert(window_spec.min_periods),
- sge.Null(),
+ # Operations like count treat even NULLs as valid observations
+ # for the sake of min_periods notnull is just used to convert
+ # null values to non-null (FALSE) values to be counted.
+ is_observation = ir._cast(
+ sge.Not(this=sge.Is(this=inputs[0], expression=sge.Null())),
+ "INT64",
+ )
+ observation_count = windows.apply_window_if_present(
+ sge.func("COUNT", is_observation), window_spec
+ )
+
+ clauses.append(
+ (
+ observation_count < sge.convert(window_spec.min_periods),
+ sge.Null(),
+ )
  )
+ if clauses:
+ when_expressions = [sge.When(this=cond, true=res) for cond, res in clauses]
+ window_op = sge.Case(ifs=when_expressions, default=window_op)
+
+ # TODO: check if we can directly window the expression.
+ result = child.window(
+ window_op=window_op,
+ output_column_id=cdef.id.sql,
  )
- if clauses:
- when_expressions = [sge.When(this=cond, true=res) for cond, res in clauses]
- window_op = sge.Case(ifs=when_expressions, default=window_op)
-
- # TODO: check if we can directly window the expression.
- return child.window(
- window_op=window_op,
- output_column_id=node.output_name.sql,
- )
+ return result
 
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):