xarray store disables blockwise fusion if compute=False #11879
Description
#11844 introduced a regression for xarray users for all store APIs if compute=False.
da.store previously called optimize explicitly on the input arrays causing them to be optimized every single time regardless of what would happen next.
However, if compute=False xarray is wrapping the result of da.store again in a delayed object (see here) which effectively disables optimization.
This change has not been released yet. This will be considered a release blocker.
Reproducer
@pytest.mark.parametrize("compute", [True, False])
def test_xarray_blockwise_fusion_store(compute):
def custom_scheduler_get(dsk, keys, expected, **kwargs):
dsk = dsk.__dask_graph__()
assert (
len(dsk) == expected
), f"False number of tasks got {len(dsk)} but expected {expected}"
return [42 for _ in keys]
# First test that this mocking stuff works as expecged
with pytest.raises(AssertionError, match="False number of tasks"):
scheduler = partial(custom_scheduler_get, expected=42)
dask.compute(da.ones(10), scheduler=scheduler)
coord = da.arange(8, chunks=-1)
data = da.random.random((8, 8), chunks=-1) + 1
x = xr.DataArray(data, coords={"x": coord, "y": coord}, dims=["x", "y"])
y = ((x + 1) * 2) / 2 - 1
# Everything fused into one compute task
# one finalize Alias
expected = 2
scheduler = partial(custom_scheduler_get, expected=expected)
dask.compute(y, scheduler=scheduler)
with tmpdir() as dirname:
if compute:
with dask.config.set(scheduler=scheduler):
y.to_zarr(dirname, compute=True)
else:
# There's a delayed finalize store smashed on top which won't be fused by
# default
expected += 1
scheduler = partial(custom_scheduler_get, expected=expected)
stored = y.to_zarr(dirname, compute=False)
dask.compute(stored, scheduler=scheduler)
this is somewhat related to #11854
I will attempt to address this by changing the delayed interface such that optimizations are indeed called on all collections that are passed to the delayed object. This would be consistent to how we're handling Expr objects already and I don't see an obvious reason why this should be any different for traditional HLG objects.
If that fix fails, I will contribute something to xarray upstream to ensure they will be able to cope with this appropriately.
cc @dcherian for visibility