Add chunks='auto' support for cftime datasets

xarray/namedarray/daskmanager.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 
+from xarray.core.common import _contains_cftime_datetimes
 from xarray.core.indexing import ImplicitToExplicitIndexingAdapter
 from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint, T_ChunkedArray
 from xarray.namedarray.utils import is_duck_dask_array, module_available
@@ -16,6 +17,7 @@
         _NormalizedChunks,
         duckarray,
     )
+    from xarray.namedarray.parallelcompat import _Chunks
 
     try:
         from dask.array import Array as DaskArray
@@ -264,3 +266,50 @@ def shuffle(
         if chunks != "auto":
             raise NotImplementedError("Only chunks='auto' is supported at present.")
         return dask.array.shuffle(x, indexer, axis, chunks="auto")
+
+    def rechunk(  # type: ignore[override]
+        self,
+        data: T_ChunkedArray,
+        chunks: _NormalizedChunks | tuple[int, ...] | _Chunks,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Changes the chunking pattern of the given array.
+
+        Called when the .chunk method is called on an xarray object that is already chunked.
+
+        Parameters
+        ----------
+        data : dask array
+            Array to be rechunked.
+        chunks :  int, tuple, dict or str, optional
+            The new block dimensions to create. -1 indicates the full size of the
+            corresponding dimension. Default is "auto" which automatically
+            determines chunk sizes.
+
+        Returns
+        -------
+        chunked array
+
+        See Also
+        --------
+        dask.array.Array.rechunk
+        cubed.Array.rechunk
+        """
+
+        if _contains_cftime_datetimes(data):
+            # Preprocess chunks if they're cftime
+
+            from dask import config as dask_config
+            from dask.utils import parse_bytes
+
+            from xarray.namedarray.utils import build_chunkspec
+
+            target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))
+
+            chunks = build_chunkspec(
+                data,
+                target_chunksize=target_chunksize,
+            )
+
+        return data.rechunk(chunks, **kwargs)

xarray/namedarray/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import sys
 import warnings
 from collections.abc import Hashable, Iterable, Iterator, Mapping
 from functools import lru_cache
@@ -16,6 +17,8 @@
 
     from numpy.typing import NDArray
 
+    from xarray.namedarray.parallelcompat import T_ChunkedArray
+
     try:
         from dask.array.core import Array as DaskArray
         from dask.typing import DaskCollection
@@ -195,6 +198,30 @@ def either_dict_or_kwargs(
     return pos_kwargs
 
 
+def build_chunkspec(
+    data: T_ChunkedArray,
+    target_chunksize: int,
+) -> tuple[int, ...]:
+    """
+    Try to make chunks roughly cubic. This needs to be a bit smarter, it
+    really ought to account for xr.structure.chunks._getchunk and try to
+    use the default encoding to set the chunk size.
+    """
+    from xarray.core.formatting import first_n_items
+
+    cftime_nbytes_approx: int = sys.getsizeof(first_n_items(data, 1))  # type: ignore[no-untyped-call]
+    elements_per_chunk = target_chunksize // cftime_nbytes_approx
+    ndim = data.ndim  # type:ignore[attr-defined]
+    shape = data.shape  # type:ignore[attr-defined]
+    if ndim > 0:
+        chunk_size_per_dim = int(elements_per_chunk ** (1.0 / ndim))
+        chunks = tuple(min(chunk_size_per_dim, dim_size) for dim_size in shape)
+    else:
+        chunks = ()
+
+    return chunks
+
+
 class ReprObject:
     """Object that prints as the given value, for use with sentinel values."""
 

xarray/structure/chunks.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union, overload
 
 from xarray.core import utils
+from xarray.core.common import _contains_cftime_datetimes
 from xarray.core.utils import emit_user_level_warning
 from xarray.core.variable import IndexVariable, Variable
 from xarray.namedarray.parallelcompat import (
@@ -83,9 +84,27 @@ def _get_chunk(var: Variable, chunks, chunkmanager: ChunkManagerEntrypoint):
         for dim, preferred_chunk_sizes in zip(dims, preferred_chunk_shape, strict=True)
     )
 
-    chunk_shape = chunkmanager.normalize_chunks(
-        chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape
-    )
+    if _contains_cftime_datetimes(var):
+        # If we have cftime datetimes, need to preprocess them - we can't pass
+        # an object dtype into DaskManager.normalize_chunks.
+        from dask import config as dask_config
+        from dask.utils import parse_bytes
+
+        from xarray.namedarray.utils import build_chunkspec
+
+        target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))
+        chunk_shape = build_chunkspec(var, target_chunksize=target_chunksize)
+
+        chunk_shape = chunkmanager.normalize_chunks(
+            chunk_shape, shape=shape, previous_chunks=preferred_chunk_shape
+        )
+    else:
+        chunk_shape = chunkmanager.normalize_chunks(
+            chunk_shape,
+            shape=shape,
+            dtype=var.dtype,
+            previous_chunks=preferred_chunk_shape,
+        )
 
     # Warn where requested chunks break preferred chunks, provided that the variable
     # contains data.

xarray/tests/test_dask.py

@@ -1060,6 +1060,29 @@ def make_da():
     return da
 
 
+def make_da_cftime():
+    yrs = np.arange(2000, 2120)
+    cftime_dates = xr.date_range(
+        start=f"{yrs[0]}-01-01",
+        end=f"{yrs[-1]}-12-31",
+        freq="1YE",
+        use_cftime=True,
+    )
+    yr_array = np.tile(cftime_dates.values, (10, 1))
+    da = xr.DataArray(
+        yr_array,
+        dims=["x", "t"],
+        coords={"x": np.arange(10), "t": cftime_dates},
+        name="a",
+    ).chunk({"x": 4, "t": 5})
+    da.x.attrs["long_name"] = "x"
+    da.attrs["test"] = "test"
+    da.coords["c2"] = 0.5
+    da.coords["ndcoord"] = da.x * 2
+
+    return da
+
+
 def make_ds():
     map_ds = xr.Dataset()
     map_ds["a"] = make_da()
@@ -1141,6 +1164,14 @@ def test_auto_chunk_da(obj):
     assert actual.chunks == expected.chunks
 
 
+@pytest.mark.parametrize("obj", [make_da_cftime()])
+def test_auto_chunk_da_cftime(obj):
+    actual = obj.chunk("auto").data
+    expected = obj.data.rechunk({0: 10, 1: 120})
+    np.testing.assert_array_equal(actual, expected)
+    assert actual.chunks == expected.chunks
+
+
 def test_map_blocks_error(map_da, map_ds):
     def bad_func(darray):
         return (darray * darray.x + 5 * darray.y)[:1, :1]