feat: Support mixed scalar-analytic expressions

Author: TrevorBergeron

bigframes/core/agg_expressions.py

@@ -19,7 +19,7 @@
 import functools
 import itertools
 import typing
-from typing import Callable, Mapping, TypeVar
+from typing import Callable, Mapping, Tuple, TypeVar
 
 from bigframes import dtypes
 from bigframes.core import expression, window_spec
@@ -63,6 +63,10 @@ def inputs(
     ) -> typing.Tuple[expression.Expression, ...]:
         ...
 
+    @property
+    def children(self) -> Tuple[expression.Expression, ...]:
+        return self.inputs
+
     @property
     def free_variables(self) -> typing.Tuple[str, ...]:
         return tuple(
@@ -73,6 +77,10 @@ def free_variables(self) -> typing.Tuple[str, ...]:
     def is_const(self) -> bool:
         return all(child.is_const for child in self.inputs)
 
+    @functools.cached_property
+    def is_scalar_expr(self) -> bool:
+        return False
+
     @abc.abstractmethod
     def replace_args(self: TExpression, *arg) -> TExpression:
         ...
@@ -176,8 +184,13 @@ def output_type(self) -> dtypes.ExpressionType:
     def inputs(
         self,
     ) -> typing.Tuple[expression.Expression, ...]:
+        # TODO: Maybe make the window spec itself an expression?
         return (self.analytic_expr, *self.window.expressions)
 
+    @property
+    def children(self) -> Tuple[expression.Expression, ...]:
+        return self.inputs
+
     @property
     def free_variables(self) -> typing.Tuple[str, ...]:
         return tuple(
@@ -188,6 +201,10 @@ def free_variables(self) -> typing.Tuple[str, ...]:
     def is_const(self) -> bool:
         return all(child.is_const for child in self.inputs)
 
+    @functools.cached_property
+    def is_scalar_expr(self) -> bool:
+        return False
+
     def transform_children(
         self: WindowExpression,
         t: Callable[[expression.Expression], expression.Expression],

bigframes/core/array_value.py

@@ -16,19 +16,24 @@
 from dataclasses import dataclass
 import datetime
 import functools
+import itertools
 import typing
 from typing import Iterable, List, Mapping, Optional, Sequence, Tuple
 
 import google.cloud.bigquery
 import pandas
 import pyarrow as pa
 
-from bigframes.core import agg_expressions, bq_data
+from bigframes.core import (
+    agg_expressions,
+    bq_data,
+    expression_factoring,
+    join_def,
+    local_data,
+)
 import bigframes.core.expression as ex
 import bigframes.core.guid
 import bigframes.core.identifiers as ids
-import bigframes.core.join_def as join_def
-import bigframes.core.local_data as local_data
 import bigframes.core.nodes as nodes
 from bigframes.core.ordering import OrderingExpression
 import bigframes.core.ordering as orderings
@@ -261,6 +266,23 @@ def compute_values(self, assignments: Sequence[ex.Expression]):
             col_ids,
         )
 
+    def compute_general_expression(self, assignments: Sequence[ex.Expression]):
+        named_exprs = [
+            expression_factoring.NamedExpression(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
+        fragments = tuple(
+            itertools.chain.from_iterable(
+                expression_factoring.fragmentize_expression(expr)
+                for expr in named_exprs
+            )
+        )
+        target_ids = tuple(named_expr.name for named_expr in named_exprs)
+        new_root = expression_factoring.push_into_tree(self.node, fragments, target_ids)
+        return (ArrayValue(new_root), target_ids)
+
     def project_to_id(self, expression: ex.Expression):
         array_val, ids = self.compute_values(
             [expression],

bigframes/core/block_transforms.py

@@ -399,15 +399,18 @@ def pct_change(block: blocks.Block, periods: int = 1) -> blocks.Block:
     window_spec = windows.unbound()
 
     original_columns = block.value_columns
-    block, shift_columns = block.multi_apply_window_op(
-        original_columns, agg_ops.ShiftOp(periods), window_spec=window_spec
-    )
     exprs = []
-    for original_col, shifted_col in zip(original_columns, shift_columns):
-        change_expr = ops.sub_op.as_expr(original_col, shifted_col)
-        pct_change_expr = ops.div_op.as_expr(change_expr, shifted_col)
+    for original_col in original_columns:
+        shift_expr = agg_expressions.WindowExpression(
+            agg_expressions.UnaryAggregation(
+                agg_ops.ShiftOp(periods), ex.deref(original_col)
+            ),
+            window_spec,
+        )
+        change_expr = ops.sub_op.as_expr(original_col, shift_expr)
+        pct_change_expr = ops.div_op.as_expr(change_expr, shift_expr)
         exprs.append(pct_change_expr)
-    return block.project_exprs(exprs, labels=column_labels, drop=True)
+    return block.project_block_exprs(exprs, labels=column_labels, drop=True)
 
 
 def rank(

bigframes/core/blocks.py

@@ -1154,6 +1154,26 @@ def project_exprs(
             index_labels=self._index_labels,
         )
 
+    # This is a new experimental version of the project_exprs that supports mixing analytic and scalar expressions
+    def project_block_exprs(
+        self,
+        exprs: Sequence[ex.Expression],
+        labels: Union[Sequence[Label], pd.Index],
+        drop=False,
+    ) -> Block:
+        new_array, _ = self.expr.compute_general_expression(exprs)
+        if drop:
+            new_array = new_array.drop_columns(self.value_columns)
+
+        return Block(
+            new_array,
+            index_columns=self.index_columns,
+            column_labels=labels
+            if drop
+            else self.column_labels.append(pd.Index(labels)),
+            index_labels=self._index_labels,
+        )
+
     def apply_window_op(
         self,
         column: str,

bigframes/core/expression.py

@@ -15,11 +15,12 @@
 from __future__ import annotations
 
 import abc
+import collections
 import dataclasses
 import functools
 import itertools
 import typing
-from typing import Callable, Generator, Mapping, TypeVar, Union
+from typing import Callable, Dict, Generator, Mapping, Tuple, TypeVar, Union
 
 import pandas as pd
 
@@ -43,6 +44,7 @@ def free_var(id: str) -> UnboundVariableExpression:
     return UnboundVariableExpression(id)
 
 
+T = TypeVar("T")
 TExpression = TypeVar("TExpression", bound="Expression")
 
 
@@ -136,6 +138,11 @@ def is_identity(self) -> bool:
         """True for identity operation that does not transform input."""
         return False
 
+    @functools.cached_property
+    def is_scalar_expr(self) -> bool:
+        """True if expression represents scalar value or expression over scalar values (no windows or aggregations)"""
+        return all(expr.is_scalar_expr for expr in self.children)
+
     @abc.abstractmethod
     def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
         ...
@@ -150,6 +157,57 @@ def walk(self) -> Generator[Expression, None, None]:
         for child in self.children:
             yield from child.children
 
+    def unique_nodes(
+        self: Expression,
+    ) -> Generator[Expression, None, None]:
+        """Walks the tree for unique nodes"""
+        seen = set()
+        stack: list[Expression] = [self]
+        while stack:
+            item = stack.pop()
+            if item not in seen:
+                yield item
+                seen.add(item)
+                stack.extend(item.children)
+
+    def iter_nodes_topo(
+        self: Expression,
+    ) -> Generator[Expression, None, None]:
+        """Returns nodes in reverse topological order, using Kahn's algorithm."""
+        child_to_parents: Dict[Expression, list[Expression]] = collections.defaultdict(
+            list
+        )
+        out_degree: Dict[Expression, int] = collections.defaultdict(int)
+
+        queue: collections.deque["Expression"] = collections.deque()
+        for node in list(self.unique_nodes()):
+            num_children = len(node.children)
+            out_degree[node] = num_children
+            if num_children == 0:
+                queue.append(node)
+            for child in node.children:
+                child_to_parents[child].append(node)
+
+        while queue:
+            item = queue.popleft()
+            yield item
+            parents = child_to_parents.get(item, [])
+            for parent in parents:
+                out_degree[parent] -= 1
+                if out_degree[parent] == 0:
+                    queue.append(parent)
+
+    def reduce_up(self, reduction: Callable[[Expression, Tuple[T, ...]], T]) -> T:
+        """Apply a bottom-up reduction to the tree."""
+        results: dict[Expression, T] = {}
+        for node in list(self.iter_nodes_topo()):
+            # child nodes have already been transformed
+            child_results = tuple(results[child] for child in node.children)
+            result = reduction(node, child_results)
+            results[node] = result
+
+        return results[self]
+
 
 @dataclasses.dataclass(frozen=True)
 class ScalarConstantExpression(Expression):