Generalize handling of chunked array types

setup.cfg

@@ -185,6 +185,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-cftime.*]
 ignore_missing_imports = True
+[mypy-cubed.*]
+ignore_missing_imports = True
 [mypy-cupy.*]
 ignore_missing_imports = True
 [mypy-dask.*]

xarray/core/dataarray.py

@@ -88,6 +88,7 @@
         QueryParserOptions,
         ReindexMethodOptions,
         SideOptions,
+        T_CubedSpec,
         T_DataArray,
         T_Xarray,
     )
@@ -1156,6 +1157,8 @@ def chunk(
         token: str | None = None,
         lock: bool = False,
         inline_array: bool = False,
+        manager: Literal["dask", "cubed"] = "dask",
+        spec: T_CubedSpec = None,
         **chunks_kwargs: Any,
     ) -> T_DataArray:
         """Coerce this array's data into a dask arrays with the given chunks.

xarray/core/dataset.py

@@ -118,6 +118,7 @@
         QueryParserOptions,
         ReindexMethodOptions,
         SideOptions,
+        T_CubedSpec,
         T_Xarray,
     )
     from .weighted import DatasetWeighted
@@ -261,18 +262,27 @@ def _maybe_chunk(
     name_prefix="xarray-",
     overwrite_encoded_chunks=False,
     inline_array=False,
+    manager: Literal["dask", "cubed"] = "dask",
+    spec: T_CubedSpec = None,
 ):
     from dask.base import tokenize
 
     if chunks is not None:
         chunks = {dim: chunks[dim] for dim in var.dims if dim in chunks}
     if var.ndim:
         # when rechunking by different amounts, make sure dask names change
-        # by provinding chunks as an input to tokenize.
+        # by providing chunks as an input to tokenize.
         # subtle bugs result otherwise. see GH3350
         token2 = tokenize(name, token if token else var._data, chunks)
         name2 = f"{name_prefix}{name}-{token2}"
-        var = var.chunk(chunks, name=name2, lock=lock, inline_array=inline_array)
+        var = var.chunk(
+            chunks,
+            name=name2,
+            lock=lock,
+            inline_array=inline_array,
+            manager=manager,
+            spec=spec,
+        )
 
         if overwrite_encoded_chunks and var.chunks is not None:
             var.encoding["chunks"] = tuple(x[0] for x in var.chunks)
@@ -2166,6 +2176,8 @@ def chunk(
         token: str | None = None,
         lock: bool = False,
         inline_array: bool = False,
+        manager: Literal["dask", "cubed"] = "dask",
+        spec: T_CubedSpec = None,
         **chunks_kwargs: Any,
     ) -> T_Dataset:
         """Coerce all arrays in this dataset into dask arrays with the given
@@ -2228,7 +2240,9 @@ def chunk(
             )
 
         variables = {
-            k: _maybe_chunk(k, v, chunks, token, lock, name_prefix)
+            k: _maybe_chunk(
+                k, v, chunks, token, lock, name_prefix, manager=manager, spec=spec
+            )
             for k, v in self.variables.items()
         }
         return self._replace(variables)

xarray/core/parallelcompat.py

@@ -0,0 +1,199 @@
+"""
+The code in this module is an experiment in going from N=1 to N=2 parallel computing frameworks in xarray.
+It could later be used as the basis for a public interface allowing any N frameworks to interoperate with xarray,
+but for now it is just a private experiment.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import numpy as np
+
+from . import indexing, utils
+from .pycompat import DuckArrayModule, is_duck_dask_array
+
+CHUNK_MANAGERS = {}
+
+
+def _get_chunk_manager(name: str) -> "ChunkManager":
+    if name in CHUNK_MANAGERS:
+        chunkmanager = CHUNK_MANAGERS[name]
+        return chunkmanager()
+    else:
+        raise ImportError(f"ChunkManager {name} has not been defined")
+
+
+class ChunkManager(ABC):
+    """
+    Adapter between a particular parallel computing framework and xarray.
+
+    Attributes
+    ----------
+    array_type
+        Type of the array class this parallel computing framework provides.
+
+        Parallel frameworks need to provide an array class that supports the array API standard.
+        Used for type checking.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def chunks(arr):
+        ...
+
+    @staticmethod
+    @abstractmethod
+    def chunk(data: np.ndarray, chunks, **kwargs):
+        ...
+
+    @staticmethod
+    @abstractmethod
+    def rechunk(data: Any, chunks, **kwargs):
+        ...
+
+    @staticmethod
+    @abstractmethod
+    def compute(arr, **kwargs) -> np.ndarray:
+        ...
+
+    @staticmethod
+    @abstractmethod
+    def apply_ufunc():
+        """
+        Called inside xarray.apply_ufunc, so must be supplied for vast majority of xarray computations to be supported.
+        """
+        ...
+
+    @staticmethod
+    def map_blocks():
+        """Called by xarray.map_blocks."""
+        raise NotImplementedError()
+
+    @staticmethod
+    def blockwise():
+        """Called by some niche functions in xarray."""
+        raise NotImplementedError()
+
+
+class DaskManager(ChunkManager):
+    def __init__(self):
+        from dask.array import Array
+
+        self.array_type = Array
+
+    @staticmethod
+    def chunks(arr: "dask.array.Array"):
+        return arr.chunks
+
+    @staticmethod
+    def chunk(data: Any, chunks, **kwargs):
+        import dask.array as da
+
+        # dask-specific kwargs
+        name = kwargs.pop("name", None)
+        lock = kwargs.pop("lock", False)
+        inline_array = kwargs.pop("inline_array", False)
+
+        if is_duck_dask_array(data):
+            data = data.rechunk(chunks)
+        elif isinstance(data, DuckArrayModule("cubed").type):
+            raise TypeError("Trying to rechunk a cubed array using dask")
+        else:
+            if isinstance(data, indexing.ExplicitlyIndexed):
+                # Unambiguously handle array storage backends (like NetCDF4 and h5py)
+                # that can't handle general array indexing. For example, in netCDF4 you
+                # can do "outer" indexing along two dimensions independent, which works
+                # differently from how NumPy handles it.
+                # da.from_array works by using lazy indexing with a tuple of slices.
+                # Using OuterIndexer is a pragmatic choice: dask does not yet handle
+                # different indexing types in an explicit way:
+                # https://github.com/dask/dask/issues/2883
+                data = indexing.ImplicitToExplicitIndexingAdapter(
+                    data, indexing.OuterIndexer
+                )
+
+                # All of our lazily loaded backend array classes should use NumPy
+                # array operations.
+                dask_kwargs = {"meta": np.ndarray}
+            else:
+                dask_kwargs = {}
+
+            if utils.is_dict_like(chunks):
+                chunks = tuple(chunks.get(n, s) for n, s in enumerate(data.shape))
+
+            data = da.from_array(
+                data,
+                chunks,
+                name=name,
+                lock=lock,
+                inline_array=inline_array,
+                **dask_kwargs,
+            )
+        return data
+
+    @staticmethod
+    def rechunk(chunks, **kwargs):
+        ...
+
+    @staticmethod
+    def compute(arr, **kwargs):
+        return arr.compute(**kwargs)
+
+    @staticmethod
+    def apply_ufunc():
+        from dask.array.gufunc import apply_gufunc
+
+        ...
+
+    @staticmethod
+    def map_blocks():
+        from dask.array import map_blocks
+
+        ...
+
+    @staticmethod
+    def blockwise():
+        from dask.array import blockwise
+
+        ...
+
+
+try:
+    import dask
+
+    CHUNK_MANAGERS["dask"] = DaskManager
+except ImportError:
+    pass
+
+
+class CubedManager(ChunkManager):
+    def __init__(self):
+        from cubed import Array
+
+        self.array_type = Array
+
+    def chunk(self, data: np.ndarray, chunks, **kwargs):
+        import cubed  # type: ignore
+
+        spec = kwargs.pop("spec", None)
+
+        if isinstance(data, cubed.Array):
+            data = data.rechunk(chunks)
+        elif is_duck_dask_array(data):
+            raise TypeError("Trying to rechunk a dask array using cubed")
+        else:
+            data = cubed.from_array(
+                data,
+                chunks,
+                spec=spec,
+            )
+
+        return data
+
+
+try:
+    import cubed
+
+    CHUNK_MANAGERS["cubed"] = CubedManager
+except ImportError:
+    pass

xarray/core/pycompat.py

@@ -31,6 +31,8 @@ def __init__(self, mod):
                 duck_array_type = (duck_array_module.ndarray,)
             elif mod == "sparse":
                 duck_array_type = (duck_array_module.SparseArray,)
+            elif mod == "cubed":
+                duck_array_type = (duck_array_module.Array,)
             else:
                 raise NotImplementedError
 
@@ -67,3 +69,8 @@ def is_dask_collection(x):
 
 def is_duck_dask_array(x):
     return is_duck_array(x) and is_dask_collection(x)
+
+
+def is_chunked_array(x):
+    cubed_array_type = DuckArrayModule("cubed").type
+    return is_duck_dask_array(x) or isinstance(x, cubed_array_type)

xarray/core/types.py

@@ -19,6 +19,11 @@
     except ImportError:
         DaskArray = np.ndarray  # type: ignore
 
+    try:
+        from cubed import Spec
+    except ImportError:
+        Spec = None  # type: ignore
+
     # TODO: Turn on when https://github.com/python/mypy/issues/11871 is fixed.
     # Can be uncommented if using pyright though.
     # import sys
@@ -57,6 +62,8 @@
 VarCompatible = Union["Variable", "ScalarOrArray"]
 GroupByIncompatible = Union["Variable", "GroupBy"]
 
+T_CubedSpec = TypeVar("T_CubedSpec", bound="Spec")
+
 ErrorOptions = Literal["raise", "ignore"]
 ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"]
 

xarray/core/utils.py

@@ -290,15 +290,15 @@ def either_dict_or_kwargs(
 
 
 def _is_scalar(value, include_0d):
-    from .variable import NON_NUMPY_SUPPORTED_ARRAY_TYPES
+    from .variable import _get_NON_NUMPY_SUPPORTED_ARRAY_TYPES
 
     if include_0d:
         include_0d = getattr(value, "ndim", None) == 0
     return (
         include_0d
         or isinstance(value, (str, bytes))
         or not (
-            isinstance(value, (Iterable,) + NON_NUMPY_SUPPORTED_ARRAY_TYPES)
+            isinstance(value, (Iterable,) + _get_NON_NUMPY_SUPPORTED_ARRAY_TYPES())
             or hasattr(value, "__array_function__")
             or hasattr(value, "__array_namespace__")
         )