[Bug] Optimization is now much slower in TPCh benchmarks

[rjzamora]

I tried running the TPCh benchmarks in https://github.com/coiled/benchmarks and noticed that dask-expr was much slower than before (both with "pandas" and "cudf" as the backend). After looking at a performance report, it was clear that the slow down was not in compute. I confirmed that the regression was coming from optimize() (a change from ~1s to >10s for query 1).

I did a cursory bisection, and found that the regression was definitely introduced in #395

I haven't had time to figure out why 395 is slowing things down yet.

[fjetter]

Haven't been looking for the cause but the slowness is there because the expression graph for Q1 is huge. We encountered a similar problem recently in (#294; IIRC #769 relates to a partial fix since it was removing a couple of redundant Assign expressions)

Printing the entire expr takes 5k lines so it's hard to follow what is happening in detail.

The slowness is clearly that the expression tree should not be that deep for this relatively simple query. However, it's also interesting where all the time goes.
On my machine, the Q1 optimize for scale 100 takes currently about 11s. All of this time is due to deep recursion.

In particular (rough numbers here)

• 5s for meta of the Groupby expression
• 2.5s simplify_down which is mostly recursion. What's not recursion is again meta
• 3.5s rewrite (recursion -> meta)
• 800ms optimize_fusion (which points very deeply to recursive divisions)

(Note how the python stack itself is particularly deep because every attribute access has to go via the getattr magic in the Expr baseclass. These are the brown __get__ frames; This is not causing the slowness but it definitely isn't helping)

pyspy-speedscope-tpch-q1.svg.zip

[fjetter]

This is issue #786

Apparently, self-assign expressions, eg. df['b'] = df['a'] + 1 are effectively a diamond shaped graph which in dask-expr is broken into a tree, i.e. such an expression branches off into two branches. Therefore, another assignment of this type is branching off again by a factor of two. The number of branches therefore grows exponentially and in query 1 we have a lot of those assignment operations (some are mere renames, others are defining a new column)

[fjetter]

The size of the graph is unchanged and the exponential growth has been around since before #395 I haven't confirmed this yet but I assume it is indeed the runtime complexity I warned about in #395 (review)
There was the assumption that the prior main would exhibit the same runtime behavior (#395 (comment)) although this was apparently a false assumption. Also, back then we assumed that expr graph sizes would be "reasonably small" which is clearly a false statement if the size grows exponentially this easily.

[rjzamora]

Thanks for digging in to this @fjetter !

So, it looks like dask-expr indeed has a problem with the expression-graph size (#786), but that problem is not new. Rather, the regression is probably due to the fact that the new optimization approach introduces a higher level of computational complexity wrt that graph size.

I will try to take some time today to learn about the recent optimization changes. However, please feel free to suggest fixes/mitigations for this regression if you have any ideas :)

[mrocklin]

Haven't been looking for the cause but the slowness is there because the expression graph for Q1 is huge

I wouldn't expect this. I'd expect it to be a relatively deep DAG of ~50 nodes. It sounds like that's not the case, and that things are splitting/multiplying somehow? Or maybe it just looks really large because we're thinking about things as a tree rather than a dag sometimes?

Printing the entire expr takes 5k lines so it's hard to follow what is happening in detail

Probably this is a case where the visualize method might be better (it's able to capture)

My guess as to what is slow is that the dependencies / dependents calculation happens everywhere at every change (I think that this was expected to be slow, and patrick decided to wait until it was an issue). Just so we're all on the same page, my expectations would be that ...

1. Q1 of TPC-H is considered actually pretty small
2. Query Optimization for even fairly moderate graphs should probably happen sub-second

I don't think people disagree with these expectations, but I thought I'd mention them in case folks aren't on the same page so that we can then talk about it (maybe I'm wrong).

[fjetter]

Or maybe it just looks really large because we're thinking about things as a tree rather than a dag sometimes?

Yes, the DAG has just 32 nodes but since everything is forced into a tree this causes certain graph topologies to split out, e.g. diamonds. In this specific case it is about assignments #786

Is there a reason why this is forced into a tree?

[fjetter]

My guess as to what is slow is that the dependencies / dependents calculation happens everywhere at every change

There are actually a couple of different things why this is slow and it's not the dependency thing from what I can tell.

1. The simplify step is slow because we're recalculating _meta over and over again on different instances of the same expression. Even though we're caching many properties like meta on the instance, a quite fundamental assumption of dask-expr is that we can regenerate an expression over and over again (relying on the deterministic token). However, when doing so, the cache does no longer exist (using functools.cached_property). During optimize for this query, we're instantiating about 480 unique Expr object but about 20k in total. I see some getitem and assign expression being rebuilt almost 3k times with the same name.
2. Depending on what commit we're looking at this and this cause significant slowdown. the up/down_name in expr.__dir__() check is costly since __dir__ returns a list of all available attributes (I don't know what changed in the past about this. Going back to 0.3.5 this shows up as a large bottleneck, see profile above but it doesn't on current main. __dir__ is not sorted and if not implemented represents a sequence of the internal object dictionary. Therefore, more recent commits may just be sorted in a better way for us, idk, we should get rid of this)

[fjetter]

I opened #798 with two fixes (POC do not merge) that address the above things reducing the runtime to slightly more than a second. This is a little concerning to me and I'll dig a little deeper

[fjetter]

I added now a "is simplified" cache to #798 and we're at 0.2s for the optimize which is a little odd.Even with all duplication removed (by using #798) almost all of the remaining time is spent in the recursive definition of _meta.

For instance, a rather heavy one that is taking up almost half the time is the Blockwise meta operation which applies the blockwise operation to a non-empty meta object (after constructing the meta of all input arguments first, of course). In this specific example, this is a groupby operation so we see a lot of pandas take operations

Another example is a projection on a series of assigns where Assign isn't implemented very well and is calling multiple make_meta_non_empty https://github.com/dask-contrib/dask-expr/blob/13af21d1be9e0a3393f4971a0d95382188f6f248/dask_expr/_expr.py#L1770-L1774

Some of this stuff can be optimized but we're still dealing with ~500 unique expressions during this optimization step (we start with 32 nodes in the expression graph but the optimization steps are so fine grained that we have many intermediates). It's incredibly difficult to avoid calling _meta on any of those expressions so every _meta computation that is not essentially instant will hurt us badly. Most of the _meta calls do happen almost instantaneously but there are a couple that run for a couple ms (top individual offender in my profile ran for 40ms; mostly recursing into the __getattr__ implementation of _core.Expr)

Apart from a couple of micro optimizations we can apply here, the fundamental issue appears to be that the optimizer is not converging fast enough and has to take too many rounds. To me it looks as if the optimizer does not converge quickly enough.

I do not believe that a cache like it is being proposed in #797 is not viable. At least it will alter how we are optimizing but I don't have a reproducer for this. I will comment over there

[rjzamora]

It's incredibly difficult to avoid calling _meta on any of those expressions so every _meta computation that is not essentially instant will hurt us badly.

Agree that generating _meta is where all the observable time is, and that any case where we don't need non-empty metadata (e.g. Assign), we should avoid it.

Apart from a couple of micro optimizations we can apply here, the fundamental issue appears to be that the optimizer is not converging fast enough and has to take too many rounds.

That makes sense to me. Right now we either need optimizations to require very few "rounds", or we need _meta calculations to be fast or avoided altogether. When we need many rounds and the _meta isn't instant, we are in trouble.

I do not believe that a cache like it is being proposed in #797 is viable. At least it will alter how we are optimizing but I don't have a reproducer for this. I will comment over there

That doesn't surprise me, but the fact that we need to generate so many unique expression-dependent combinations is a bit scary.

[fjetter]

I do not believe that a cache like it is being proposed in #797 is viable. At least it will alter how we are optimizing but I don't have a reproducer for this. I will comment over there

That doesn't surprise me, but the fact that we need to generate so many unique expression-dependent combinations is a bit scary.

Sorry, I was meaning to say not viable. I edited my earlier comment