diff --git a/doc/whats-new.rst b/doc/whats-new.rst index d1505bfa081..6254046f6f0 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -104,6 +104,11 @@ Bug Fixes entries; ``to_dataframe`` indexes by the union of stored entries across all sparse variables sharing the same dims (:issue:`4007`). By `patnr `_. +- ``get_chunked_array_type`` no longer raises ``TypeError`` when a ``Dataset`` holds + several chunked array types that one chunk manager recognizes. Arrays are now grouped + by the chunk manager that claims them rather than by their type, so the error is + raised only for a genuine mix of frameworks such as dask and cubed (:issue:`11539`). + By `Clay Dugo `_. Documentation diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 8a68f5e9562..c28aa097a12 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -144,7 +144,8 @@ def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: """ Detects which parallel backend should be used for given set of arrays. - Also checks that all arrays are of same chunking type (i.e. not a mix of cubed and dask). + Also checks that all arrays are handled by the same chunk manager (i.e. not a mix of + cubed and dask). """ # TODO this list is probably redundant with something inside xarray.apply_ufunc @@ -156,13 +157,7 @@ def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: if is_chunked_array(a) and type(a) not in ALLOWED_NON_CHUNKED_TYPES ] - # Asserts all arrays are the same type (or numpy etc.) - chunked_array_types = {type(a) for a in chunked_arrays} - if len(chunked_array_types) > 1: - raise TypeError( - f"Mixing chunked array types is not supported, but received multiple types: {chunked_array_types}" - ) - elif len(chunked_array_types) == 0: + if not chunked_arrays: raise TypeError("Expected a chunked array but none were found") # iterate over defined chunk managers, seeing if each recognises this array type @@ -187,6 +182,13 @@ def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: elif len(selected) >= 2: raise TypeError(f"Multiple ChunkManagers recognise type {type(chunked_arr)}") else: + # Asserts the rest of the arrays are handled by it too. One chunk manager may + # recognise several array types, so their types need not all be identical. + if not all(selected[0].is_chunked_array(a) for a in chunked_arrays[1:]): + chunked_array_types = {type(a) for a in chunked_arrays} + raise TypeError( + f"Mixing chunked array types is not supported, but received multiple types: {chunked_array_types}" + ) return selected[0] diff --git a/xarray/tests/test_parallelcompat.py b/xarray/tests/test_parallelcompat.py index 456c4cae46d..9a6db074325 100644 --- a/xarray/tests/test_parallelcompat.py +++ b/xarray/tests/test_parallelcompat.py @@ -56,6 +56,14 @@ def rechunk(self, chunks, **kwargs): return copied +class OtherDummyChunkedArray(np.ndarray): + """A second, unrelated chunked array type handled by the same chunk manager.""" + + @property + def chunks(self) -> T_NormalizedChunks: + return tuple((size,) for size in self.shape) + + class DummyChunkManager(ChunkManagerEntrypoint): """Mock-up of ChunkManager class for DummyChunkedArray""" @@ -63,7 +71,7 @@ def __init__(self): self.array_cls = DummyChunkedArray def is_chunked_array(self, data: Any) -> bool: - return isinstance(data, DummyChunkedArray) + return isinstance(data, DummyChunkedArray | OtherDummyChunkedArray) def chunks(self, data: DummyChunkedArray) -> T_NormalizedChunks: return data.chunks @@ -251,6 +259,15 @@ def test_detect_dask_if_installed(self) -> None: chunk_manager = get_chunked_array_type(dask_arr) assert isinstance(chunk_manager, DaskManager) + def test_detect_several_array_types_from_one_chunkmanager( + self, register_dummy_chunkmanager + ) -> None: + dummy_arr = DummyChunkedArray([1, 2, 3]) + other_arr = OtherDummyChunkedArray([1, 2, 3]) + + chunk_manager = get_chunked_array_type(dummy_arr, other_arr) + assert isinstance(chunk_manager, DummyChunkManager) + @requires_dask def test_raise_on_mixed_array_types(self, register_dummy_chunkmanager) -> None: import dask.array as da