[draft] Accept Cubed arrays instead of dask

flox/aggregations.py

@@ -11,6 +11,7 @@
 
 from . import aggregate_flox, aggregate_npg, xrutils
 from . import xrdtypes as dtypes
+from .duck_array_ops import asarray
 
 if TYPE_CHECKING:
  FuncTuple = tuple[Callable | str, ...]
@@ -64,7 +65,7 @@ def generic_aggregate(
  f"Expected engine to be one of ['flox', 'numpy', 'numba']. Received {engine} instead."
  )
 
- group_idx = np.asarray(group_idx, like=array)
+ group_idx = asarray(group_idx, like=array)
 
  with warnings.catch_warnings():
  warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")

flox/core.py

@@ -35,7 +35,8 @@
  generic_aggregate,
 )
 from .cache import memoize
-from .xrutils import is_duck_array, is_duck_dask_array, isnull
+from .duck_array_ops import reshape
+from .xrutils import is_chunked_array, is_duck_array, is_duck_dask_array, isnull
 
 if TYPE_CHECKING:
  try:
@@ -764,7 +765,7 @@ def chunk_reduce(
  group_idx = np.broadcast_to(group_idx, array.shape[-by.ndim :])
  # always reshape to 1D along group dimensions
  newshape = array.shape[: array.ndim - by.ndim] + (math.prod(array.shape[-by.ndim :]),)
- array = array.reshape(newshape)
+ array = reshape(array, newshape)
  group_idx = group_idx.reshape(-1)
 
  assert group_idx.ndim == 1
@@ -1294,6 +1295,9 @@ def dask_groupby_agg(
 ) -> tuple[DaskArray, tuple[np.ndarray | DaskArray]]:
  import dask.array
  from dask.array.core import slices_from_chunks
+ from xarray.core.parallelcompat import get_chunked_array_type
+
+ chunkmanager = get_chunked_array_type(array)
 
  # I think _tree_reduce expects this
  assert isinstance(axis, Sequence)
@@ -1314,14 +1318,18 @@ def dask_groupby_agg(
  # Unifying chunks is necessary for argreductions.
  # We need to rechunk before zipping up with the index
  # let's always do it anyway
- if not is_duck_dask_array(by):
+ if not is_chunked_array(by):
  # chunk numpy arrays like the input array
  # This removes an extra rechunk-merge layer that would be
  # added otherwise
  chunks = tuple(array.chunks[ax] if by.shape[ax] != 1 else (1,) for ax in range(-by.ndim, 0))
 
- by = dask.array.from_array(by, chunks=chunks)
- _, (array, by) = dask.array.unify_chunks(array, inds, by, inds[-by.ndim :])
+ by = chunkmanager.from_array(
+ by,
+ chunks=chunks,
+ spec=array.spec, # cubed needs all arguments to blockwise to have same Spec
+ )
+ _, (array, by) = chunkmanager.unify_chunks(array, inds, by, inds[-by.ndim :])
 
  # preprocess the array:
  # - for argreductions, this zips the index together with the array block
@@ -1363,7 +1371,7 @@ def dask_groupby_agg(
  blockwise_method = tlz.compose(_expand_dims, blockwise_method)
 
  # apply reduction on chunk
- intermediate = dask.array.blockwise(
+ intermediate = chunkmanager.blockwise(
  partial(
  blockwise_method,
  axis=axis,
@@ -1381,9 +1389,9 @@ def dask_groupby_agg(
  inds[-by.ndim :],
  concatenate=False,
  dtype=array.dtype, # this is purely for show
- meta=array._meta,
+ # meta=array._meta,
  align_arrays=False,
- name=f"{name}-chunk-{token}",
+ # name=f"{name}-chunk-{token}",
  )
 
  group_chunks: tuple[tuple[int | float, ...]]
@@ -1392,14 +1400,13 @@ def dask_groupby_agg(
  combine: Callable[..., IntermediateDict]
  if do_simple_combine:
  combine = partial(_simple_combine, reindex=reindex)
- combine_name = "simple-combine"
  else:
  combine = partial(_grouped_combine, engine=engine, sort=sort)
- combine_name = "grouped-combine"
 
  tree_reduce = partial(
- dask.array.reductions._tree_reduce,
- name=f"{name}-reduce-{method}-{combine_name}",
+ chunkmanager.reduction,
+ func=lambda x: x,
+ # name=f"{name}-reduce-{method}-{combine_name}",
  dtype=array.dtype,
  axis=axis,
  keepdims=True,
@@ -1479,7 +1486,7 @@ def dask_groupby_agg(
  reduced = _collapse_blocks_along_axes(reduced, axis, group_chunks)
 
  # Can't use map_blocks because it forces concatenate=True along drop_axes,
- result = dask.array.blockwise(
+ result = chunkmanager.blockwise(
  _extract_result,
  out_inds,
  reduced,
@@ -1889,7 +1896,7 @@ def groupby_reduce(
  axis_ = np.core.numeric.normalize_axis_tuple(axis, array.ndim) # type: ignore
  nax = len(axis_)
 
- has_dask = is_duck_dask_array(array) or is_duck_dask_array(by_)
+ has_dask = is_chunked_array(array) or is_duck_dask_array(by_)
 
  if _is_first_last_reduction(func):
  if has_dask and nax != 1:

flox/duck_array_ops.py

@@ -0,0 +1,18 @@
+import numpy as np
+
+
+def get_array_namespace(x):
+ if hasattr(x, "__array_namespace__"):
+ return x.__array_namespace__()
+ else:
+ return np
+
+
+def reshape(array, shape):
+ xp = get_array_namespace(array)
+ return xp.reshape(array, shape)
+
+
+def asarray(obj, like):
+ xp = get_array_namespace(like)
+ return xp.asarray(obj)

flox/xrutils.py

@@ -34,11 +34,18 @@ def is_duck_array(value: Any) -> bool:
  hasattr(value, "ndim")
  and hasattr(value, "shape")
  and hasattr(value, "dtype")
- and hasattr(value, "__array_function__")
- and hasattr(value, "__array_ufunc__")
+ and (
+ (hasattr(value, "__array_function__") and hasattr(value, "__array_ufunc__"))
+ or hasattr(value, "__array_namespace__")
+ )
  )
 
 
+def is_chunked_array(x) -> bool:
+ """True if dask or cubed"""
+ return is_duck_dask_array(x) or (is_duck_array(x) and hasattr(x, "chunks"))
+
+
 def is_dask_collection(x):
  try:
  import dask