perf: Fragmentize multiple col expressions as one (#2333)

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -625,21 +625,7 @@ def skew(
     # counts, moment3 for each column
     aggregations = []
     for col in original_columns:
-        delta3_expr = _mean_delta_to_power(3, col)
-        count_agg = agg_expressions.UnaryAggregation(
-            agg_ops.count_op,
-            ex.deref(col),
-        )
-        moment3_agg = agg_expressions.UnaryAggregation(
-            agg_ops.mean_op,
-            delta3_expr,
-        )
-        variance_agg = agg_expressions.UnaryAggregation(
-            agg_ops.PopVarOp(),
-            ex.deref(col),
-        )
-        skew_expr = _skew_from_moments_and_count(count_agg, moment3_agg, variance_agg)
-        aggregations.append(skew_expr)
+        aggregations.append(skew_expr(ex.deref(col)))
 
     block = block.aggregate(
         aggregations, grouping_column_ids, column_labels=column_labels
@@ -663,16 +649,7 @@ def kurt(
     # counts, moment4 for each column
     kurt_exprs = []
     for col in original_columns:
-        delta_4_expr = _mean_delta_to_power(4, col)
-        count_agg = agg_expressions.UnaryAggregation(agg_ops.count_op, ex.deref(col))
-        moment4_agg = agg_expressions.UnaryAggregation(agg_ops.mean_op, delta_4_expr)
-        variance_agg = agg_expressions.UnaryAggregation(
-            agg_ops.PopVarOp(), ex.deref(col)
-        )
-
-        # Corresponds to order of aggregations in preceding loop
-        kurt_expr = _kurt_from_moments_and_count(count_agg, moment4_agg, variance_agg)
-        kurt_exprs.append(kurt_expr)
+        kurt_exprs.append(kurt_expr(ex.deref(col)))
 
     block = block.aggregate(
         kurt_exprs, grouping_column_ids, column_labels=column_labels
@@ -686,13 +663,38 @@ def kurt(
     return block
 
 
+def skew_expr(expr: ex.Expression) -> ex.Expression:
+    delta3_expr = _mean_delta_to_power(3, expr)
+    count_agg = agg_expressions.UnaryAggregation(
+        agg_ops.count_op,
+        expr,
+    )
+    moment3_agg = agg_expressions.UnaryAggregation(
+        agg_ops.mean_op,
+        delta3_expr,
+    )
+    variance_agg = agg_expressions.UnaryAggregation(
+        agg_ops.PopVarOp(),
+        expr,
+    )
+    return _skew_from_moments_and_count(count_agg, moment3_agg, variance_agg)
+
+
+def kurt_expr(expr: ex.Expression) -> ex.Expression:
+    delta_4_expr = _mean_delta_to_power(4, expr)
+    count_agg = agg_expressions.UnaryAggregation(agg_ops.count_op, expr)
+    moment4_agg = agg_expressions.UnaryAggregation(agg_ops.mean_op, delta_4_expr)
+    variance_agg = agg_expressions.UnaryAggregation(agg_ops.PopVarOp(), expr)
+    return _kurt_from_moments_and_count(count_agg, moment4_agg, variance_agg)
+
+
 def _mean_delta_to_power(
     n_power: int,
-    val_id: str,
+    col_expr: ex.Expression,
 ) -> ex.Expression:
     """Calculate (x-mean(x))^n. Useful for calculating moment statistics such as skew and kurtosis."""
-    mean_expr = agg_expressions.UnaryAggregation(agg_ops.mean_op, ex.deref(val_id))
-    delta = ops.sub_op.as_expr(val_id, mean_expr)
+    mean_expr = agg_expressions.UnaryAggregation(agg_ops.mean_op, col_expr)
+    delta = ops.sub_op.as_expr(col_expr, mean_expr)
     return ops.pow_op.as_expr(delta, ex.const(n_power))
 
 

bigframes/core/expression.py

@@ -15,12 +15,11 @@
 from __future__ import annotations
 
 import abc
-import collections
 import dataclasses
 import functools
 import itertools
 import typing
-from typing import Callable, Dict, Generator, Mapping, Tuple, TypeVar, Union
+from typing import Callable, Generator, Mapping, TypeVar, Union
 
 import pandas as pd
 
@@ -162,57 +161,6 @@ 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):

bigframes/core/expression_factoring.py

@@ -18,7 +18,10 @@
 import functools
 import itertools
 from typing import (
+    Callable,
     cast,
+    Dict,
+    Generator,
     Hashable,
     Iterable,
     Iterator,
@@ -40,18 +43,72 @@
 
 _MAX_INLINE_COMPLEXITY = 10
 
+T = TypeVar("T")
+
+
+def unique_nodes(
+    roots: Sequence[expression.Expression],
+) -> Generator[expression.Expression, None, None]:
+    """Walks the tree for unique nodes"""
+    seen = set()
+    stack: list[expression.Expression] = list(roots)
+    while stack:
+        item = stack.pop()
+        if item not in seen:
+            yield item
+            seen.add(item)
+            stack.extend(item.children)
+
+
+def iter_nodes_topo(
+    roots: Sequence[expression.Expression],
+) -> Generator[expression.Expression, None, None]:
+    """Returns nodes in reverse topological order, using Kahn's algorithm."""
+    child_to_parents: Dict[
+        expression.Expression, list[expression.Expression]
+    ] = collections.defaultdict(list)
+    out_degree: Dict[expression.Expression, int] = collections.defaultdict(int)
+
+    queue: collections.deque[expression.Expression] = collections.deque()
+    for node in unique_nodes(roots):
+        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(
+    roots: Sequence[expression.Expression],
+    reduction: Callable[[expression.Expression, Tuple[T, ...]], T],
+) -> Tuple[T, ...]:
+    """Apply a bottom-up reduction to the forest."""
+    results: dict[expression.Expression, T] = {}
+    for node in list(iter_nodes_topo(roots)):
+        # 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 tuple(results[root] for root in roots)
+
 
 def apply_col_exprs_to_plan(
     plan: nodes.BigFrameNode, col_exprs: Sequence[nodes.ColumnDef]
 ) -> nodes.BigFrameNode:
-    # TODO: Jointly fragmentize expressions to more efficiently reuse common sub-expressions
     target_ids = tuple(named_expr.id for named_expr in col_exprs)
 
-    fragments = tuple(
-        itertools.chain.from_iterable(
-            fragmentize_expression(expr) for expr in col_exprs
-        )
-    )
+    fragments = fragmentize_expression(col_exprs)
     return push_into_tree(plan, fragments, target_ids)
 
 
@@ -101,14 +158,26 @@ class FactoredExpression:
     sub_exprs: Tuple[nodes.ColumnDef, ...]
 
 
-def fragmentize_expression(root: nodes.ColumnDef) -> Sequence[nodes.ColumnDef]:
+def fragmentize_expression(
+    roots: Sequence[nodes.ColumnDef],
+) -> Sequence[nodes.ColumnDef]:
     """
     The goal of this functions is to factor out an expression into multiple sub-expressions.
     """
-
-    factored_expr = root.expression.reduce_up(gather_fragments)
-    root_expr = nodes.ColumnDef(factored_expr.root_expr, root.id)
-    return (root_expr, *factored_expr.sub_exprs)
+    # TODO: Fragmentize a bit less aggressively
+    factored_exprs = reduce_up([root.expression for root in roots], gather_fragments)
+    root_exprs = (
+        nodes.ColumnDef(factored.root_expr, root.id)
+        for factored, root in zip(factored_exprs, roots)
+    )
+    return (
+        *root_exprs,
+        *dedupe(
+            itertools.chain.from_iterable(
+                factored_expr.sub_exprs for factored_expr in factored_exprs
+            )
+        ),
+    )
 
 
 @dataclasses.dataclass(frozen=True, eq=False)