Move chunks-related functions to a new file

Part of pydata#10089

Author: max-sixty

Diff for: xarray/__init__.py

@@ -24,7 +24,6 @@
     cross,
     dot,
     polyval,
-    unify_chunks,
     where,
 )
 from xarray.conventions import SerializationWarning, decode_cf
@@ -52,6 +51,7 @@
 from xarray.core.variable import IndexVariable, Variable, as_variable
 from xarray.namedarray.core import NamedArray
 from xarray.structure.alignment import align, broadcast
+from xarray.structure.chunks import unify_chunks
 from xarray.structure.combine import combine_by_coords, combine_nested
 from xarray.structure.concat import concat
 from xarray.structure.merge import Context, MergeError, merge

Diff for: xarray/backends/api.py

@@ -35,7 +35,6 @@
 from xarray.backends.locks import _get_scheduler
 from xarray.coders import CFDatetimeCoder, CFTimedeltaCoder
 from xarray.core import indexing
-from xarray.core.chunk import _get_chunk, _maybe_chunk
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
 from xarray.core.datatree import DataTree
@@ -45,6 +44,7 @@
 from xarray.core.utils import is_remote_uri
 from xarray.namedarray.daskmanager import DaskManager
 from xarray.namedarray.parallelcompat import guess_chunkmanager
+from xarray.structure.chunks import _get_chunk, _maybe_chunk
 from xarray.structure.combine import (
     _infer_concat_order_from_positions,
     _nested_combine,

Diff for: xarray/coding/calendar_ops.py

@@ -9,7 +9,7 @@
     _should_cftime_be_used,
     convert_times,
 )
-from xarray.computation.apply_ufunc import apply_ufunc
+
 from xarray.core.common import (
     _contains_datetime_like_objects,
     full_like,
@@ -333,6 +333,8 @@ def _decimal_year(times):
     else:
         function = _decimal_year_numpy
         kwargs = {"dtype": times.dtype}
+    from xarray.computation.apply_ufunc import apply_ufunc
+
     return apply_ufunc(
         function,
         times,

Diff for: xarray/computation/computation.py

@@ -13,12 +13,11 @@
     Callable,
     Hashable,
 )
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Literal, cast, overload
 
 import numpy as np
 
 from xarray.compat.array_api_compat import to_like_array
-from xarray.computation.apply_ufunc import apply_ufunc
 from xarray.core import dtypes, duck_array_ops, utils
 from xarray.core.common import zeros_like
 from xarray.core.duck_array_ops import datetime_to_numeric
@@ -467,6 +466,8 @@ def cross(
                 " dimensions without coordinates must have have a length of 2 or 3"
             )
 
+    from xarray.computation.apply_ufunc import apply_ufunc
+
     c = apply_ufunc(
         duck_array_ops.cross,
         a,
@@ -629,6 +630,8 @@ def dot(
     # subscripts should be passed to np.einsum as arg, not as kwargs. We need
     # to construct a partial function for apply_ufunc to work.
     func = functools.partial(duck_array_ops.einsum, subscripts, **kwargs)
+    from xarray.computation.apply_ufunc import apply_ufunc
+
     result = apply_ufunc(
         func,
         *arrays,
@@ -729,6 +732,8 @@ def where(cond, x, y, keep_attrs=None):
         keep_attrs = _get_keep_attrs(default=False)
 
     # alignment for three arguments is complicated, so don't support it yet
+    from xarray.computation.apply_ufunc import apply_ufunc
+
     result = apply_ufunc(
         duck_array_ops.where,
         cond,
@@ -951,80 +956,3 @@ def _calc_idxminmax(
     res.attrs = indx.attrs
 
     return res
-
-
-_T = TypeVar("_T", bound=Union["Dataset", "DataArray"])
-_U = TypeVar("_U", bound=Union["Dataset", "DataArray"])
-_V = TypeVar("_V", bound=Union["Dataset", "DataArray"])
-
-
-@overload
-def unify_chunks(__obj: _T) -> tuple[_T]: ...
-
-
-@overload
-def unify_chunks(__obj1: _T, __obj2: _U) -> tuple[_T, _U]: ...
-
-
-@overload
-def unify_chunks(__obj1: _T, __obj2: _U, __obj3: _V) -> tuple[_T, _U, _V]: ...
-
-
-@overload
-def unify_chunks(*objects: Dataset | DataArray) -> tuple[Dataset | DataArray, ...]: ...
-
-
-def unify_chunks(*objects: Dataset | DataArray) -> tuple[Dataset | DataArray, ...]:
-    """
-    Given any number of Dataset and/or DataArray objects, returns
-    new objects with unified chunk size along all chunked dimensions.
-
-    Returns
-    -------
-    unified (DataArray or Dataset) – Tuple of objects with the same type as
-    *objects with consistent chunk sizes for all dask-array variables
-
-    See Also
-    --------
-    dask.array.core.unify_chunks
-    """
-    from xarray.core.dataarray import DataArray
-
-    # Convert all objects to datasets
-    datasets = [
-        obj._to_temp_dataset() if isinstance(obj, DataArray) else obj.copy()
-        for obj in objects
-    ]
-
-    # Get arguments to pass into dask.array.core.unify_chunks
-    unify_chunks_args = []
-    sizes: dict[Hashable, int] = {}
-    for ds in datasets:
-        for v in ds._variables.values():
-            if v.chunks is not None:
-                # Check that sizes match across different datasets
-                for dim, size in v.sizes.items():
-                    try:
-                        if sizes[dim] != size:
-                            raise ValueError(
-                                f"Dimension {dim!r} size mismatch: {sizes[dim]} != {size}"
-                            )
-                    except KeyError:
-                        sizes[dim] = size
-                unify_chunks_args += [v._data, v._dims]
-
-    # No dask arrays: Return inputs
-    if not unify_chunks_args:
-        return objects
-
-    chunkmanager = get_chunked_array_type(*list(unify_chunks_args))
-    _, chunked_data = chunkmanager.unify_chunks(*unify_chunks_args)
-    chunked_data_iter = iter(chunked_data)
-    out: list[Dataset | DataArray] = []
-    for obj, ds in zip(objects, datasets, strict=True):
-        for k, v in ds._variables.items():
-            if v.chunks is not None:
-                ds._variables[k] = v.copy(data=next(chunked_data_iter))
-        out.append(obj._from_temp_dataset(ds) if isinstance(obj, DataArray) else ds)
-
-    return tuple(out)

Diff for: xarray/core/accessor_str.py

@@ -51,7 +51,6 @@
 
 import numpy as np
 
-from xarray.computation.computation import apply_ufunc
 from xarray.core import duck_array_ops
 from xarray.core.types import T_DataArray
 
@@ -127,6 +126,8 @@ def _apply_str_ufunc(
     if output_sizes is not None:
         dask_gufunc_kwargs["output_sizes"] = output_sizes
 
+    from xarray.computation.apply_ufunc import apply_ufunc
+
     return apply_ufunc(
         func,
         obj,

Diff for: xarray/core/common.py

@@ -491,7 +491,7 @@ def clip(
         --------
         numpy.clip : equivalent function
         """
-        from xarray.computation.computation import apply_ufunc
+        from xarray.computation.apply_ufunc import apply_ufunc
 
         if keep_attrs is None:
             # When this was a unary func, the default was True, so retaining the
@@ -1311,7 +1311,7 @@ def isnull(self, keep_attrs: bool | None = None) -> Self:
         array([False,  True, False])
         Dimensions without coordinates: x
         """
-        from xarray.computation.computation import apply_ufunc
+        from xarray.computation.apply_ufunc import apply_ufunc
 
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=False)
@@ -1354,7 +1354,7 @@ def notnull(self, keep_attrs: bool | None = None) -> Self:
         array([ True, False,  True])
         Dimensions without coordinates: x
         """
-        from xarray.computation.computation import apply_ufunc
+        from xarray.computation.apply_ufunc import apply_ufunc
 
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=False)
@@ -1393,7 +1393,7 @@ def isin(self, test_elements: Any) -> Self:
         --------
         numpy.isin
         """
-        from xarray.computation.computation import apply_ufunc
+        from xarray.computation.apply_ufunc import apply_ufunc
         from xarray.core.dataarray import DataArray
         from xarray.core.dataset import Dataset
         from xarray.core.variable import Variable
@@ -1476,7 +1476,7 @@ def astype(
         dask.array.Array.astype
         sparse.COO.astype
         """
-        from xarray.computation.computation import apply_ufunc
+        from xarray.computation.apply_ufunc import apply_ufunc
 
         kwargs = dict(order=order, casting=casting, subok=subok, copy=copy)
         kwargs = {k: v for k, v in kwargs.items() if v is not None}

Diff for: xarray/core/dataarray.py

@@ -31,7 +31,6 @@
 from xarray.coding.cftimeindex import CFTimeIndex
 from xarray.computation import computation, ops
 from xarray.computation.arithmetic import DataArrayArithmetic
-from xarray.computation.computation import unify_chunks
 from xarray.core import dtypes, indexing, utils
 from xarray.core._aggregations import DataArrayAggregations
 from xarray.core.accessor_dt import CombinedDatetimelikeAccessor
@@ -88,6 +87,7 @@
     _get_broadcast_dims_map_common_coords,
     align,
 )
+from xarray.structure.chunks import unify_chunks
 from xarray.structure.merge import PANDAS_TYPES, MergeError
 from xarray.util.deprecation_helpers import _deprecate_positional_args, deprecate_dims
 

Diff for: xarray/core/dataset.py

@@ -27,9 +27,9 @@
 import numpy as np
 from pandas.api.types import is_extension_array_dtype
 
-from xarray.core.chunk import _maybe_chunk
 from xarray.core.dataset_utils import _get_virtual_variable, _LocIndexer
 from xarray.core.dataset_variables import DataVariables
+from xarray.structure.chunks import _maybe_chunk
 
 # remove once numpy 2.0 is the oldest supported version
 try:
@@ -44,7 +44,7 @@
 from xarray.compat.array_api_compat import to_like_array
 from xarray.computation import ops
 from xarray.computation.arithmetic import DatasetArithmetic
-from xarray.computation.computation import _ensure_numeric, unify_chunks
+from xarray.computation.computation import _ensure_numeric
 from xarray.core import dtypes as xrdtypes
 from xarray.core import (
     duck_array_ops,
@@ -125,6 +125,7 @@
     _get_broadcast_dims_map_common_coords,
     align,
 )
+from xarray.structure.chunks import unify_chunks
 from xarray.structure.merge import (
     dataset_merge_method,
     dataset_update_method,

Diff for: xarray/groupers.py

@@ -17,7 +17,7 @@
 from numpy.typing import ArrayLike
 
 from xarray.coding.cftime_offsets import BaseCFTimeOffset, _new_to_legacy_freq
-from xarray.computation.computation import apply_ufunc
+from xarray.computation.apply_ufunc import apply_ufunc
 from xarray.core.coordinates import Coordinates, _coordinates_from_variable
 from xarray.core.dataarray import DataArray
 from xarray.core.duck_array_ops import array_all, isnull