expr

[dcherian]

This is a running thread to discuss dask-expr optimizations. I intend to add small examples in separate messages

cc @mrocklin

[dcherian]

Here's a case where I would expect slices to be pushed down through Blockwise

import dask
dask.config.set({'array.query-planning': True})

import xarray as xr

# This `108` chunksize is important; it is large enough to trigger automatic rechunking by flox
# FIXME: be more aggressive about rechunking with expr
ds = xr.tutorial.open_dataset("air_temperature", chunks={"time": 108})

arr = ds.air

# this rechunks so that the groupby-reduction can be applied blockwise
# this output has chunksize of 1 along `time`
clim = arr.resample(time="ME").mean()

# let's select just three months, so 3 chunks
subset = clim.sel(time=slice("2013-01", "2013-03"))

# I would expect the slice to be pushed down to both FromArray layers, but it is not
subset.data.expr.simplify()

PS: I'm happy to take a stab at resolving this tomorrow.

[mrocklin]

Very happy to take a look (or have Claude take a look). Your example isn't working for me though:

• is arr supposed to be ds.air?
• I'm not getting dask arrays even with the chunks= keyword and I'm not sure why

Also, hopefully if you install rich you'll be able to get nicer reprs that will help explain what's going on.

[keewis]

Without having more context:

is arr supposed to be ds.air

indeed.

I'm not getting dask arrays even with the chunks= keyword and I'm not sure why

you need to replace load_dataset with open_dataset (otherwise the dask arrays will be computed immediately)

[dcherian]

Apologies, I was summarizing a slightly messy notebook. The example works now

[mrocklin]

Sorry I'm still struggling. How did you produce the air.zarr file. You all may be assuming that I have more familiarity with xarray than I actually do.

[mrocklin]

Never mind. Good now. Thanks.

[dcherian]

🤦🏾‍♂️ fixed again; the zarr vs netcdf doesn't matter here; and i should be able to reproduce to pure dask.array blockwise + slicing later tonight.

[mrocklin]

OK, so we have this:

>>> subset.expr

 Operation                Shape        Bytes 
 Getitem                  3×25×53     31 kiB 
 └ Transpose              24×25×53   248 kiB 
   └ Groupby Nanmean      25×53×24   248 kiB 
     └ Groupby Nanmean    25×53×2920  30 MiB 
       ├ Rechunk          25×53×2920  30 MiB 
       │ └ Transpose      25×53×2920  30 MiB 
       │   └ Open Dataset 2920×25×53  30 MiB 
       └ Array            2920        23 kiB 

>>> subset.expr.simplify()

 Operation              Shape        Bytes 
 Transpose              3×25×53     31 kiB 
 └ Getitem              25×53×3     31 kiB 
   └ Groupby Nanmean    25×53×24   248 kiB 
     └ Groupby Nanmean  25×53×2920  30 MiB 
       ├ Transpose      25×53×2920  30 MiB 
       │ └ Open Dataset 2920×25×53  30 MiB 
       └ Array          2920        23 kiB 

Now you're saying that you want the getitem to traverse down through the groupby nanmean. It looks like the Getitem/Slice is slicing in the last dimension, but the groupby nanmean is also affecting that dimension, which is why dask array is being nervous about moving it through.

In dask array parlance the groupby nanmean is a blockwise (sorry, this is going to get a little ugly for a moment, but I figure it's good education).

>>> print(subset.data.expr.simplify().tree_repr())

Transpose: axes=(2, 0, 1)
  SliceSlicesIntegers: index=(slice(None, None, None), slice(None, None, None), slice(None, 3, None)) allow_getitem_optimization=True
    Blockwise: func='_extract_result' out_ind=(0, 1, 2) name='groupby_nanmean-11e1f256e8e6e275b89d6f0da934d110' adjust_chunks={0: (25,), 1: (53,), 2: (2, 1, 2, 2, 1, 2, 2, 1, 2, 1, 2, 2, 1, 2, 1)} new_axes={} concatenate=False _meta_provided='<array>' kwargs={'key': 'nanmean'}(0, 1, 2)
      Blockwise: func='_reduce_blockwise' out_ind=(0, 1, 2) name='groupby_nanmean-chunk-11e1f256e8e6e275b89d6f0da934d110' dtype='float64' new_axes={} align_arrays=False concatenate=False _meta_provided='<array>' kwargs={}(0, 1, 2)(2,)
        Transpose: axes=(1, 2, 0)
          FromArray: array='<ImplicitToExplicitIndexingAdapter>' _chunks=((236, 124, 244, 244, 124, 244, 244, 124, 236, 120, 244, 248, 120, 244, 124), (25,), (53,)) meta='<array>' _name_override='open_dataset-air-a5b61afff50733bab97ed2df986e37c3'
        FromArray: array='<array>' _chunks=((236, 124, 244, 244, 124, 244, 244, 124, 236, 120, 244, 248, 120, 244, 124),) _name_override='array'

Now, we're happy in general to push a slice through a blockwise, but we're nervous if that blockwise adjusts the chunking of the underlying array, which this certainly does. [n:m] has to be translated somehow into [f(n): f(m)] and we don't know how to do that here in general. Probably Xarray does though, through the use of the time index. This might be a case where we need some sort of Xarray expression rather than a Dask array expression. I think that Dask array (or any underlying ndarray abstraction) doesn't have the necessary metadata to do this transformation safely. I suspect that this optimization has to be done within the Xarray abstraction.

Open to thoughts. My main question right now is how much of the value of expression optimization requires Xarray knowledge, and how much of it can be done purely within the ndarray abstaction.

[mrocklin]

Or, if we didn't want to Expr-ify Xarray, we could make some other traversal system that went through and applied optimizations. This would probably be less invasive and not that hard, but it would have to live here I think.

[dcherian]

My thinking is that the blockwise takes in a chunk of size m, outputs a chunk of size n, so a slice that aligns with chunk boundaries e.g. slice(n, 3n) is fine to push down as slice(m, 3m) (or appropriately modified). Basically, block-slicing of blockwise is fine in theory. But now I can see that this is a little hard with general blockwise since there is an arbitrary function to transform chunks (the m->n mapping).

One thing I've been toying around with is whether flox needs to insert its own expr layer, instead of relying on general blockwise. This seems like a reason to do so.

[mrocklin]

Having your own Expr layer thats more specific than Blockwise would make this much easier to resolve.

My thinking is that the blockwise takes in a chunk of size m, outputs a chunk of size n, so a slice that aligns with chunk boundaries e.g. slice(n, 3n) is fine to push down as slice(m, 3m)

Yeah, historically we used to allow adjust_chunks to be a function. I could imagine something like what you're describing being more possible (but still probably not doable in general) in that world. In our case we just store a dict of new chunks now (it's hard to send a callable to the scheduler).

FWIW I suspect that Claude, given the context in the dask repo, could probably implement that new expression for you if given enough description of what you're trying to accomplish.

[dcherian]

In our case we just store a dict of new chunks now

In that case if the applied slice indices are present in cumsum(new_chunks), you can back out which "source" chunks are needed on all source arrays.

Happy to do this in flox, but I bet someone is going to expect this. It's also a pretty high-impact optimization I would think (blockwise being so pervasive)

[mrocklin]

In that case if the applied slice indices are present in cumsum(new_chunks), you can back out which "source" chunks are needed on all source arrays.

Maybe I need to better understand what flox is doing. This is a cumreduction? If so, why did the shape change? I'm quite happy to add more optimizations into dask array expressions (that's what it's there for) I just don't understand this situation well enough to understand what the pattern is. My understanding currently is that Flox is doing some blockwise operation with highly custom chunk changing. In that situation there's not much we can do. If Flox is actually doing some other more specific dask array operation (like calling cumreduction) and if there's some optimization we can do there I'm more than happy to work on it. I just need to learn more.

[dcherian]

I can work on a pure dask.array example tonight. Here is a visual from the docs [each square is a single array element; each color is a different group; gaps between groups of elements indicate a chunk boundary]:

If I apply a slicer aligned with blocks on the output (say slice(0, 4) => slice out the first two blocks); then to compute that i only need slice(0, 16) on the input (first two blocks of the input after rechunking). The blockwise applies a groupby-reduction per block which is why the chunk sizes change. Because of the intelligently chosen chunking, it happens to be a reduction without cross-block dependencies, so block-aligned slicing can be pushed down.

Does the expr branch still cull or are we relying on pushdown to remove the need to cull?

[mrocklin]

Does the expr branch still cull or are we relying on pushdown to remove the need to cull?

We're relying on slice pushdown to remove the need to cull. Beyond this, maybe the scheduler handles it. It's not something that task graphs are thinking about though.

If I apply a slicer aligned with blocks on the output (say slice(0, 4) => slice out the first two blocks); then to compute that i only need slice(0, 16) on the input (first two blocks of the input after rechunking). The blockwise applies a groupby-reduction per block which is why the chunk sizes change. Because of the intelligently chosen chunking, it happens to be a reduction without cross-block dependencies, so block-aligned slicing can be pushed down.

We only know that we only need the first two blocks in this case because we get to look at the colors though, right? This [:4] -> [:16] result doesn't hold generally for all groupby-reduction operations, right? I suspect that in this case the colors are coming from the index, is that right? So we need to have the index on-hand in order to make this optimization safely. If so, my guess is that the Flox expression will have access to that Index as it performs the optimization. That's the piece of information that the Dask array expressions are missing.