Fix idxmax/idxmin silently promoting integer label coordinate dtype to float64 (GH#7527)

doc/whats-new.rst

@@ -105,6 +105,15 @@ Bug Fixes
   sparse variables sharing the same dims (:issue:`4007`).
   By `patnr <https://github.com/patnr>`_.
 
+- :py:meth:`DataArray.idxmax` / :py:meth:`DataArray.idxmin` and the
+  :py:class:`Dataset` equivalents no longer promote the dtype of the returned
+  label coordinate when no ``fill_value`` is needed. Previously, on floating-point
+  data with at least one fully-valid reduction slice, the integer coordinate
+  labels were silently cast to ``float64`` (e.g. ``idxmax(dim="y")`` returned
+  ``float64`` labels even when ``y`` was ``int64``). The coordinate dtype is now
+  preserved unless a slice is actually all-``NaN`` and gets filled
+  (:issue:`7527`).
+  By `Shurong Cao <https://github.com/CAOShurong>`_.
 
 Documentation
 ~~~~~~~~~~~~~

xarray/computation/computation.py

@@ -1015,8 +1015,10 @@ def _calc_idxminmax(
     res = index._replace(coord[(index.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
-        res = res.where(~allna, fill_value)
+        # Put the NaN values back in after removing them.
+        # We attempt to preserve dtype where we can.
+        if is_chunked_array(allna.data) or allna.any():
+            res = res.where(~allna, fill_value)
 
     # Copy attributes from argmin/argmax, if any
     res.attrs = index.attrs

xarray/core/dataset.py

@@ -9417,13 +9417,13 @@ def idxmin(
             int      int64 8B 4
             float    (y) int64 24B 4 0 2
         >>> ds.idxmin(dim="x")
-        <xarray.Dataset> Size: 52B
+        <xarray.Dataset> Size: 40B
         Dimensions:  (y: 3)
         Coordinates:
           * y        (y) int64 24B -1 0 1
         Data variables:
             int      <U1 4B 'e'
-            float    (y) object 24B 'e' 'a' 'c'
+            float    (y) <U1 12B 'e' 'a' 'c'
         """
         return self.map(
             methodcaller(
@@ -9515,13 +9515,13 @@ def idxmax(
             int      int64 8B 1
             float    (y) int64 24B 0 2 2
         >>> ds.idxmax(dim="x")
-        <xarray.Dataset> Size: 52B
+        <xarray.Dataset> Size: 40B
         Dimensions:  (y: 3)
         Coordinates:
           * y        (y) int64 24B -1 0 1
         Data variables:
             int      <U1 4B 'b'
-            float    (y) object 24B 'a' 'c' 'c'
+            float    (y) <U1 12B 'a' 'c' 'c'
         """
         return self.map(
             methodcaller(

xarray/tests/test_dataarray.py

@@ -7972,3 +7972,33 @@ def test_pyarrow_table_transposed_coords(self):
         np.testing.assert_array_equal(
             table["data"].to_pylist(), np.arange(6, dtype=float)
         )
+
+
+class TestIdxminmaxCoordDtype:
+    def test_idxmax_preserves_coord_dtype(self) -> None:
+        # Regression test for GH#7527: idxmax/idxmin should not promote the
+        # coordinate dtype of the returned labels when no fill value is needed.
+        array = xr.DataArray(
+            [
+                [2.0, 1.0, 2.0, 0.0, -2.0],
+                [-4.0, np.nan, 2.0, np.nan, -2.0],
+                [np.nan, np.nan, 1.0, np.nan, np.nan],
+            ],
+            dims=["y", "x"],
+            coords={"y": [-1, 0, 1], "x": np.arange(5.0) ** 2},
+        )
+        # No fully-masked slices, so the integer coordinate dtype is kept.
+        # The result is named after the reduced dimension and carries its
+        # (integer) label dtype.
+        assert array.idxmax(dim="y").dtype == np.int64
+        assert array.idxmin(dim="y").dtype == np.int64
+        # x is already a float64 coordinate, so reducing over it stays float64
+        assert array.idxmin(dim="x").dtype == np.float64
+        # Targeted fill only happens for the all-NaN slices, and the result is
+        # still float there (fill_value default is NaN).
+        allna = xr.DataArray(
+            [np.nan, np.nan],
+            dims=["x"],
+            coords={"x": np.array([1, 2], dtype=np.int64)},
+        )
+        assert allna.idxmax().dtype == np.float64