Repository: pydata/xarray

Issue:
Add xr.unify_chunks top level method
This should handle multiple DataArrays and Datasets. 

Implemented in #3276 as `Dataset.unify_chunks` and `DataArray.unify_chunks`

Discussion:
pep8speaks: Hello @malmans2! Thanks for opening this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:



* In the file [`xarray/core/dask_array_ops.py`](https://github.com/pydata/xarray/blob/0d20cde2eaa35ab953a376cf7d8c030e92a8d1bf/xarray/core/dask_array_ops.py):

> [Line 6:5](https://github.com/pydata/xarray/blob/0d20cde2eaa35ab953a376cf7d8c030e92a8d1bf/xarray/core/dask_array_ops.py#L6): [F401](https://duckduckgo.com/?q=pep8%20F401) '.dataarray.DataArray' imported but unused
> [Line 7:5](https://github.com/pydata/xarray/blob/0d20cde2eaa35ab953a376cf7d8c030e92a8d1bf/xarray/core/dask_array_ops.py#L7): [F401](https://duckduckgo.com/?q=pep8%20F401) '.dataset.Dataset' imported but unused
> [Line 99:5](https://github.com/pydata/xarray/blob/0d20cde2eaa35ab953a376cf7d8c030e92a8d1bf/xarray/core/dask_array_ops.py#L99): [F811](https://duckduckgo.com/?q=pep8%20F811) redefinition of unused 'DataArray' from line 6








jhamman: pinging @dcherian and/or @crusaderky who may have thoughts on this PR. 

Patch:
diff --git a/doc/api.rst b/doc/api.rst
index bddccb20..bb3a99bf 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -36,6 +36,7 @@ Top-level functions
    map_blocks
    show_versions
    set_options
+   unify_chunks
 
 Dataset
 =======
diff --git a/doc/tutorials-and-videos.rst b/doc/tutorials-and-videos.rst
index 6189fb74..22d501a4 100644
--- a/doc/tutorials-and-videos.rst
+++ b/doc/tutorials-and-videos.rst
@@ -19,7 +19,13 @@ Videos
   :card: text-center
 
   ---
+  Xdev Python Tutorial Seminar Series 2021 seminar introducing Xarray (1 of 2) | Anderson Banihirwe
+  ^^^
+  .. raw:: html
+
+    <iframe width="100%" src="https://www.youtube.com/embed/Ss4ryKukhi4" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
 
+  ---
   Xarray's virtual tutorial | October 2020 | Anderson Banihirwe, Deepak Cherian, and Martin Durant
   ^^^
   .. raw:: html
@@ -49,7 +55,6 @@ Videos
 
     <iframe width="100%" src="https://www.youtube.com/embed/J9ypQOnt5l8" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
 
-
 Books, Chapters and Articles
 -----------------------------
 
diff --git a/doc/user-guide/indexing.rst b/doc/user-guide/indexing.rst
index 14af176c..263bd4f4 100644
--- a/doc/user-guide/indexing.rst
+++ b/doc/user-guide/indexing.rst
@@ -234,9 +234,6 @@ arrays). However, you can do normal indexing with dimension names:
     ds[dict(space=[0], time=[0])]
     ds.loc[dict(time="2000-01-01")]
 
-Using indexing to *assign* values to a subset of dataset (e.g.,
-``ds[dict(space=0)] = 1``) is not yet supported.
-
 Dropping labels and dimensions
 ------------------------------
 
@@ -536,6 +533,34 @@ __ https://docs.scipy.org/doc/numpy/user/basics.indexing.html#assigning-values-t
       da.isel(x=[0, 1, 2])[1] = -1
       da
 
+You can also assign values to all variables of a :py:class:`Dataset` at once:
+
+.. ipython:: python
+
+    ds_org = xr.tutorial.open_dataset("eraint_uvz").isel(
+        latitude=slice(56, 59), longitude=slice(255, 258), level=0
+    )
+    # set all values to 0
+    ds = xr.zeros_like(ds_org)
+    ds
+
+    # by integer
+    ds[dict(latitude=2, longitude=2)] = 1
+    ds["u"]
+    ds["v"]
+
+    # by label
+    ds.loc[dict(latitude=47.25, longitude=[11.25, 12])] = 100
+    ds["u"]
+
+    # dataset as new values
+    new_dat = ds_org.loc[dict(latitude=48, longitude=[11.25, 12])]
+    new_dat
+    ds.loc[dict(latitude=47.25, longitude=[11.25, 12])] = new_dat
+    ds["u"]
+
+The dimensions can differ between the variables in the dataset, but all variables need to have at least the dimensions specified in the indexer dictionary.
+The new values must be either a scalar, a :py:class:`DataArray` or a :py:class:`Dataset` itself that contains all variables that also appear in the dataset to be modified.
 
 .. _more_advanced_indexing:
 
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index ff67ea20..f592f55a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -21,6 +21,8 @@ v0.18.3 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- New top-level function :py:func:`unify_chunks`.
+  By `Mattia Almansi <https://github.com/malmans2>`_.
 - Allow assigning values to a subset of a dataset using positional or label-based
   indexing (:issue:`3015`, :pull:`5362`).
   By `Matthias Göbel <https://github.com/matzegoebel>`_.
diff --git a/xarray/__init__.py b/xarray/__init__.py
index b4b71350..8321aba4 100644
--- a/xarray/__init__.py
+++ b/xarray/__init__.py
@@ -18,7 +18,7 @@ from .conventions import SerializationWarning, decode_cf
 from .core.alignment import align, broadcast
 from .core.combine import combine_by_coords, combine_nested
 from .core.common import ALL_DIMS, full_like, ones_like, zeros_like
-from .core.computation import apply_ufunc, corr, cov, dot, polyval, where
+from .core.computation import apply_ufunc, corr, cov, dot, polyval, unify_chunks, where
 from .core.concat import concat
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
@@ -74,6 +74,7 @@ __all__ = (
     "save_mfdataset",
     "set_options",
     "show_versions",
+    "unify_chunks",
     "where",
     "zeros_like",
     # Classes
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index fe6d8ab3..474a7218 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -1,6 +1,8 @@
 """
 Functions for applying functions that act on arrays to xarray's labeled data.
 """
+from __future__ import annotations
+
 import functools
 import itertools
 import operator
@@ -19,6 +21,7 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
+    TypeVar,
     Union,
 )
 
@@ -34,8 +37,11 @@ from .variable import Variable
 
 if TYPE_CHECKING:
     from .coordinates import Coordinates  # noqa
+    from .dataarray import DataArray
     from .dataset import Dataset
 
+    T_DSorDA = TypeVar("T_DSorDA", DataArray, Dataset)
+
 _NO_FILL_VALUE = utils.ReprObject("<no-fill-value>")
 _DEFAULT_NAME = utils.ReprObject("<default-name>")
 _JOINS_WITHOUT_FILL_VALUES = frozenset({"inner", "exact"})
@@ -1721,3 +1727,61 @@ def _calc_idxminmax(
     res.attrs = indx.attrs
 
     return res
+
+
+def unify_chunks(*objects: T_DSorDA) -> Tuple[T_DSorDA, ...]:
+    """
+    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 .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 argumets 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
+
+    # Run dask.array.core.unify_chunks
+    from dask.array.core import unify_chunks
+
+    _, dask_data = unify_chunks(*unify_chunks_args)
+    dask_data_iter = iter(dask_data)
+    out = []
+    for obj, ds in zip(objects, datasets):
+        for k, v in ds._variables.items():
+            if v.chunks is not None:
+                ds._variables[k] = v.copy(data=next(dask_data_iter))
+        out.append(obj._from_temp_dataset(ds) if isinstance(obj, DataArray) else ds)
+
+    return tuple(out)
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
index dffe97f7..eab4413d 100644
--- a/xarray/core/dataarray.py
+++ b/xarray/core/dataarray.py
@@ -44,6 +44,7 @@ from .alignment import (
 )
 from .arithmetic import DataArrayArithmetic
 from .common import AbstractArray, DataWithCoords
+from .computation import unify_chunks
 from .coordinates import (
     DataArrayCoordinates,
     assert_coordinate_consistent,
@@ -3686,8 +3687,8 @@ class DataArray(AbstractArray, DataWithCoords, DataArrayArithmetic):
         --------
         dask.array.core.unify_chunks
         """
-        ds = self._to_temp_dataset().unify_chunks()
-        return self._from_temp_dataset(ds)
+
+        return unify_chunks(self)[0]
 
     def map_blocks(
         self,
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index d03b806f..48925b70 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -53,6 +53,7 @@ from . import (
 from .alignment import _broadcast_helper, _get_broadcast_dims_map_common_coords, align
 from .arithmetic import DatasetArithmetic
 from .common import DataWithCoords, _contains_datetime_like_objects
+from .computation import unify_chunks
 from .coordinates import (
     DatasetCoordinates,
     assert_coordinate_consistent,
@@ -1487,7 +1488,7 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
         if hashable(key):
             return self._construct_dataarray(key)
         else:
-            return self._copy_listed(np.asarray(key))
+            return self._copy_listed(key)
 
     def __setitem__(self, key: Union[Hashable, List[Hashable], Mapping], value) -> None:
         """Add an array to this dataset.
@@ -4987,7 +4988,10 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
                     variables[name] = var
             else:
                 if (
-                    not numeric_only
+                    # Some reduction functions (e.g. std, var) need to run on variables
+                    # that don't have the reduce dims: PR5393
+                    not reduce_dims
+                    or not numeric_only
                     or np.issubdtype(var.dtype, np.number)
                     or (var.dtype == np.bool_)
                 ):
@@ -6563,37 +6567,7 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
         dask.array.core.unify_chunks
         """
 
-        try:
-            self.chunks
-        except ValueError:  # "inconsistent chunks"
-            pass
-        else:
-            # No variables with dask backend, or all chunks are already aligned
-            return self.copy()
-
-        # import dask is placed after the quick exit test above to allow
-        # running this method if dask isn't installed and there are no chunks
-        import dask.array
-
-        ds = self.copy()
-
-        dims_pos_map = {dim: index for index, dim in enumerate(ds.dims)}
-
-        dask_array_names = []
-        dask_unify_args = []
-        for name, variable in ds.variables.items():
-            if isinstance(variable.data, dask.array.Array):
-                dims_tuple = [dims_pos_map[dim] for dim in variable.dims]
-                dask_array_names.append(name)
-                dask_unify_args.append(variable.data)
-                dask_unify_args.append(dims_tuple)
-
-        _, rechunked_arrays = dask.array.core.unify_chunks(*dask_unify_args)
-
-        for name, new_array in zip(dask_array_names, rechunked_arrays):
-            ds.variables[name]._data = new_array
-
-        return ds
+        return unify_chunks(self)[0]
 
     def map_blocks(
         self,


Test patch:
diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py
index e80cc2f4..f790587e 100644
--- a/xarray/tests/test_dask.py
+++ b/xarray/tests/test_dask.py
@@ -1069,12 +1069,26 @@ def test_unify_chunks(map_ds):
     with pytest.raises(ValueError, match=r"inconsistent chunks"):
         ds_copy.chunks
 
-    expected_chunks = {"x": (4, 4, 2), "y": (5, 5, 5, 5), "z": (4,)}
+    expected_chunks = {"x": (4, 4, 2), "y": (5, 5, 5, 5)}
     with raise_if_dask_computes():
         actual_chunks = ds_copy.unify_chunks().chunks
-    expected_chunks == actual_chunks
+    assert actual_chunks == expected_chunks
     assert_identical(map_ds, ds_copy.unify_chunks())
 
+    out_a, out_b = xr.unify_chunks(ds_copy.cxy, ds_copy.drop_vars("cxy"))
+    assert out_a.chunks == ((4, 4, 2), (5, 5, 5, 5))
+    assert out_b.chunks == expected_chunks
+
+    # Test unordered dims
+    da = ds_copy["cxy"]
+    out_a, out_b = xr.unify_chunks(da.chunk({"x": -1}), da.T.chunk({"y": -1}))
+    assert out_a.chunks == ((4, 4, 2), (5, 5, 5, 5))
+    assert out_b.chunks == ((5, 5, 5, 5), (4, 4, 2))
+
+    # Test mismatch
+    with pytest.raises(ValueError, match=r"Dimension 'x' size mismatch: 10 != 2"):
+        xr.unify_chunks(da, da.isel(x=slice(2)))
+
 
 @pytest.mark.parametrize("obj", [make_ds(), make_da()])
 @pytest.mark.parametrize(
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 7cade410..6c7cebce 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -3305,6 +3305,11 @@ class TestDataset:
         with pytest.raises(KeyError, match=r"('var1', 'var2')"):
             data[("var1", "var2")]
 
+    def test_getitem_multiple_dtype(self):
+        keys = ["foo", 1]
+        dataset = Dataset({key: ("dim0", range(1)) for key in keys})
+        assert_identical(dataset, dataset[keys])
+
     def test_virtual_variables_default_coords(self):
         dataset = Dataset({"foo": ("x", range(10))})
         expected = DataArray(range(10), dims="x", name="x")
@@ -4947,15 +4952,16 @@ class TestDataset:
     def test_reduce_non_numeric(self):
         data1 = create_test_data(seed=44)
         data2 = create_test_data(seed=44)
-        add_vars = {"var4": ["dim1", "dim2"]}
+        add_vars = {"var4": ["dim1", "dim2"], "var5": ["dim1"]}
         for v, dims in sorted(add_vars.items()):
             size = tuple(data1.dims[d] for d in dims)
             data = np.random.randint(0, 100, size=size).astype(np.str_)
             data1[v] = (dims, data, {"foo": "variable"})
 
-        assert "var4" not in data1.mean()
+        assert "var4" not in data1.mean() and "var5" not in data1.mean()
         assert_equal(data1.mean(), data2.mean())
         assert_equal(data1.mean(dim="dim1"), data2.mean(dim="dim1"))
+        assert "var4" not in data1.mean(dim="dim2") and "var5" in data1.mean(dim="dim2")
 
     @pytest.mark.filterwarnings(
         "ignore:Once the behaviour of DataArray:DeprecationWarning"
@@ -6078,11 +6084,6 @@ class TestDataset:
 # pytest tests — new tests should go here, rather than in the class.
 
 
-@pytest.fixture(params=[None])
-def data_set(request):
-    return create_test_data(request.param)
-
-
 @pytest.mark.parametrize("test_elements", ([1, 2], np.array([1, 2]), DataArray([1, 2])))
 def test_isin(test_elements, backend):
     expected = Dataset(
@@ -6153,17 +6154,18 @@ def test_constructor_raises_with_invalid_coords(unaligned_coords):
         xr.DataArray([1, 2, 3], dims=["x"], coords=unaligned_coords)
 
 
-def test_dir_expected_attrs(data_set):
+@pytest.mark.parametrize("ds", [3], indirect=True)
+def test_dir_expected_attrs(ds):
 
     some_expected_attrs = {"pipe", "mean", "isnull", "var1", "dim2", "numbers"}
-    result = dir(data_set)
+    result = dir(ds)
     assert set(result) >= some_expected_attrs
 
 
-def test_dir_non_string(data_set):
+def test_dir_non_string(ds):
     # add a numbered key to ensure this doesn't break dir
-    data_set[5] = "foo"
-    result = dir(data_set)
+    ds[5] = "foo"
+    result = dir(ds)
     assert 5 not in result
 
     # GH2172
@@ -6173,16 +6175,16 @@ def test_dir_non_string(data_set):
     dir(x2)
 
 
-def test_dir_unicode(data_set):
-    data_set["unicode"] = "uni"
-    result = dir(data_set)
+def test_dir_unicode(ds):
+    ds["unicode"] = "uni"
+    result = dir(ds)
     assert "unicode" in result
 
 
 @pytest.fixture(params=[1])
-def ds(request):
+def ds(request, backend):
     if request.param == 1:
-        return Dataset(
+        ds = Dataset(
             dict(
                 z1=(["y", "x"], np.random.randn(2, 8)),
                 z2=(["time", "y"], np.random.randn(10, 2)),
@@ -6194,21 +6196,29 @@ def ds(request):
                 y=range(2),
             ),
         )
-
-    if request.param == 2:
-        return Dataset(
-            {
-                "z1": (["time", "y"], np.random.randn(10, 2)),
-                "z2": (["time"], np.random.randn(10)),
-                "z3": (["x", "time"], np.random.randn(8, 10)),
-            },
-            {
-                "x": ("x", np.linspace(0, 1.0, 8)),
-                "time": ("time", np.linspace(0, 1.0, 10)),
-                "c": ("y", ["a", "b"]),
-                "y": range(2),
-            },
+    elif request.param == 2:
+        ds = Dataset(
+            dict(
+                z1=(["time", "y"], np.random.randn(10, 2)),
+                z2=(["time"], np.random.randn(10)),
+                z3=(["x", "time"], np.random.randn(8, 10)),
+            ),
+            dict(
+                x=("x", np.linspace(0, 1.0, 8)),
+                time=("time", np.linspace(0, 1.0, 10)),
+                c=("y", ["a", "b"]),
+                y=range(2),
+            ),
         )
+    elif request.param == 3:
+        ds = create_test_data()
+    else:
+        raise ValueError
+
+    if backend == "dask":
+        return ds.chunk()
+
+    return ds
 
 
 def test_coarsen_absent_dims_error(ds):
@@ -6526,6 +6536,7 @@ def test_rolling_properties(ds):
 @pytest.mark.parametrize("center", (True, False, None))
 @pytest.mark.parametrize("min_periods", (1, None))
 @pytest.mark.parametrize("key", ("z1", "z2"))
+@pytest.mark.parametrize("backend", ["numpy"], indirect=True)
 def test_rolling_wrapped_bottleneck(ds, name, center, min_periods, key):
     bn = pytest.importorskip("bottleneck", minversion="1.1")
 
@@ -6551,6 +6562,7 @@ def test_rolling_wrapped_bottleneck(ds, name, center, min_periods, key):
 
 
 @requires_numbagg
+@pytest.mark.parametrize("backend", ["numpy"], indirect=True)
 def test_rolling_exp(ds):
 
     result = ds.rolling_exp(time=10, window_type="span").mean()
@@ -6558,6 +6570,7 @@ def test_rolling_exp(ds):
 
 
 @requires_numbagg
+@pytest.mark.parametrize("backend", ["numpy"], indirect=True)
 def test_rolling_exp_keep_attrs(ds):
 
     attrs_global = {"attrs": "global"}
diff --git a/xarray/tests/test_duck_array_ops.py b/xarray/tests/test_duck_array_ops.py
index 0eb00725..6d49e209 100644
--- a/xarray/tests/test_duck_array_ops.py
+++ b/xarray/tests/test_duck_array_ops.py
@@ -285,15 +285,15 @@ def assert_dask_array(da, dask):
 def test_datetime_mean(dask):
     # Note: only testing numpy, as dask is broken upstream
     da = DataArray(
-        np.array(["2010-01-01", "NaT", "2010-01-03", "NaT", "NaT"], dtype="M8"),
+        np.array(["2010-01-01", "NaT", "2010-01-03", "NaT", "NaT"], dtype="M8[ns]"),
         dims=["time"],
     )
     if dask:
         # Trigger use case where a chunk is full of NaT
         da = da.chunk({"time": 3})
 
-    expect = DataArray(np.array("2010-01-02", dtype="M8"))
-    expect_nat = DataArray(np.array("NaT", dtype="M8"))
+    expect = DataArray(np.array("2010-01-02", dtype="M8[ns]"))
+    expect_nat = DataArray(np.array("NaT", dtype="M8[ns]"))
 
     actual = da.mean()
     if dask:
@@ -889,8 +889,6 @@ def test_push_dask():
     # some chunks of size-1 with NaN
     with raise_if_dask_computes():
         actual = push(
-            dask.array.from_array(array, chunks=(1, 2, 3, 2, 2, 1, 1)),
-            axis=0,
-            n=None,
+            dask.array.from_array(array, chunks=(1, 2, 3, 2, 2, 1, 1)), axis=0, n=None
         )
     np.testing.assert_equal(actual, expected)
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index acf60465..bbd7c31f 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -2581,7 +2581,6 @@ class TestDatetimePlot(PlotTestCase):
         self.darray.plot.line()
 
 
-@pytest.mark.xfail(reason="recent versions of nc-time-axis and cftime are incompatible")
 @pytest.mark.filterwarnings("ignore:setting an array element with a sequence")
 @requires_nc_time_axis
 @requires_cftime