Unify treatment of Expr and IR nodes in cudf-polars DSL

[wence-]

Description

As part of in-progress multi-GPU work, we will likely want to:

1. Introduce additional nodes into the IR namespace;
2. Implement rewrite rules for IR trees to express needed communication patterns;
3. Write visitors that translate expressions into an appropriate description for whichever multi-GPU approach we end up taking.

It was already straightforward to write generic visitors for Expr nodes, since those uniformly have a .children property for their dependents. In contrast, the IR nodes were more ad-hoc. To solve this, pull out the generic implementation from Expr into an abstract Node class. Now Expr nodes just inherit from this, and IR nodes do so similarly.

Redoing the IR nodes is a little painful because we want to make them hashable, so we have to provide a bunch of custom get_hashable implementations (the schema dict, for example, is not hashable).

With these generic facilities in place, we can now implement traversal and visitor infrastructure. Specifically, we provide:

• a mechanism for pre-order traversal of an expression DAG, yielding each unique node exactly once. This is useful if one wants to know if an expression contains some particular node;
• a mechanism for writing recursive visitors and then wrapping a caching scheme around the outside. This is useful for rewrites.

Some example usages are shown in tests.

Checklist

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

[rjzamora]

Thanks for working on this @wence- ! I'll probably end up taking several passes over this. Leaving a few comments after my first pass.

[Matt711]

Just a small things I noticed. I haven't really been a part of the multi-GPU discussions. Can you explain at a high level how the cudf.polars IR graph will be translated to a dask task graph? I think that's the approach @rjzamora mentioned offline.

Edit: I'm asking for my own understanding and because maybe that might help us identify any other helper functions like resuse_if_unchanged that might help with the effort.

[wence-]

Just a small things I noticed. I haven't really been a part of the multi-GPU discussions. Can you explain at a high level how the cudf.polars IR graph will be translated to a dask task graph? I think that's the approach @rjzamora mentioned offline.

Not completely straightforwardly. But broadly, write a visitor that emits the correct task graph for computing a node in the plan, given that the children have already been computed (i.e. we have a task graph for them).

Edit: I'm asking for my own understanding and because maybe that might help us identify any other helper functions like resuse_if_unchanged that might help with the effort.

I think that we should wait for that until we write things and then see if we can generalise the implementation.

Right now I know I need reuse-if-unchanged because it's a useful utility for transforming an IR graph into another one of the same type with some new nodes.

[wence-]

I will rebase on top of #17014 on monday.

[rjzamora]

Just a small things I noticed. I haven't really been a part of the multi-GPU discussions. Can you explain at a high level how the cudf.polars IR graph will be translated to a dask task graph? I think that's the approach @rjzamora mentioned offline.

Not completely straightforwardly. But broadly, write a visitor that emits the correct task graph for computing a node in the plan, given that the children have already been computed (i.e. we have a task graph for them).

Yes. We will eventually want to implement a visitor that can traverse the final IR/Node graph to generate a task graph that can be scheduled with distributed/dask-cuda.

Further Notes: In practice, it will probably make sense to start by rewriting the initial IR graph into a "partition-aware" IR graph before we generate the task graph. In other words: We probably want to define and propagate leaf-node partitioning, and define boundaries between node sequences that do or do-not require inter-partition communication. Ideally, this partition-aware representation would not have anything Dask-specific in it.

[Matt711]

Not completely straightforwardly. But broadly, write a visitor that emits the correct task graph for computing a node in the plan, given that the children have already been computed (i.e. we have a task graph for them).

Yes. We will eventually want to implement a visitor that can traverse the final IR/Node graph to generate a task graph that can be scheduled with distributed/dask-cuda.

Thanks that's a nice overview.

Further Notes: In practice, it will probably make sense to start by rewriting the initial IR graph into a "partition-aware" IR graph before we generate the task graph. In other words: We probably want to define and propagate leaf-node partitioning,

So this "partition aware" graph will need to carry metadata about the partitions? And as the graph is executed this metadata will be carried upward.

and define boundaries between node sequences that do or do-not require inter-partition communication. Ideally, this partition-aware representation would not have anything Dask-specific in it.

I think this is because some operations can execute on a single partition and others can't. What does "inter-partition communication" look like? Is this where shuffling the data comes into play?

[rjzamora]

So this "partition aware" graph will need to carry metadata about the partitions? And as the graph is executed this metadata will be carried upward.

There are probably several ways to approach this, so we don't need to commit to anything specific right now. We will probably want to be able to ask a node how many partitions it will produce (i.e. how many output tasks it will define). One way to do this is to attach an npartitions-like property to all IR nodes. Most operations will simply inherit the partitioning of its child(ren), but leaf nodes and reductions will depend on the data/algorithm.

I think this is because some operations can execute on a single partition and others can't. What does "inter-partition communication" look like? Is this where shuffling the data comes into play?

Yes. If I do something like add two columns together to assign a new column, then we don't need any communication between partition-0 and partition-1. However, if we want to perform a groupby aggregation, we usually do need to reduce/shuffle data between distinct partitions. When we use dask for execution, we define the inter-partition communication as edges in the task graph.

[wence-]

Rebased onto and fixed conflicts after the expression move.

[Matt711]

Thanks @wence- for adding examples to overview.md. I had some non-blocking questions.

[rjzamora]

Thanks for working on this @wence- !

I have a few minor thoughts/questions, but I generally approve of the change.

[vyasr]

I'm not exactly sure where this is going next since none of the traversal logic seems to actually be used yet, but I'm sure I'll see it soon :)

[vyasr]

Looks great. A couple of threads are still open and it would be good to resolve them one way or another, but I don't need to review again. Thanks!

[wence-]

/merge