Optimize grouped first, last.

1. Use flox where possible.
2. Use simple indexing where possible.

Closes pydata#9647

Author: dcherian

xarray/core/groupby.py

@@ -20,6 +20,7 @@
 from xarray.core.alignment import align, broadcast
 from xarray.core.arithmetic import DataArrayGroupbyArithmetic, DatasetGroupbyArithmetic
 from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce
+from xarray.core.computation import apply_ufunc
 from xarray.core.concat import concat
 from xarray.core.coordinates import Coordinates, _coordinates_from_variable
 from xarray.core.duck_array_ops import where
@@ -1357,7 +1358,9 @@ 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: str, skipna: bool | None, keep_attrs: bool | None):
+        from xarray.core.dataarray import DataArray
+
         if all(
             isinstance(maybe_slice, slice)
             and (maybe_slice.stop == maybe_slice.start + 1)
@@ -1368,17 +1371,95 @@ 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
+
+        def _groupby_first_last_wrapper(
+            values,
+            by,
+            *,
+            op: Literal["first", "last"],
+            skipna: bool | None,
+            group_indices,
+        ):
+            no_nans = dtypes.isdtype(
+                values.dtype, "signed integer"
+            ) or dtypes.is_string(values.dtype)
+            if (skipna or skipna is None) and not no_nans:
+                skipna = True
+            else:
+                skipna = False
+
+            if TYPE_CHECKING:
+                assert isinstance(skipna, bool)
+
+            if skipna is False or (skipna and no_nans):
+                # this is an optimization: when skipna=False, we can simply index
+                # the whole object after picking the first/last member of each group
+                # in self.encoded.group_indices
+                if op == "first":
+                    indices = [
+                        (idx.start if isinstance(idx, slice) else idx[0])
+                        for idx in group_indices
+                        if idx
+                    ]
+                else:
+                    indices = [
+                        (idx.stop - 1 if isinstance(idx, slice) else idx[-1])
+                        for idx in self.encoded.group_indices
+                        if idx
+                    ]
+                return self._obj.isel({self._group_dim: indices})
+
+            elif (
+                skipna
+                and module_available("flox", minversion="0.9.14")
+                and OPTIONS["use_flox"]
+                and contains_only_chunked_or_numpy(self._obj)
+            ):
+                import flox
+
+                result, *_ = flox.groupby_reduce(
+                    values, self.group1d.data, axis=-1, func=f"nan{op}"
+                )
+                return result
+
+            else:
+                return self.reduce(
+                    getattr(duck_array_ops, op),
+                    dim=[self._group_dim],
+                    skipna=skipna,
+                    keep_attrs=keep_attrs,
+                )
+
+        result = apply_ufunc(
+            _groupby_first_last_wrapper,
+            self._obj,
+            self.group1d,
+            input_core_dims=[[self._group_dim], [self._group_dim]],
+            output_core_dims=[[self.group1d.name]],
+            dask="allowed",
+            output_sizes={self.group1d.name: len(self)},
+            exclude_dims={self._group_dim},
+            keep_attrs=keep_attrs,
+            kwargs={
+                "op": op,
+                "skipna": skipna,
+                "group_indices": self.encoded.group_indices,
+            },
         )
+        result = result.assign_coords(self.encoded.coords)
+        result = self._maybe_unstack(result)
+        result = self._maybe_restore_empty_groups(result)
+        if isinstance(result, DataArray):
+            result = self._restore_dim_order(result)
+        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)
+        return self._first_or_last("first", skipna, keep_attrs)
 
     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)
+        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

@@ -103,6 +103,22 @@ 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: str, 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)
+        result = result.rename({RESAMPLE_DIM: self._group_dim})
+        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/tests/test_groupby.py

@@ -1618,6 +1618,8 @@ def test_groupby_first_and_last(self) -> None:
         expected = array  # should be a no-op
         assert_identical(expected, actual)
 
+        # TODO: groupby_bins too
+
     def make_groupby_multidim_example_array(self) -> DataArray:
         return DataArray(
             [[[0, 1], [2, 3]], [[5, 10], [15, 20]]],