Group chunked arrays by chunk manager, not type

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 <https://github.com/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 <https://github.com/claydugo>`_.
 
 
 Documentation

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]
 
 

xarray/tests/test_parallelcompat.py

@@ -56,14 +56,22 @@ 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"""
 
     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