Generalize handling of chunked array types

xarray/core/dataarray.py

@@ -87,6 +87,7 @@
  QueryParserOptions,
  ReindexMethodOptions,
  SideOptions,
+ T_CubedSpec,
  T_DataArray,
  T_Xarray,
  )
@@ -1155,6 +1156,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

@@ -117,6 +117,7 @@
  QueryParserOptions,
  ReindexMethodOptions,
  SideOptions,
+ T_CubedSpec,
  T_Xarray,
  )
  from .weighted import DatasetWeighted
@@ -260,18 +261,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)
@@ -2165,6 +2175,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
@@ -2227,7 +2239,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/pycompat.py

@@ -31,10 +31,13 @@ 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.CoreArray,)
  else:
  raise NotImplementedError
 
- except ImportError: # pragma: no cover
+ except ImportError as err: # pragma: no cover
+ print(err)
  duck_array_module = None
  duck_array_version = Version("0.0.0")
  duck_array_type = ()
@@ -53,6 +56,10 @@ def __init__(self, mod):
 sparse_array_type = sp.type
 sparse_version = sp.version
 
+cub = DuckArrayModule("cubed")
+cubed_version = cub.version
+cubed_array_type = cub.type
+
 cupy_array_type = DuckArrayModule("cupy").type
 
 
@@ -67,3 +74,7 @@ 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):
+ return is_duck_dask_array(x) or isinstance(x, cubed_array_type)

xarray/core/types.py

@@ -18,6 +18,11 @@
  except ImportError:
  DaskArray = np.ndarray # type: ignore
 
+ try:
+ from cubed import Spec
+ except ImportError:
+ Spec = None # type: ignore
+
  # TODO: Turn on when https:/python/mypy/issues/11871 is fixed.
  # Can be uncommented if using pyright though.
  # import sys
@@ -50,6 +55,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/variable.py

@@ -38,9 +38,11 @@
 from .options import OPTIONS, _get_keep_attrs
 from .pycompat import (
  DuckArrayModule,
+ cubed_array_type,
  cupy_array_type,
  dask_array_type,
  integer_types,
+ is_chunked_array,
  is_duck_dask_array,
  sparse_array_type,
 )
@@ -64,6 +66,7 @@
  pd.Index,
  )
  + dask_array_type
+ + cubed_array_type
  + cupy_array_type
 )
 # https:/python/mypy/issues/224
@@ -74,6 +77,7 @@
  ErrorOptionsWithWarn,
  PadModeOptions,
  PadReflectOptions,
+ T_CubedSpec,
  T_Variable,
  )
 
@@ -462,7 +466,7 @@ def load(self, **kwargs):
  --------
  dask.array.compute
  """
- if is_duck_dask_array(self._data):
+ if is_chunked_array(self._data):
  self._data = as_compatible_data(self._data.compute(**kwargs))
  elif not is_duck_array(self._data):
  self._data = np.asarray(self._data)
@@ -1052,6 +1056,8 @@ def chunk(
  name: str = None,
  lock: bool = False,
  inline_array: bool = False,
+ manager: Literal["dask", "cubed"] = "dask",
+ spec: T_CubedSpec = None,
  **chunks_kwargs: Any,
  ) -> Variable:
  """Coerce this array's data into a dask array with the given chunks.
@@ -1093,54 +1099,89 @@ def chunk(
  xarray.unify_chunks
  dask.array.from_array
  """
- import dask.array as da
+ if manager == "dask":
+ import dask.array as da
+
+ if chunks is None:
+ warnings.warn(
+ "None value for 'chunks' is deprecated. "
+ "It will raise an error in the future. Use instead '{}'",
+ category=FutureWarning,
+ )
+ chunks = {}
 
- if chunks is None:
- warnings.warn(
- "None value for 'chunks' is deprecated. "
- "It will raise an error in the future. Use instead '{}'",
- category=FutureWarning,
- )
- chunks = {}
+ if isinstance(chunks, (float, str, int, tuple, list)):
+ pass # dask.array.from_array can handle these directly
+ else:
+ chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
 
- if isinstance(chunks, (float, str, int, tuple, list)):
- pass # dask.array.from_array can handle these directly
- else:
- chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
+  if utils.is_dict_like(chunks):
+  chunks = {
+  self.get_axis_num(dim): chunk for dim, chunk in chunks.items()
+  }
 
- if utils.is_dict_like(chunks):
- chunks = {self.get_axis_num(dim): chunk for dim, chunk in chunks.items()}
+ data = self._data
+ if is_duck_dask_array(data):
+ data = data.rechunk(chunks)
+ elif isinstance(data, cubed_array_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:/dask/dask/issues/2883
+ data = indexing.ImplicitToExplicitIndexingAdapter(
+ data, indexing.OuterIndexer
+ )
 
- data = self._data
- if is_duck_dask_array(data):
- data = data.rechunk(chunks)
- 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:/dask/dask/issues/2883
- data = indexing.ImplicitToExplicitIndexingAdapter(
- data, indexing.OuterIndexer
+ # All of our lazily loaded backend array classes should use NumPy
+ # array operations.
+ kwargs = {"meta": np.ndarray}
+ else:
+ kwargs = {}
+
+ if utils.is_dict_like(chunks):
+ chunks = tuple(chunks.get(n, s) for n, s in enumerate(self.shape))
+
+ data = da.from_array(
+ data,
+ chunks,
+ name=name,
+ lock=lock,
+ inline_array=inline_array,
+ **kwargs,
  )
 
- # All of our lazily loaded backend array classes should use NumPy
- # array operations.
- kwargs = {"meta": np.ndarray}
- else:
- kwargs = {}
+ elif manager == "cubed":
+ import cubed
+
+ chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
 
  if utils.is_dict_like(chunks):
- chunks = tuple(chunks.get(n, s) for n, s in enumerate(self.shape))
+ chunks = {
+ self.get_axis_num(dim): chunk for dim, chunk in chunks.items()
+ }
 
- data = da.from_array(
- data, chunks, name=name, lock=lock, inline_array=inline_array, **kwargs
+ data = self._data
+ if isinstance(data, cubed_array_type):
+ data = data.rechunk(chunks)
+ elif is_duck_dask_array(data):
+ raise TypeError("Trying to rechunk a dask array using cubed")
+
+ data = cubed.from_array(
+ data,
+ chunks,
+ spec=spec,
  )
 
+ else:
+ raise ValueError
+
  return self._replace(data=data)
 
  def to_numpy(self) -> np.ndarray:
@@ -1149,7 +1190,7 @@ def to_numpy(self) -> np.ndarray:
  data = self.data
 
  # TODO first attempt to call .to_numpy() once some libraries implement it
- if isinstance(data, dask_array_type):
+ if isinstance(data, (dask_array_type, cubed_array_type)):
  data = data.compute()
  if isinstance(data, cupy_array_type):
  data = data.get()
@@ -2760,7 +2801,15 @@ def values(self, values):
  f"Please use DataArray.assign_coords, Dataset.assign_coords or Dataset.assign as appropriate."
  )
 
- def chunk(self, chunks={}, name=None, lock=False, inline_array=False):
+ def chunk(
+ self,
+ chunks={},
+ name=None,
+ lock=False,
+ inline_array=False,
+ manager="dask",
+ spec=None,
+ ):
  # Dummy - do not chunk. This method is invoked e.g. by Dataset.chunk()
  return self.copy(deep=False)