FWIW, here's a drive-by comment on the general expression design.

I much prefer the pattern of separating the transformation passes that you might do on the expression from the classes that implement it. Right now you have regenerate and optimize, but optimize is a big hammer that it will likely make sense to split into multiple pieces. A lightweight visitor pattern with memoization of visits to shared nodes in the DAG might look something like:

from functools import singledispatch
class MemoizingVisitor:
    def __init__(self, func):
        self.func = func
        self.cache = {}

    def __call__(self, expr):
        try:
            return self.cache[expr]
        except KeyError:
            return self.cache.setdefault(expr, func(expr, self))

@singledispatch
def _optimize(expr, visitor):
    raise AssertionError(f"Unhandled type {type(expr)}")

@_optimize.register(TypeA)
def _(expr, visitor):
    transformed_children = map(visitor, expr.children)
    return expr.regenerate(transformed_children)

...

def optimize(exprs):
    visitor = MemoizingVisitor(_optimize)
    return [visitor(expr) for expr in exprs]

You could imagine writing the visitors as fixed points of the recursion and then tying the recursive knot outside, but that makes it slightly difficult to control the visitor order (pre- versus post-, or cutoff at some point).

I like this approach though because it gives transparent handling of DAGs (as opposed to trees) without expanding the perceived size, and you can share the cache across transformations if wanted.

This relies on your expressions being hashable, so the structural part isn't allowed to be mutable.

Thanks for the feedback @wence-! I was definitely planning to clean up the regenerate/optimize traversal design. I just experimented with some quick changes based on your suggestion, and I agree that it is better to decouple traversal from the CollectionOperation class. I still need to continue cleaning things up a lot more (graph generation is still a mess), but any/all design suggestions are welcome.

So, looking over this briefly my understanding is that the design here is ...

1. The same DataFrame class, useful for API
2. A new Operation attribute which holds onto an expression tree and history
3. This operation object has the relevant metadata like meta, divisions, and the graph
4. Operations are replayable?

Earlier you mentioned an idea where we removed the Mapping-style graph and instead treated the colleciton (or in this case operation) as that graph. It seems like this is not implemented here yet, correct? Do you have a plan here?

Stepping back a bit, I'm seeing a lot of low-level review on this PR, which is great. However, I'm also curious about first thinking a bit about the plan presented here and making larger changes first if we think that they're going to be useful.

So, looking over this briefly my understanding is that the design here is ...

1. The same DataFrame class, useful for API
2. A new Operation attribute which holds onto an expression tree and history
3. This operation object has the relevant metadata like meta, divisions, and the graph
4. Operations are replayable?

Yes. This particular PR is experimenting with a few things that are likely to go away or change, but the primary concept that I do think is quite promising is the idea of adding a new operation attribute (or whatever name we can agree on). The purpose of this attribute is to:

1. Manage all relevant metadata for the collection (meta, divisions, whatever) so that the “state” of a collection is completely defined by its current operation attribute
2. Expose the necessary methods for low-level graph materialization
3. Expose the necessary methods for “regeneration” (to simplify expression-tree optimization before a low-level graph is materialized)

Right now, the state of a collection is completely decoupled from graph materialization (and optimization), and the overall idea here is to improve this pain point. Without doing something like this, I don’t think we have much of a chance of ever supporting general graph optimization.

Earlier you mentioned an idea where we removed the Mapping-style graph and instead treated the collection (or in this case >operation) as that graph. It seems like this is not implemented here yet, correct?

I think the answer to this is “yes and no,” but I may be misunderstanding.

My long-term vision was that the “operation” can replace the need for the HLG/Layer API altogether. However, my ideas for graph materialization do still require an operation to know how to generate the necessary graph to produce a specified set of output keys (which is still “mapping” like - even if the operation object is not a Mapping).

Although my long-term vision was to replace HLGs, I don’t think there is a practical path for completely bailing on HLG/Layers all at once. In fact, one variation on the current POC is to allow a CollectionOperation to (optionally) manage its own HLG subgraph. That is, we could lean on the graph-materialization logic we already have (for now), but use the cleaner CollectionOperation API for "high-level" optimization passes.

Do you have a plan here?

This PR is absolutely exploratory, and so I am not proposing a clear API or plan just yet. However, I do think I will be pushing hard for a design where the collection state and graph materialization logic live in the same place (like operation).

Stepping back a bit, I'm seeing a lot of low-level review on this PR, which is great. However, I'm also curious about first thinking a bit about the plan presented here and making larger changes first if we think that they're going to be useful.

The reviews so far have been very helpful to me. However, I agree that general design suggestions are most valuable at this point.

I don’t think there is a practical path for completely bailing on HLG/Layers all at once.

Let me push on this a little. Why not? If we think that this is a good idea, what would stop us from doing this to all of at least one of the major collections (dataframe, array) before pressing the green button?

My long-term vision was that the “operation” can replace the need for the HLG/Layer API altogether. However, my ideas for graph materialization do still require an operation to know how to generate the necessary graph to produce a specified set of output keys (which is still “mapping” like - even if the operation object is not a Mapping).

Every operation could produce a mapping for the keys that define that operation. We might not need a HLG as well though? In my mind an objective of this work would be to not have to create yet-another-system. HLGs go away and get replaced by operations. I would be sad if we had dataframe store operations store hlgs store layers store tasks.

However, I do think I will be pushing hard for a design where the collection state and graph materialization logic live in the same place (like operation).

👍

The reviews so far have been very helpful to me. However, I agree that general design suggestions are most valuable at this point

Sounds great. I like the idea of going back and forth. I'd just want to make sure that at some point we take the learnings from this PR, take a step back, and propose a plan.

Update: The current plan is to implement something similar to what I have here (and much more) in the https://github.com/dask/dask/tree/collection-refactor through multiple distinct PRs. See general plan here: #9117

It's great to see a lot of thought being put into fundamental improvements to the data structures! I was also very excited to see the design ideas around high-level expressions and so am wondering how (if at all) that fits into the work here?

Even if they're off the table for now, it might be an idea to consider them in the design to make any future retro-fitting of expression graphs as easy / least disruptive as possible.

I was also very excited to see the design ideas around #7933 and so am wondering how (if at all) that fits into the work here?

Good question @dhirschfeld ! The final goal of the proposed solution here is to capture the same behavior/features as a high-level expression system. However, the primary difference is in the proposed development path for getting to the final design (and the specifics of where the expression tree lives).

The general background here is that the current state of HighLevelGraph/Layer does not provide what we need for general “high-level” optimization and/or efficient graph generation. There seem to be at least three different ways that we can attack these challenges:

1. HLE-Focused Solution - Completely redesign the existing collection APIs to become high-level expressions themselves
2. HLG/Layer-Focused Solution - Push information about the collection state into HLG/Layers and redefine the Layer contract to provide what we need for efficient graph generation/handling
3. Something in the middle (i.e. This Proposal) - Incrementally move the collection APIs toward solution (1) by adding an HLE-like property to the existing collection APIs

From my perspective, the original HLE plan stalled, because an incremental development path was just too difficult to articulate and justify. @jcrist described some of these challenges, and pointed out that we might as well improve the existing HLG/Layer infrastructure to fit our needs (solution 2 above). Therefore, when I set out to explore a tractable solution here, my original intention was to revise the existing HLG/Layer design.

The reason I eventually moved away from an HLG-based solution was that it was much easier for me to imagine an incremental development plan for something completely separate. I am still very open to other solutions. However, my gut tells me that we need to start by moving collection properties (like meta, divisions, etc) into the object that is responsible for graph materialization. My understanding is that this critical step is easiest for solution 3, but I could be wrong.

From my perspective these are all the same thing. We're just trying to
find the least painful way to get to the same place :)

Yep! The main differences are definitely in the "pathway" to the solution.