diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6d67d234020..19402a75d6a 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -117,6 +117,11 @@ Bug Fixes type error when applying reduction methods, due to the reduction methods being dynamically generated (:issue:`8136`). By `Andrew Scherer `_. +- :py:meth:`DataTree.chunk` now accepts a single chunk specification such as + ``"auto"`` or an integer and applies it to every dimension in the tree, as + :py:meth:`Dataset.chunk` does, instead of raising ``TypeError`` + (:issue:`11315`, :pull:`11569`). + By `fredrikblau `_. .. _`pandas-dev/pandas#64793`: https://github.com/pandas-dev/pandas/pull/64793 diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 98934f29b92..73beaec9097 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -2613,9 +2613,10 @@ def chunk( Parameters ---------- - chunks : int, tuple of int, "auto" or mapping of hashable to int or a TimeResampler, optional + chunks : int, "auto" or mapping of hashable to int or a TimeResampler, optional Chunk sizes along each dimension, e.g., ``5``, ``"auto"``, or ``{"x": 5, "y": 5}`` or ``{"x": 5, "time": TimeResampler(freq="YE")}``. + A single value is applied to every dimension in the tree. name_prefix : str, default: "xarray-" Prefix for the name of any new dask arrays. token : str, optional @@ -2650,15 +2651,24 @@ def chunk( xarray.unify_chunks dask.array.from_array """ - # don't support deprecated ways of passing chunks - if not isinstance(chunks, Mapping): - raise TypeError( - f"invalid type for chunks: {type(chunks)}. Only mappings are supported." - ) - combined_chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk") - all_dims = self._get_all_dims() + combined_chunks: Mapping[Any, T_ChunkDimFreq] + if not isinstance(chunks, Mapping): + # don't support deprecated ways of passing chunks: sequences of + # dimension-order sizes are deprecated for Dataset.chunk, and are + # additionally ambiguous for a tree whose groups need not share an + # ordering of their dimensions + if chunks is None or isinstance(chunks, tuple | list): + raise TypeError( + f"invalid type for chunks: {type(chunks)}. Only mappings and " + 'single chunk specifications (e.g. 5 or "auto") to be applied ' + "to every dimension are supported." + ) + combined_chunks = dict.fromkeys(all_dims, chunks) + else: + combined_chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk") + bad_dims = combined_chunks.keys() - all_dims if bad_dims: raise ValueError( diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 19133bfd88d..36a145e05ce 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -2784,3 +2784,20 @@ def test_chunk(self): with pytest.raises(ValueError, match="not found in data dimensions"): tree.chunk({"u": 2}) + + @pytest.mark.parametrize("chunks", ["auto", 5, "20B"]) + def test_chunk_single_spec(self, chunks): + ds1 = xr.Dataset({"a": ("x", np.arange(10))}) + ds2 = xr.Dataset({"b": ("y", np.arange(6))}) + + tree = xr.DataTree.from_dict({"/": ds1, "/group1": ds2}) + actual = tree.chunk(chunks) + + expected = xr.DataTree.from_dict( + {"/": ds1.chunk(chunks), "/group1": ds2.chunk(chunks)} + ) + + assert_identical(actual, expected) + assert actual.chunksizes == expected.chunksizes + assert set(actual.chunksizes["/"]) == {"x"} + assert set(actual.chunksizes["/group1"]) == {"y"}