Optimisation order restricts what we can do - e.g. move Lower to a later stage

[phofl]

I am currently sad about the order of our optimizations, e.g how simplify, lower, combine similar and others are executed after another.

I am looking into different approaches how we can change the expression tree when we have reordered things as we want (e.g. currently after combine_similar). Adding the IO Fusion brought a nice performance improvement.

The step currently happens after combine_similar, e.g. also after lower, which is problematic because changing things like the number of partitions can have impact on other decision we would like to make, e.g. the merge algorithm if we can reduce to 1 partition for example or use a broadcast join if we can reduce the partitions on one side. Lower locks these algorithms In before we can touch the tree though.

I want to move lower to a later stage in our optimization approach, e.g. after combine_similar and also after we modify the expression as we see fit. There is no real downside as far as I can see, the only difference is that we can't run simplify again after lower happens, but I don't think that we have a use-case for this at the moment and I also can't see one coming up anytime soon.

[rjzamora]

I want to move lower to a later stage in our optimization approach

This has been my opinion since very early in the project. So, you can probably consider me "on board".

the only difference is that we can't run simplify again after lower happens

Yeah, I think the user should be allowed to call .simplify() after lowering, but I agree that this would often be a "bad idea," and that we wouldn't want to do it in optimize.

My general understanding is that we should be able to do most (if not all) column-projection and predicate-pushdown optimizations before any lowering is performed or any divisions-related information needs to be queried. After the "abstract expression graph" is already optimized, it seems like lowering should not produce expressions that need further simplify logic, but I can imagine cases where a new expression is "similar" to one that already exists (and so a follow-up combine_similar pass may make sense).

[phofl]

My general understanding is that we should be able to do most (if not all) column-projection and predicate-pushdown optimizations before any lowering is performed or any divisions-related information needs to be queried. After the "abstract expression graph" is already optimized, it seems like lowering should not produce expressions that need further simplify logic, but I can imagine cases where a new expression is "similar" to one that already exists (and so a follow-up combine_similar pass may make sense).

agree, yes

Yeah, I think the user should be allowed to call .simplify() after lowering, but I agree that this would often be a "bad idea," and that we wouldn't want to do it in optimize.

Yeah this is a very bad idea without another combine_similar at the moment

[rjzamora]

Update: This issue can be addressed by #338 + #339, but the changes in #338 illustrate why the current distinction between Expr.simplify and Expr._lower are probably insufficient. That is, we find that some _lower operations produce Expr classes that need further "simplification" to enable dask-specific optimizations like partition culling.

#338 attacks this problem in a very reasonable way: It moves the current _lower logic into simplify so that we get the dask-specific Partitions expression before lowering is performed. While I think this approach makes sense, I would certainly prefer that we avoid lowering from general (abstract) expressions to dask-specific expressions before _lower is called. I can think of two ways of solving this that align with our future plans to delay any interactions with partitions/divisions:

Alternative 1 (Illustrated in rjzamora:simplify-lower-simplify): We introduce a new LoweredExpr class to better distinguish between "abstract" and "dask-specific" expressions. With this class in place, we can start restricting the second simplify pass within optimize to only rewrite LoweredExpr instances.

Alternative 2 (Illustrated in rjzamora:simplify-lower-simplify): We generalize the logic used to traverse the expression graph in Expr.simplify, and move that logic to a new Expr.rewrite(kind: str) method. This makes it relatively easy to define new/specialized optimization passes by defining Expr._{kind}_down and Expr._{kind}_up methods for any expression that should be "rewritten" when Expr.rewrite({kind}) is executed. For example, we can introduce a new "culling" optimization pass by adding a Partitions._cull_down method that just uses the existing Partitions._simplify_down logic.

@mrocklin - I'd like to get your high-level thoughts on the ideas @phofl and I are batting around. I know we had similar discussions earlier in the project, and I think we need to detangle the optimize pass a bit to reduce the cognitive load needed to debug and improve query optimization. It would be nice if we could move in a direction where we are delaying interactions with divisions/npartitions as long as possible.

[mrocklin]

Thanks for pinging me on issues like this. Changes to either the core class hierarchy, or optimization passes will certainly be necessary over the lifetime of this project, but they also make me nervous because they are opportunities to paint ourselves into a corner.

My gut reaction without thinking deeply or understanding current concerns is to prefer Alternative 2 above. Mostly this is because it is much more similar to other systems that I've seen be successful in the past.

When I was last active on this repository the optimize function looked like this:

while expression changed:
    simplify expression as much as possible
    lower a little

My sense from what I'm reading today is that the optimize function looks like this:

simplify expression as much as possible
lower as much as possible
combine similar

Is that correct?

If so, I definitely empathize with the idea that we should be able to simplify after lowering. Often compile chains lower through several different representations.

I think what I'm hearing as a suggestion above is something like the following:

while changed:
    simplify expression as much as possible
    lower a bit
    combine similar

Would this make sense? Maybe it's bad because combine and simplify tend to cancel each other out.

It sounds like alternative 2 above is something like the following:

for optimization in optimizations:
    while changed:
        optimize as much as possible with a specific set of optimizations

[mrocklin]

Also happy to chat live if that would help. I find that these conversations are aided with specific examples if those are easy to produce.

[phofl]

I think I prefer 2 as well, I am a bit concerned that we end up with 20 different optimizations, but not concerned enough to block on this. This fits nicely into what I want to do with general modifications of the expression graph

[mrocklin]

I'd welcome more conversation on this topic as well. I suspect that there exists an Alternative 3 out there somewhere.

Where would I go to learn more about combine_similar?

[mrocklin]

I spoke briefly with @phofl live. I guess Alternative 2 would result in an optimization function that looked something like the following:

    result = expr
    expr = expr.optimize("general", "projection", "filter")
    expr = combine_similar(expr)
    expr = io_fusion(expr)
    expr = expr.optimize("general", "lower")

    result = optimize_blockwise_fusion(result)

Where each name, like "general", refers to a family of optimizations.

[rjzamora]

My sense from what I'm reading today is that the optimize function looks like this:

simplify expression as much as possible
lower as much as possible
combine similar

Is that correct?

Today, we are still doing something that you are probably familiar with:

while changed:
    simplify expression as much as possible
    lower a bit

# Optional optimizations
combine similar # (rolls back sub-optimal simplify behavior)
IO fusion
blockwise fusion

The primary idea being discussed in this issue is whether it is time to break the "lowering" step out of the while loop (and how to accomplish this if so). The problem with the status quo is that the intermingling of simplify and lower makes optimizations very difficult to think through, and also makes it difficult to enforce different behavior for general "abstract" expressions and dask-specific "lowered" expressions.

What I'd personally like in the near/long term is a system where we can fully optimize an expression graph without "lowering" dask-agnostic expressions down to dask-specific expressions (i.e. expressions with divisions/partitions/_layer support). What I'd like in the short term is to be able to fully execute column-projection, predicate pushdown, and combine-similar optimizations before simple abstract expressions have been lowered down to uglier expression groups. This short-term goal simply makes it easier for us to think through the most common/critical optimizations we expect to need in practice (we are using tpch has a proxy for this).

I spoke briefly with @phofl live. I guess Alternative 2 would result in an optimization function that looked something like the following: ...

Yes, I think that roughly captures one variation of the idea I have in mind.

Also happy to chat live if that would help... I suspect that there exists an Alternative 3 out there somewhere.

Happy to make time for this, since I agree a real-time chat may be fruitful.

[mrocklin]

and also makes it difficult to enforce different behavior for general "abstract" expressions and dask-specific "lowered" expressions

My intuition says to not think about a distinction here. They are all expressions. Some of them are more specified than others. I think that making two classes here creates some inertia. For example, what if in the future we need three classes or four classes?

Maybe what I'm hearing from you is some desire to think about prioritization among optimizations? Or maybe we're thinking about how to give expressions attributes that we can reason about, like uses_divisions or something. That way any optimization can quickly check if a child/parent uses_divisions and make a quick judgement about how to operate. Something like that?

What I'd like in the short term is to be able to fully execute column-projection, predicate pushdown, and combine-similar optimizations before simple abstract expressions have been lowered down to uglier expression groups. This short-term goal simply makes it easier for us to think through the most common/critical optimizations we expect to need in practice (we are using tpch has a proxy for this).

Can you say more about why you want this? This is for RAPIDS work I'm guessing? Can you give more insight about what the precise sticking points are for you? (sorry if you've said this elsewhere and I haven't been paying attention).

[mrocklin]

Happy to make time for this, since I agree a real-time chat may be fruitful

I'm around any time today. Also happy to wait until tomorrow if we want to include Patrick.

[rjzamora]

My intuition says to not think about a distinction here. They are all expressions. Some of them are more specified than others. I think that making two classes here creates some inertia. For example, what if in the future we need three classes or four classes?

I can understand this perspective. With that said, I feel somewhat strongly that Expr currently has too much dask-specific "stuff" attached to it. So, while I agree that both are expressions, I also feel that we are actually creating a lot more inertial by mixing everything together than we would if we made the top-level Expr class more agnostic to dask.

What I'd like in the short term is to be able to fully execute column-projection, predicate pushdown, and combine-similar >>optimizations before simple abstract expressions have been lowered down to uglier expression groups. This short-term >>goal simply makes it easier for us to think through the most common/critical optimizations we expect to need in practice >>(we are using tpch has a proxy for this).

Can you say more about why you want this? This is for RAPIDS work I'm guessing?

I think this is the piece that both @phofl and I desire, but I could be wrong. My interest in making column-projection, predicate pushdown, and combine-similar less dask specific comes from a desire to use dask-expr for query optimization in cudf- and pandas-only workflows (something strongly desired by RAPIDS users today).

[mrocklin]

I also feel that we are actually creating a lot more inertial by mixing everything together than we would if we made the top-level Expr class more agnostic to dask.

I don't disagree that there could be a design that separates things. This is maybe more similar to an Ibis-like infrastructure. I'm open to something if you can come up with something that doesn't create a lot of duplication. For example would you have two Add classes? My sense is that for as long as things are semantically equivalent (Add and PandasAdd and DaskAdd or something) then they should just be the same class, and different systems can use protocols (like __dask_layer__ or something) to compose.

Maybe you're suggesting something different though. I'm not sure.

My sense is that having __dask_foo__ protocols doesn't harm anything, and so I'm fine attaching them everywhere. My sense is that relying on divisions is bad to do early on. If you wanted to add a property like uses_divisions then I'd be totally fine with things.

So, summarizing here, I think it's good to avoid assuming knowledge about dask stuff until you have to. I'm ok differentiating stuff. I'm even ok adding class hierarchies, but I'll likely get very very picky about this latter bit. My guess is that it's more complex, but I'm open to being demonstrated otherwise.

desire to use dask-expr for query optimization in cudf- and pandas-only workflows

Got it. And this is what I expected. Can you help motivate more about the kinds of challenges that you run into when doing this today?

For example, if I wanted dask-expr to also support Pandas I think that that would be pretty easy today. I would just add a _apply_pandas method or something everywhere and it would be good. Of course, this might get strange if Pandas had different operations or different optimizations it wanted to do. That could get pretty interesting then, we'd want to have the ability to switch between different optimization systems entirely! Or maybe there would be a different approach that could keep things in the same system. I don't know; I'd have to learn more about the differences here to make a sound judgement.

With that summary, I'll ask again, can you give more insight about what the precise sticking points are for you?

[rjzamora]

can you give more insight about what the precise sticking points are for you?... Can you help motivate more about the kinds of challenges that you run into when doing this today?

On the question of "what is stoping you from using dask-expr for lazy cudf execution," the sticking points are that (1) we need something like an __apply__ protocol, and (2) we need a clean way to apply simplify/combine_simlar optimizations without "lowering." This issue is mostly related to sticking point (2), because I get the sense that there are many ways to approach this problem (and I certainly want to know your thoughts).

The only sticking point I'm seeing with current dask-expr development under the existing design is the fact that Patrick's recent tpch-motivated changes to combine_similar have been difficult to "think" through. In our various discussions we have mentioned several times that it would be nice to keep operations like Merge and SortValues from being lowered until after the combine-simliar stage. There is no fundamental reason that we would need to do this right now besides allowing more people to understand the code.

I don't disagree that there could be a design that separates things. This is maybe more similar to an Ibis-like infrastructure. I'm open to something if you can come up with something that doesn't create a lot of duplication. For example would you have two Add classes? My sense is that for as long as things are semantically equivalent (Add and PandasAdd and DaskAdd or something) then they should just be the same class, and different systems can use protocols (like dask_layer or something) to compose.

Right, exactly. I agree that we don't want to duplicate simple classes all over the place, so I probably agree that a class hierarchy might be the wrong way to go. Perhaps a protocol-based design is better-aligned with what we have so far.

For example, if I wanted dask-expr to also support Pandas I think that that would be pretty easy today. I would just add a _apply_pandas method or something everywhere and it would be good. Of course, this might get strange if Pandas had different operations or different optimizations it wanted to do. That could get pretty interesting then, we'd want to have the ability to switch between different optimization systems entirely! Or maybe there would be a different approach that could keep things in the same system. I don't know; I'd have to learn more about the differences here to make a sound judgement.

Yeah, I get the feeling that Expr.optimize will continue evolving. I'm thinking this evolution will probably be easiest to manage if the simplify (and perhaps lower) logic is extensible. I suppose if we go down the protocol path, the different optimization strategies could be encoded in there? I'm very open to ideas.