Enable multi-partition Select operations containing basic aggregations

[rjzamora]

Description

The overall goal is to enable us to decompose arbitrary Expr graphs containing one or more "non-pointwise" nodes. In order to achieve this, I propose that we add an experimental FusedExpr class (and related Expr-graph decomposition utilities). The general idea is that we can iteratively traverse an Expr-graph in reverse-topological order, and rewrite the graph until it is entirely composed of FusedExpr nodes. From there, it becomes relatively simple to build the task graph for each FusedExpr node independently.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[rjzamora]

Note: The rewrite approach used in this PR is illustrated in the figure below.

[wence-]

Partial comments

[wence-]

"Renaming" feels like the wrong thing, because the options for one agg might not apply to the options of another.

[rjzamora]

the options for one agg might not apply to the options of another

Yeah, that's totally true. We definitely need to figure out how the options "plumbing" should work here. I was hoping the options would normally translate in a trivial way, but I don't have a great sense for the range of possibilities.

[wence-]

I think we just need pattern rules that map a global Agg into a pair of (local_aggs, finalise_agg) or whatever it looks like.

[rjzamora]

Yeah, for this iteration we can just pass in None for the "new" options by default (since that's actually the "correct" thing to do for the supported aggregations that require "renaming").

[wence-]

It seems like it should be possible to write this function as a single bottom up visit rather than this approach which does a fixed-point iteration with three full tree-traversals per iteration.

Can you describe what rewrite patterns at a node are?

I think it's something like:

def _decompose(expr: Expr, rec: ExprTransformer):
    new_children = tuple(map(rec, expr.children))
    fused_children = tuple(c for c in new_children if isinstance(c, FusedExpr))
    if fused_children:
        return FusedExpr(expr.dtype, expr, fused_children)
    elif not e.is_pointwise:
        # check for supported case of `Agg`...
        return FusedExpr(expr.dtype, expr, ())
    else:
        # pointwise, no fused children
        return expr

def decompose_expr_graph(expr):
   mapper = CachingVisitor(_decompose)
   return mapper(expr)

Do I have it right?

[rjzamora]

Yeah, kind of. I played with this for a bit, and I think I have something working now.

[rjzamora]

Update: This PR now depends on #18405

[wence-]

Lots of questions/comments. I think this is pretty good but there are a few places where I'm confused about the logic

[wence-]

These should be independent of the size of the graph, do you mean "large expressions"?

[wence-]

Do you mean "child" here, or do you mean the partition count of the input frame?

[rjzamora]

I use the term child-IR to mean the input dataframe in most places, since this IR node is the child of the overarching Select node (that owns this expression we are processing). This is indeed, not an expression child.

Would you prefer "input IR"? I just want to distinguish this input IR node from the IR node that "owns" the expression. Otherwise, I don't care much about the naming.

[wence-]

If you add:

diff --git a/python/cudf_polars/cudf_polars/dsl/traversal.py b/python/cudf_polars/cudf_polars/dsl/traversal.py
index 9c45a68812..50a643e2cc 100644
--- a/python/cudf_polars/cudf_polars/dsl/traversal.py
+++ b/python/cudf_polars/cudf_polars/dsl/traversal.py
@@ -49,6 +49,39 @@ def traversal(nodes: Sequence[NodeT]) -> Generator[NodeT, None, None]:
                 lifo.append(child)
 
 
+def cutoff_traversal(
+    nodes: Sequence[NodeT], *, cutoff_types: tuple[type[NodeT], ...]
+) -> Generator[NodeT, None, None]:
+    """
+    Pre-order traversal of nodes in an expression.
+
+    Parameters
+    ----------
+    nodes
+        Roots of expressions to traverse.
+    cutoff_types
+        Types to terminate traversal at. If a type is in this tuple
+        then we do not yield any of its children.
+
+    Yields
+    ------
+    Unique nodes in the expressions, parent before child, children
+    in-order from left to right.
+    """
+    seen = set(nodes)
+    lifo = list(nodes)
+
+    while lifo:
+        node = lifo.pop()
+        yield node
+        if isinstance(node, cutoff_types):
+            continue
+        for child in reversed(node.children):
+            if child not in seen:
+                seen.add(child)
+                lifo.append(child)
+
+
 def reuse_if_unchanged(node: NodeT, fn: GenericTransformer[NodeT, NodeT]) -> NodeT:
     """
     Recipe for transforming nodes that returns the old object if unchanged.

and then use (at the top of this function):

traverse = partial(cutoff_traversal, cutoff_types=(FusedExpr,)) else traversal
...
for e in exprs:
   for node in list(traverse([expr]))[::-1]:
       ...

You won't needlessly traverse the full tree when skipping fused exprs in the children.

[rjzamora]

Yeah, I was tempted to do this before, but held off. Shall I just add an optional cutoff_types argument to traversal rather than duplicate most of the logic?

[wence-]

Yeah, I think that's reasonable.

[wence-]

I would prefer to validate this once before lowering to a graph. Rather than on construction. This is where it's "annoying" that FusedExprs have two types of children. The actual sub expressions and the FusedExpr "children".

[wence-]

This seems very wrong, it's extracting all columns that are leaves of on (an arbitrary expression that we are shuffling on) and using those as the schema of the Select call, which only selects a single column (shuffle_on).

What are you trying to do here?

I suppose what you're trying to do is determine which columns of the child are required to evaluate shuffle_on?

[rjzamora]

I can avoid all this for now if you'd rather keep things simple (and there could definitely be a mistake).

The general idea is that we need to shuffle the input IR node to evaluate this expression, but the output is only a single column. So, in order to avoid shuffling a bunch of columns that we don't actually need to evaluate the expression, we (1) look for the underlying columns needed by the expression, (2) drop all other columns from the input IR node, and (3) shuffle only the columns we need.

[wence-]

That makes sense, but those columns need to be in the list of expressions you're selecting

[rjzamora]

but those columns need to be in the list of expressions you're selecting

I don't think I understand this comment very well. Are you saying that we cannot drop columns because we might need it in the overarching Select, even if we don't need it for the current FusedExpr we are processing in this function?

When we drop columns and shuffle IR, the modified/shuffled IR is only used for this specific FusedExpr. If other FusedExpr nodes depended on other columns, they would still be referencing the original input IR.

[wence-]

As I understand it, you're shuffling on some expression A. This expression depends on some columns (B, C). The child input might have additional columns (D, E, ...). And you want to ensure that you only move around the data you need to evaluate A.

So I think this is:

necessary_columns = [e for e in traversal([on]) if isinstance(e, expr.Col)]

input = Select({c.name: c.dtype for c in necessary_columns},
               necessary_columns,
               False,
               child)

shuffled = Shuffle(..., shuffle_on, input)

No?

[wence-]

question: it seems like this Select might induce a shuffle itself (if, for example, shuffle_on contains an n_unique aggregation, or later a sort). Will that break anything? Or is it guaranteed not to be the case?

[rjzamora]

Will that break anything? Or is it guaranteed not to be the case?

This is guaranteed not to be the case, because we have decomposed everything into a graph of FusedExpr nodes. A FusedExpr may only contain a single non-pointwise expression.

[rjzamora]

Actually, it is possible that two distinct FusedExpr nodes will end up shuffling the same input IR multiple times (on the same or different columns). I haven't thought through ways to optimize this case, but I don't think anything should "break".

[wence-]

The schema of this select is {named_expr.name: named_expr.value.dtype}

[rjzamora]

Right, but this is function is creating a new IR node that is shuffled (and that the original Select can be called on in a pointwise fashion). This temporary schema may have multiple columns (e.g. if we are selecting the uniques of some binary op between two columns).

[wence-]

But that's not what this Select expression does: it produces a dataframe with a single output column that is the result of evaluating shuffle_on on the input.

[rjzamora]

Update: I decided that I don't like how much work we are doing during graph-construction time. I prefer handling the decomposition entirely during the lowering/re-write process. I'm exploring a slight refactoring of this PR in rjzamora:complex-aggregations-refactor.

It may take me a few more hours tomorrow morning to clean that up and get shuffling implemented for n_unique (It also seems fine to leave n_unique for a follow-up).

[rjzamora]

Update: It is likely that #18492 will be superseding this PR. That PR accomplishes what we need without introducing any new (or Expr-specific) task-graph logic.

[rjzamora]

Closing in favor of #18492