refactor: WindowOpNode can create multiple cols

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -405,17 +405,14 @@ def project_window_expr(
         self,
         expression: agg_expressions.Aggregation,
         window: WindowSpec,
-        skip_reproject_unsafe: bool = False,
     ):
         output_name = self._gen_namespaced_uid()
         return (
             ArrayValue(
                 nodes.WindowOpNode(
                     child=self.node,
-                    expression=expression,
+                    agg_exprs=(nodes.ColumnDef(expression, ids.ColumnId(output_name)),),
                     window_spec=window,
-                    output_name=ids.ColumnId(output_name),
-                    skip_reproject_unsafe=skip_reproject_unsafe,
                 )
             ),
             output_name,

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 = child.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 = df.with_columns(agg_pl.alias(cdef.id.sql))
+                else:  # row-bounded window
+                    window_result = self._calc_row_analytic_func(
+                        df, 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):

bigframes/core/expression_factoring.py

@@ -187,7 +187,9 @@ def graph_extract_window_expr() -> Optional[
         while result := graph_extract_window_expr():
             id, window_expr = result
             curr_root = nodes.WindowOpNode(
-                curr_root, window_expr.analytic_expr, window_expr.window, output_name=id
+                curr_root,
+                (nodes.ColumnDef(window_expr.analytic_expr, id),),
+                window_expr.window,
             )
         if len(graph.nodes) >= pre_size:
             raise ValueError("graph didn't shrink")