Merge branch 'main' into custom-groupers

* main:
  Refactoring/fixing zarr-python v3 incompatibilities in xarray datatrees (pydata#10020)
  Refactor calendar fixtures (pydata#10150)
  Use flox for grouped first, last. (pydata#10148)
  Update flaky pydap test (pydata#10149)

Author: dcherian

doc/whats-new.rst

@@ -41,6 +41,11 @@ New Features
 - Support reading to `GPU memory with Zarr <https://zarr.readthedocs.io/en/stable/user-guide/gpu.html>`_ (:pull:`10078`).
   By `Deepak Cherian <https://github.com/dcherian>`_.
 
+Performance
+~~~~~~~~~~~
+- :py:meth:`DatasetGroupBy.first` and :py:meth:`DatasetGroupBy.last` can now use ``flox`` if available. (:issue:`9647`)
+  By `Deepak Cherian <https://github.com/dcherian>`_.
+
 Breaking changes
 ~~~~~~~~~~~~~~~~
 - Rolled back code that would attempt to catch integer overflow when encoding
@@ -61,6 +66,10 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+
+- Fix ``open_datatree`` incompatibilities with Zarr-Python V3 and refactor
+  ``TestZarrDatatreeIO`` accordingly (:issue:`9960`, :pull:`10020`).
+  By `Alfonso Ladino-Rincon <https://github.com/aladinor>`_.
 - Default to resolution-dependent optimal integer encoding units when saving
   chunked non-nanosecond :py:class:`numpy.datetime64` or
   :py:class:`numpy.timedelta64` arrays to disk. Previously units of
@@ -92,6 +101,7 @@ Bug fixes
   datetimes and timedeltas (:issue:`8957`, :pull:`10050`).
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 
+
 Documentation
 ~~~~~~~~~~~~~
 - Better expose the :py:class:`Coordinates` class in API reference (:pull:`10000`)
@@ -174,8 +184,6 @@ New Features
   :py:class:`pandas.DatetimeIndex` (:pull:`9965`). By `Spencer Clark
   <https://github.com/spencerkclark>`_ and `Kai Mühlbauer
   <https://github.com/kmuehlbauer>`_.
-- :py:meth:`DatasetGroupBy.first` and :py:meth:`DatasetGroupBy.last` can now use ``flox`` if available. (:issue:`9647`)
-  By `Deepak Cherian <https://github.com/dcherian>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/backends/zarr.py

@@ -666,10 +666,21 @@ def open_store(
             use_zarr_fill_value_as_mask=use_zarr_fill_value_as_mask,
             zarr_format=zarr_format,
         )
+
+        from zarr import Group
+
+        group_members: dict[str, Group] = {}
         group_paths = list(_iter_zarr_groups(zarr_group, parent=group))
-        return {
+        for path in group_paths:
+            if path == group:
+                group_members[path] = zarr_group
+            else:
+                rel_path = path.removeprefix(f"{group}/")
+                group_members[path] = zarr_group[rel_path.removeprefix("/")]
+
+        out = {
             group: cls(
-                zarr_group.get(group),
+                group_store,
                 mode,
                 consolidate_on_close,
                 append_dim,
@@ -680,8 +691,9 @@ def open_store(
                 use_zarr_fill_value_as_mask,
                 cache_members=cache_members,
             )
-            for group in group_paths
+            for group, group_store in group_members.items()
         }
+        return out
 
     @classmethod
     def open_group(
@@ -1034,8 +1046,6 @@ def store(
         if self._consolidate_on_close:
             kwargs = {}
             if _zarr_v3():
-                # https://github.com/zarr-developers/zarr-python/pull/2113#issuecomment-2386718323
-                kwargs["path"] = self.zarr_group.name.lstrip("/")
                 kwargs["zarr_format"] = self.zarr_group.metadata.zarr_format
             zarr.consolidate_metadata(self.zarr_group.store, **kwargs)
 
@@ -1662,8 +1672,6 @@ def open_groups_as_dict(
         zarr_version=None,
         zarr_format=None,
     ) -> dict[str, Dataset]:
-        from xarray.core.treenode import NodePath
-
         filename_or_obj = _normalize_path(filename_or_obj)
 
         # Check for a group and make it a parent if it exists
@@ -1686,7 +1694,6 @@ def open_groups_as_dict(
         )
 
         groups_dict = {}
-
         for path_group, store in stores.items():
             store_entrypoint = StoreBackendEntrypoint()
 
@@ -1762,44 +1769,57 @@ def _get_open_params(
             consolidated = False
 
     if _zarr_v3():
-        missing_exc = ValueError
+        # TODO: replace AssertionError after https://github.com/zarr-developers/zarr-python/issues/2821 is resolved
+        missing_exc = AssertionError
     else:
         missing_exc = zarr.errors.GroupNotFoundError
 
-    if consolidated is None:
-        try:
-            zarr_group = zarr.open_consolidated(store, **open_kwargs)
-        except (ValueError, KeyError):
-            # ValueError in zarr-python 3.x, KeyError in 2.x.
+    if consolidated in [None, True]:
+        # open the root of the store, in case there is metadata consolidated there
+        group = open_kwargs.pop("path")
+
+        if consolidated:
+            # TODO: an option to pass the metadata_key keyword
+            zarr_root_group = zarr.open_consolidated(store, **open_kwargs)
+        elif consolidated is None:
+            # same but with more error handling in case no consolidated metadata found
             try:
-                zarr_group = zarr.open_group(store, **open_kwargs)
-                emit_user_level_warning(
-                    "Failed to open Zarr store with consolidated metadata, "
-                    "but successfully read with non-consolidated metadata. "
-                    "This is typically much slower for opening a dataset. "
-                    "To silence this warning, consider:\n"
-                    "1. Consolidating metadata in this existing store with "
-                    "zarr.consolidate_metadata().\n"
-                    "2. Explicitly setting consolidated=False, to avoid trying "
-                    "to read consolidate metadata, or\n"
-                    "3. Explicitly setting consolidated=True, to raise an "
-                    "error in this case instead of falling back to try "
-                    "reading non-consolidated metadata.",
-                    RuntimeWarning,
-                )
-            except missing_exc as err:
-                raise FileNotFoundError(
-                    f"No such file or directory: '{store}'"
-                ) from err
-    elif consolidated:
-        # TODO: an option to pass the metadata_key keyword
-        zarr_group = zarr.open_consolidated(store, **open_kwargs)
+                zarr_root_group = zarr.open_consolidated(store, **open_kwargs)
+            except (ValueError, KeyError):
+                # ValueError in zarr-python 3.x, KeyError in 2.x.
+                try:
+                    zarr_root_group = zarr.open_group(store, **open_kwargs)
+                    emit_user_level_warning(
+                        "Failed to open Zarr store with consolidated metadata, "
+                        "but successfully read with non-consolidated metadata. "
+                        "This is typically much slower for opening a dataset. "
+                        "To silence this warning, consider:\n"
+                        "1. Consolidating metadata in this existing store with "
+                        "zarr.consolidate_metadata().\n"
+                        "2. Explicitly setting consolidated=False, to avoid trying "
+                        "to read consolidate metadata, or\n"
+                        "3. Explicitly setting consolidated=True, to raise an "
+                        "error in this case instead of falling back to try "
+                        "reading non-consolidated metadata.",
+                        RuntimeWarning,
+                    )
+                except missing_exc as err:
+                    raise FileNotFoundError(
+                        f"No such file or directory: '{store}'"
+                    ) from err
+
+        # but the user should still receive a DataTree whose root is the group they asked for
+        if group and group != "/":
+            zarr_group = zarr_root_group[group.removeprefix("/")]
+        else:
+            zarr_group = zarr_root_group
     else:
         if _zarr_v3():
             # we have determined that we don't want to use consolidated metadata
             # so we set that to False to avoid trying to read it
             open_kwargs["use_consolidated"] = False
         zarr_group = zarr.open_group(store, **open_kwargs)
+
     close_store_on_close = zarr_group.store is not store
 
     # we use this to determine how to handle fill_value

xarray/core/datatree.py

@@ -16,7 +16,6 @@
     TYPE_CHECKING,
     Any,
     Concatenate,
-    Literal,
     NoReturn,
     ParamSpec,
     TypeVar,
@@ -1741,7 +1740,7 @@ def to_zarr(
         consolidated: bool = True,
         group: str | None = None,
         write_inherited_coords: bool = False,
-        compute: Literal[True] = True,
+        compute: bool = True,
         **kwargs,
     ):
         """

xarray/core/groupby.py

@@ -998,7 +998,7 @@ def _flox_reduce(
         dim: Dims,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ):
+    ) -> T_Xarray:
         """Adaptor function that translates our groupby API to that of flox."""
         import flox
         from flox.xarray import xarray_reduce
@@ -1121,6 +1121,8 @@ def _flox_reduce(
                     # flox always assigns an index so we must drop it here if we don't need it.
                     to_drop.append(grouper.name)
                     continue
+                # TODO: We can't simply use `self.encoded.coords` here because it corresponds to `unique_coord`,
+                # NOT `full_index`. We would need to construct a new Coordinates object, that corresponds to `full_index`.
                 new_coords.append(
                     # Using IndexVariable here ensures we reconstruct PandasMultiIndex with
                     # all associated levels properly.
@@ -1363,7 +1365,12 @@ def where(self, cond, other=dtypes.NA) -> T_Xarray:
         """
         return ops.where_method(self, cond, other)
 
-    def _first_or_last(self, op, skipna, keep_attrs):
+    def _first_or_last(
+        self,
+        op: Literal["first" | "last"],
+        skipna: bool | None,
+        keep_attrs: bool | None,
+    ):
         if all(
             isinstance(maybe_slice, slice)
             and (maybe_slice.stop == maybe_slice.start + 1)
@@ -1374,17 +1381,65 @@ def _first_or_last(self, op, skipna, keep_attrs):
             return self._obj
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=True)
-        return self.reduce(
-            op, dim=[self._group_dim], skipna=skipna, keep_attrs=keep_attrs
-        )
+        if (
+            module_available("flox", minversion="0.10.0")
+            and OPTIONS["use_flox"]
+            and contains_only_chunked_or_numpy(self._obj)
+        ):
+            result = self._flox_reduce(
+                dim=None, func=op, skipna=skipna, keep_attrs=keep_attrs
+            )
+        else:
+            result = self.reduce(
+                getattr(duck_array_ops, op),
+                dim=[self._group_dim],
+                skipna=skipna,
+                keep_attrs=keep_attrs,
+            )
+        return result
 
-    def first(self, skipna: bool | None = None, keep_attrs: bool | None = None):
-        """Return the first element of each group along the group dimension"""
-        return self._first_or_last(duck_array_ops.first, skipna, keep_attrs)
+    def first(
+        self, skipna: bool | None = None, keep_attrs: bool | None = None
+    ) -> T_Xarray:
+        """
+        Return the first element of each group along the group dimension
 
-    def last(self, skipna: bool | None = None, keep_attrs: bool | None = None):
-        """Return the last element of each group along the group dimension"""
-        return self._first_or_last(duck_array_ops.last, skipna, keep_attrs)
+        Parameters
+        ----------
+        skipna : bool or None, optional
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for float dtypes; other dtypes either do not
+            have a sentinel missing value (int) or ``skipna=True`` has not been
+            implemented (object, datetime64 or timedelta64).
+        keep_attrs : bool or None, optional
+            If True, ``attrs`` will be copied from the original
+            object to the new one.  If False, the new object will be
+            returned without attributes.
+
+        """
+        return self._first_or_last("first", skipna, keep_attrs)
+
+    def last(
+        self, skipna: bool | None = None, keep_attrs: bool | None = None
+    ) -> T_Xarray:
+        """
+        Return the last element of each group along the group dimension
+
+        Parameters
+        ----------
+        skipna : bool or None, optional
+            If True, skip missing values (as marked by NaN). By default, only
+            skips missing values for float dtypes; other dtypes either do not
+            have a sentinel missing value (int) or ``skipna=True`` has not been
+            implemented (object, datetime64 or timedelta64).
+        keep_attrs : bool or None, optional
+            If True, ``attrs`` will be copied from the original
+            object to the new one.  If False, the new object will be
+            returned without attributes.
+
+
+        """
+        return self._first_or_last("last", skipna, keep_attrs)
 
     def assign_coords(self, coords=None, **coords_kwargs):
         """Assign coordinates by group.

xarray/core/resample.py

@@ -2,7 +2,7 @@
 
 import warnings
 from collections.abc import Callable, Hashable, Iterable, Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from xarray.core._aggregations import (
     DataArrayResampleAggregations,
@@ -55,8 +55,11 @@ def _flox_reduce(
         keep_attrs: bool | None = None,
         **kwargs,
     ) -> T_Xarray:
-        result = super()._flox_reduce(dim=dim, keep_attrs=keep_attrs, **kwargs)
-        result = result.rename({RESAMPLE_DIM: self._group_dim})
+        result: T_Xarray = (
+            super()
+            ._flox_reduce(dim=dim, keep_attrs=keep_attrs, **kwargs)
+            .rename({RESAMPLE_DIM: self._group_dim})  # type: ignore[assignment]
+        )
         return result
 
     def shuffle_to_chunks(self, chunks: T_Chunks = None):
@@ -103,6 +106,21 @@ def shuffle_to_chunks(self, chunks: T_Chunks = None):
         (grouper,) = self.groupers
         return self._shuffle_obj(chunks).drop_vars(RESAMPLE_DIM)
 
+    def _first_or_last(
+        self, op: Literal["first", "last"], skipna: bool | None, keep_attrs: bool | None
+    ) -> T_Xarray:
+        from xarray.core.dataset import Dataset
+
+        result = super()._first_or_last(op=op, skipna=skipna, keep_attrs=keep_attrs)
+        if isinstance(result, Dataset):
+            # Can't do this in the base class because group_dim is RESAMPLE_DIM
+            # which is not present in the original object
+            for var in result.data_vars:
+                result._variables[var] = result._variables[var].transpose(
+                    *self._obj._variables[var].dims
+                )
+        return result
+
     def _drop_coords(self) -> T_Xarray:
         """Drop non-dimension coordinates along the resampled dimension."""
         obj = self._obj

xarray/groupers.py

@@ -254,7 +254,11 @@ def _factorize_unique(self) -> EncodedGroups:
         unique_coord = Variable(
             dims=codes.name, data=unique_values, attrs=self.group.attrs
         )
-        full_index = pd.Index(unique_values)
+        full_index = (
+            unique_values
+            if isinstance(unique_values, pd.MultiIndex)
+            else pd.Index(unique_values)
+        )
 
         return EncodedGroups(
             codes=codes,