[jsignell]

• Towards opening a zarr dataset taking so much time with dask #8902
  • it doesn't really close it, but it is a better alternative than chunks={} or chunks="auto". Since Remove special mapping of auto to {} in open_zarr #11010 got in xarra could eventually change the default on open_zarr to map it to "preserve" rather than to {} if dask is available.
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

Proposal

A new chunks option that is only allowed to use encoded chunks or multiples of them. No chunk splitting allowed.

Demo

Current behavior when chunks="auto"

ds = xr.open_zarr("gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", chunks="auto")

This PR introduces a new option: chunks="preserve"

ds = xr.open_zarr("gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", chunks="preserve")

Context

I originally set out to update the auto_chunks function in dask, but it felt like my goals were actually quite different. The goal of the dask auto_chunks function is to guarantee that the chunksize will be under a configurable limit while preserving the aspect ratio of previous_chunks (previous_chunks == encoding). This PR instead guarantees that encoded chunks are never split but it will multiply them by some factor to try to get the chunksize close to a targetsize. It doesn't try to preserve the aspect ratio of the chunks. Instead it goes after the dim where there is the greatest number of chunks and it tries to take those in bigger bites.

Also:

• I'm not quite sure how the interface should work and I am definitely not attached to the word "preserve".
• I originally put this in the DaskManager, but moved it to the base ChunkManagerEntrypoint class once I realized there was nothing dasky about it. I'm not sure if there is really supposed to be logic in methods on that class though.

[jsignell]

These tests are kind of slow.

[jsignell]

@dcherian is this something you would be able to review? I'd love to get more people trying it.

[jsignell]

Ok @dcherian this is now a PR that changes the behavior of chunks="auto". The changes to tests were minimal with this change, but if we think it's important to give people a way to get the prior behavior we can add an option so that people can set xr.set_options(use_dask_auto=True)

This would be what the commit looks like to add the option jsignell@84cca20

[jsignell]

I can make this not change the existing behavior at all if we want to do a slow rollout where people would do:

import xarray as xr

with xr.set_options(use_dask_auto=False):
    ds = xr.open_zarr("gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", chunks="auto")

[jsignell]

I was just reminded of this PR and remembered that I still care about it! I am going to update this branch and include the option setting so that people can get back to the previous behavior.

[dcherian]

Sorry, I can review when you're done. I don't think an option is warranted, by saying "auto:, the user has already opted in to magic.

[github-actions]

😴 No significant changes. The benchmarks took a nap.

Comparing 7dcf3989 (this PR merged into the base branch) with 960d99a8 (base branch).
Only changes by more than a factor of 1.5 are shown, see the workflow run for all results.

[headtr1ck]

Thanks for picking this up again! Never splitting encoded chunks is a great default for "auto". But since "auto" now always goes through preserve_chunks, it also changes cases without encoded chunks:

import dask
import numpy as np
import xarray as xr

ds = xr.Dataset({"v": (("t", "y", "x"), np.random.rand(100, 100, 100))})  # 8 MB
ds.to_netcdf("contiguous.nc")

with dask.config.set({"array.chunk-size": "1MiB"}):
    # 1. no encoded chunks: never split
    print(xr.open_dataset("contiguous.nc", chunks="auto").v.data.chunksize)
    # 2. existing dask chunks are not split by .chunk("auto")
    print(ds.chunk(-1).chunk("auto").v.data.chunksize)
    # 3. explicit chunks are ignored along dims with uneven chunks
    uneven = ds.chunk({"y": (10, 20, 30, 40)})
    print(uneven.chunk({"t": "auto", "y": 30, "x": "auto"}).v.chunks[1])

	main	this PR
1.	(50, 50, 50)	(100, 100, 100)
2.	(50, 50, 50)	(100, 100, 100)
3.	(30, 30, 30, 10)	(10, 20, 30, 40)

1. Without encoded chunks (contiguous netCDF, scipy, ...), _get_chunk falls back to the full shape as preferred chunks, which preserve_chunks then never splits. So chunks="auto" on a large contiguous variable now creates a single chunk of the whole variable. I think preserve_chunks should only be used when the backend reports preferred_chunks, and otherwise fall back to dask's auto chunking.
2. In _maybe_chunk, the existing dask chunks are passed as previous_chunks, so .chunk("auto") can no longer split chunks that are too large, a common way to fix bad chunks. This is about dask chunks rather than encoded chunks, so I'd keep dask's behavior here.
3. In preserve_chunks, a dimension with non-uniform previous_chunks keeps them even if an explicit chunk size was given: that branch only sets max_chunks, and the explicit chunk is only applied when previous_chunk is an int.

Smaller things:

• The uniformity check previous_chunk[1:-1] == previous_chunk[:-2] treats every 2-tuple as uniform, e.g. (128, 512).
• The docstrings say multiples are taken "over the largest dimension", but the code now grows the last dimension first.
• The docs build fails because the whats-new entry needs a blank line before the "Breaking Changes" heading. Since this changes existing behavior, the entry probably belongs under "Breaking Changes" anyway.
• The comma changes in the chunks error messages and the open_zarr docstring are unrelated (and go in opposite directions).
• The PR description still describes chunks="preserve".

[jsignell]

Thanks for taking a look @headtr1ck. You can see how the thinking evolved on this PR and created some gaps. I'll update it.

[jsignell]

@headtr1ck I had thought about that issue when calling open_dataset(...).chunk("auto") as opposed to open_dataset(..., chunks="auto") (#11060 (comment)). I guess I had decided that they should have the same behavior and they should both prefer not to split encoded chunks. But do you think it would be better to have the behavior be different? Or do you prefer the original approach this PR took of introducing a new value for chunks: "preserve"?