Repository: pydata/xarray

Issue:
idxmin / idxmax is not parallel friendly
### Discussed in https://github.com/pydata/xarray/discussions/9421

<div type='discussions-op-text'>

<sup>Originally posted by **KBodolai** September  2, 2024</sup>
Hi there! I have a question about the chunking behaviour when using idxmin / idxmax for a chunked array.

What is the expected behaviour for the chunks after we run `idxmin` over one of the dimensions?  Naively I'd expect it to keep the chunks along the other dimensions, but that doesn't seem to be what happens: (Example below with time, x, y)

```python
import numpy as np
import xarray as xr

# create some dummy data and chunk
x, y, t = 1000, 1000, 57
rang = np.arange(t*x*y)
da = xr.DataArray(rang.reshape(t, x, y), coords={'time':range(t), 'x': range(x), 'y':range(y)})
da = da.chunk(dict(time=-1, x=256, y=256))
```

Now when I look at the array, it looks something like this:

<img width="542" alt="Screenshot 2024-09-02 at 17 06 22" src="https://github.com/user-attachments/assets/846fec75-3b45-4219-b38c-3119a933e05b">


```python
da.idxmin('time')
```

But after doing `idxmin` I get the outputs below
<img width="478" alt="Screenshot 2024-09-02 at 17 00 17" src="https://github.com/user-attachments/assets/cf4e701f-c015-45ce-979b-ab43656649c3">

My understanding is that it seems to be trying to keep the number of chunks. But oddly, when we do it for floats:

```python
da = da.astype('float32')
```

before and after doing the `idxmin` looks like this:


<img width="542" alt="Screenshot 2024-09-02 at 17 10 25" src="https://github.com/user-attachments/assets/9623c28b-b605-43d7-ad9a-8325c972d240">
<img width="477" alt="Screenshot 2024-09-02 at 17 10 11" src="https://github.com/user-attachments/assets/eab09f05-fb01-4243-8b7b-36064181b9d7">

Is this the expected behavour for this operation? I'm guessing the reshaping in the source code happens [here](https://github.com/pydata/xarray/blob/12d8cfa317b1a08748cb6a0a1b10a6efe3a7d3d2/xarray/core/computation.py#L2180), but I haven't been able to figure out how yet.

Thanks!
K.
</div>

Discussion:
dcherian: Merging on Thursday if there are no comments

Patch:
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index 6e233425..29de030a 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -32,7 +32,12 @@ from xarray.core.indexes import Index, filter_indexes_from_coords
 from xarray.core.merge import merge_attrs, merge_coordinates_without_align
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.types import Dims, T_DataArray
-from xarray.core.utils import is_dict_like, is_scalar, parse_dims_as_set, result_name
+from xarray.core.utils import (
+    is_dict_like,
+    is_scalar,
+    parse_dims_as_set,
+    result_name,
+)
 from xarray.core.variable import Variable
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
@@ -2166,19 +2171,17 @@ def _calc_idxminmax(
     indx = func(array, dim=dim, axis=None, keep_attrs=keep_attrs, skipna=skipna)
 
     # Handle chunked arrays (e.g. dask).
+    coord = array[dim]._variable.to_base_variable()
     if is_chunked_array(array.data):
         chunkmanager = get_chunked_array_type(array.data)
-        chunks = dict(zip(array.dims, array.chunks, strict=True))
-        dask_coord = chunkmanager.from_array(array[dim].data, chunks=chunks[dim])
-        data = dask_coord[duck_array_ops.ravel(indx.data)]
+        coord_array = chunkmanager.from_array(
+            array[dim].data, chunks=((array.sizes[dim],),)
+        )
+        coord = coord.copy(data=coord_array)
     else:
-        arr_coord = to_like_array(array[dim].data, array.data)
-        data = arr_coord[duck_array_ops.ravel(indx.data)]
+        coord = coord.copy(data=to_like_array(array[dim].data, array.data))
 
-    # rebuild like the argmin/max output, and rename as the dim name
-    data = duck_array_ops.reshape(data, indx.shape)
-    res = indx.copy(data=data)
-    res.name = dim
+    res = indx._replace(coord[(indx.variable,)]).rename(dim)
 
     if skipna or (skipna is None and array.dtype.kind in na_dtypes):
         # Put the NaN values back in after removing them
diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py
index 0a7b94a5..f185a05c 100644
--- a/xarray/core/indexing.py
+++ b/xarray/core/indexing.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import enum
 import functools
+import math
 import operator
 from collections import Counter, defaultdict
 from collections.abc import Callable, Hashable, Iterable, Mapping
@@ -472,12 +473,6 @@ class VectorizedIndexer(ExplicitIndexer):
         for k in key:
             if isinstance(k, slice):
                 k = as_integer_slice(k)
-            elif is_duck_dask_array(k):
-                raise ValueError(
-                    "Vectorized indexing with Dask arrays is not supported. "
-                    "Please pass a numpy array by calling ``.compute``. "
-                    "See https://github.com/dask/dask/issues/8958."
-                )
             elif is_duck_array(k):
                 if not np.issubdtype(k.dtype, np.integer):
                     raise TypeError(
@@ -1508,6 +1503,7 @@ class NumpyIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
         return self.array[key]
 
     def _vindex_get(self, indexer: VectorizedIndexer):
+        _assert_not_chunked_indexer(indexer.tuple)
         array = NumpyVIndexAdapter(self.array)
         return array[indexer.tuple]
 
@@ -1607,6 +1603,28 @@ class ArrayApiIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
         return xp.permute_dims(self.array, order)
 
 
+def _apply_vectorized_indexer_dask_wrapper(indices, coord):
+    from xarray.core.indexing import (
+        VectorizedIndexer,
+        apply_indexer,
+        as_indexable,
+    )
+
+    return apply_indexer(
+        as_indexable(coord), VectorizedIndexer((indices.squeeze(axis=-1),))
+    )
+
+
+def _assert_not_chunked_indexer(idxr: tuple[Any, ...]) -> None:
+    if any(is_chunked_array(i) for i in idxr):
+        raise ValueError(
+            "Cannot index with a chunked array indexer. "
+            "Please chunk the array you are indexing first, "
+            "and drop any indexed dimension coordinate variables. "
+            "Alternatively, call `.compute()` on any chunked arrays in the indexer."
+        )
+
+
 class DaskIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
     """Wrap a dask array to support explicit indexing."""
 
@@ -1630,7 +1648,35 @@ class DaskIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
             return value
 
     def _vindex_get(self, indexer: VectorizedIndexer):
-        return self.array.vindex[indexer.tuple]
+        try:
+            return self.array.vindex[indexer.tuple]
+        except IndexError as e:
+            # TODO: upstream to dask
+            has_dask = any(is_duck_dask_array(i) for i in indexer.tuple)
+            # this only works for "small" 1d coordinate arrays with one chunk
+            # it is intended for idxmin, idxmax, and allows indexing with
+            # the nD array output of argmin, argmax
+            if (
+                not has_dask
+                or len(indexer.tuple) > 1
+                or math.prod(self.array.numblocks) > 1
+                or self.array.ndim > 1
+            ):
+                raise e
+            (idxr,) = indexer.tuple
+            if idxr.ndim == 0:
+                return self.array[idxr.data]
+            else:
+                import dask.array
+
+                return dask.array.map_blocks(
+                    _apply_vectorized_indexer_dask_wrapper,
+                    idxr[..., np.newaxis],
+                    self.array,
+                    chunks=idxr.chunks,
+                    drop_axis=-1,
+                    dtype=self.array.dtype,
+                )
 
     def __getitem__(self, indexer: ExplicitIndexer):
         self._check_and_raise_if_non_basic_indexer(indexer)
@@ -1770,6 +1816,7 @@ class PandasIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
         | np.datetime64
         | np.timedelta64
     ):
+        _assert_not_chunked_indexer(indexer.tuple)
         key = self._prepare_key(indexer.tuple)
 
         if getattr(key, "ndim", 0) > 1:  # Return np-array if multidimensional


Test patch:
diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py
index 068f57ed..b970781f 100644
--- a/xarray/tests/test_dask.py
+++ b/xarray/tests/test_dask.py
@@ -1815,3 +1815,16 @@ def test_minimize_graph_size():
         # all the other dimensions.
         # e.g. previously for 'x',  actual == numchunks['y'] * numchunks['z']
         assert actual == numchunks[var], (actual, numchunks[var])
+
+
+def test_idxmin_chunking():
+    # GH9425
+    x, y, t = 100, 100, 10
+    rang = np.arange(t * x * y)
+    da = xr.DataArray(
+        rang.reshape(t, x, y), coords={"time": range(t), "x": range(x), "y": range(y)}
+    )
+    da = da.chunk(dict(time=-1, x=25, y=25))
+    actual = da.idxmin("time")
+    assert actual.chunksizes == {k: da.chunksizes[k] for k in ["x", "y"]}
+    assert_identical(actual, da.compute().idxmin("time"))
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index ea5186e5..f684fd06 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -4949,7 +4949,15 @@ class TestReduce1D(TestReduce):
 
         assert_identical(result2, expected2)
 
-    @pytest.mark.parametrize("use_dask", [True, False])
+    @pytest.mark.parametrize(
+        "use_dask",
+        [
+            pytest.param(
+                True, marks=pytest.mark.skipif(not has_dask, reason="no dask")
+            ),
+            False,
+        ],
+    )
     def test_idxmin(
         self,
         x: np.ndarray,
@@ -4958,16 +4966,11 @@ class TestReduce1D(TestReduce):
         nanindex: int | None,
         use_dask: bool,
     ) -> None:
-        if use_dask and not has_dask:
-            pytest.skip("requires dask")
-        if use_dask and x.dtype.kind == "M":
-            pytest.xfail("dask operation 'argmin' breaks when dtype is datetime64 (M)")
         ar0_raw = xr.DataArray(
             x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs
         )
-
         if use_dask:
-            ar0 = ar0_raw.chunk({})
+            ar0 = ar0_raw.chunk()
         else:
             ar0 = ar0_raw
 
diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py
index 698849cb..d9784e6a 100644
--- a/xarray/tests/test_indexing.py
+++ b/xarray/tests/test_indexing.py
@@ -974,7 +974,7 @@ def test_indexing_1d_object_array() -> None:
 
 
 @requires_dask
-def test_indexing_dask_array():
+def test_indexing_dask_array() -> None:
     import dask.array
 
     da = DataArray(
@@ -988,13 +988,14 @@ def test_indexing_dask_array():
 
 
 @requires_dask
-def test_indexing_dask_array_scalar():
+def test_indexing_dask_array_scalar() -> None:
     # GH4276
     import dask.array
 
     a = dask.array.from_array(np.linspace(0.0, 1.0))
     da = DataArray(a, dims="x")
     x_selector = da.argmax(dim=...)
+    assert not isinstance(x_selector, DataArray)
     with raise_if_dask_computes():
         actual = da.isel(x_selector)
     expected = da.isel(x=-1)
@@ -1002,7 +1003,7 @@ def test_indexing_dask_array_scalar():
 
 
 @requires_dask
-def test_vectorized_indexing_dask_array():
+def test_vectorized_indexing_dask_array() -> None:
     # https://github.com/pydata/xarray/issues/2511#issuecomment-563330352
     darr = DataArray(data=[0.2, 0.4, 0.6], coords={"z": range(3)}, dims=("z",))
     indexer = DataArray(
@@ -1010,12 +1011,30 @@ def test_vectorized_indexing_dask_array():
         coords={"y": range(4), "x": range(2)},
         dims=("y", "x"),
     )
-    with pytest.raises(ValueError, match="Vectorized indexing with Dask arrays"):
-        darr[indexer.chunk({"y": 2})]
+    expected = darr[indexer]
+
+    # fails because we can't index pd.Index lazily (yet).
+    # We could make this succeed by auto-chunking the values
+    # and constructing a lazy index variable, and not automatically
+    # create an index for it.
+    with pytest.raises(ValueError, match="Cannot index with"):
+        with raise_if_dask_computes():
+            darr.chunk()[indexer.chunk({"y": 2})]
+    with pytest.raises(ValueError, match="Cannot index with"):
+        with raise_if_dask_computes():
+            actual = darr[indexer.chunk({"y": 2})]
+
+    with raise_if_dask_computes():
+        actual = darr.drop_vars("z").chunk()[indexer.chunk({"y": 2})]
+    assert_identical(actual, expected.drop_vars("z"))
+
+    with raise_if_dask_computes():
+        actual_variable = darr.variable.chunk()[indexer.variable.chunk({"y": 2})]
+    assert_identical(actual_variable, expected.variable)
 
 
 @requires_dask
-def test_advanced_indexing_dask_array():
+def test_advanced_indexing_dask_array() -> None:
     # GH4663
     import dask.array as da