Dask graph rewriting/ expression optimisation

Is your feature request related to a problem?

I've been working on this problem for a while after @ianhi encouraged me to have a crack at it at Scipy back in July.

The TLDR; is that my Claude assisted prototype (https://github.com/charles-turner-1/xrexpr) has gotten pretty far along (heads up for anyone looking at that - I've been thrashing the code out with Claude. I haven't let it decide any of the design but I have delegated pretty much all the code writing to it).

Anyhow, after running into flox last week and a bit of playing around over the weekend, I think I've reached a bit of a fork in the road where it might be worth publicly asking about this prototype and what's the best direction to go with it.

My prototype basically uses an accessor to capture attribute access on a dataset/dataarray and turn it into an IR which it then reasons about - it does things like predicate/projection pushdown, removes unnecessary rechunks, etc, only allowing provably safe rewrites. It then just replays these back out and lets xarray reconstruct the outcome of the chain, delegating everything to the existing machinery.

Under this way of doing things, we can only speed expression chains up as far as we could get from a 'perfectly written' chain, but no more.

As a simple canonical example:

ds.mean(dim="lat").mean(dim="lon").isel(time=0)

becomes

ds.plan.mean(dim="lat").mean(dim="lon").isel(time=0)

would get rewritten to

ds.isel(time=0).mean(dim=["lat","lon"])

If we have a time dim of len 100, this is going to give us an ~99% runtime reduction when we do compute, as we're removing the 99% of unnecessary work.

Anyhow, the fork in the road that I alluded to was (long & boring detailed exploratory notebook full of claude-isms here https://github.com/charles-turner-1/xrexpr/pull/184) that after prodding around a bit in how flox works, I realised that this expression rewriting approach would let us do a generic map-reduction, which can generate a speedup, even for 'perfectly written' expressions.

Something like (ds['temperature'] * 1.8 + 32).mean() is basically a map-reduction. My current understanding is we can't really generically turn this into a map reduction the way it's currently implemented, because - in a loose sense- the *1.8 + 32 and .mean() parts of the computation don't really know about each other. Flox can, because it's working on eg. DatasetGroupBy, so the two parts of the computation do know about each other.

Anyhow, with an Intermediate Representation that has these nodes, we can so something like:

• (ds['temperature'] * 1.8 + 32).mean() becomes
• [Projection(args=('temperature'),ElementWise(op='__mul__',val=1.8),ElementWise(op='__add__',val=32),Reduce(op='mean',dims='all')]
  which we can rewrite to
• [Projection(args=('temperature'),ElementWiseAffineMap(add=32,mul=1.8),Reduce(op='mean',dims='all')] and finally
• [Projection(args=('temperature'),MapReduce(op=Affine(1.8,32),dims='all')]

This lets us push the computation straight into the dask blocks, flox style (as I've understood it. I guess if I've misunderstood I've discovered a different optimisation?), which gets a ~2x speedup in my benchmark in the marimo notebook tests I linked - but has to graft a new dask graph back on in a fairly hacky way.

Anyhow, this kind of optimisation is getting out of the scope of the prototype I'd produced, as it's the sort of thing that requires doing more than just replaying operations in a different order and/or merging them. I figured I'd open an issue up to get some feedback on the idea - I'm sure there are plenty of smart people in here who have thought about these sort of ideas before.

(Obviously there's been a lot of aggressive AI usage to get this prototype to the state is it at the speed I have. The whole thing is pretty verbose so I'm happy to summarise any relevant parts for anyone interested - I have actually read and verified - as well as I'm able to 😅 - all of the code)

Describe the solution you'd like

No response

Describe alternatives you've considered

No response

Additional context

No response

Begin with the xrexpr prototype and exploratory notebook linked in the issue, then compare their intermediate-representation and rewrite approach with xarray's existing computation machinery. Before implementation, establish a concrete design and acceptance criteria for graph rewrites or map-reductions; this issue currently requests feedback and names no target files or tests.