diff --git a/doc/whats-new.rst b/doc/whats-new.rst index d1505bfa081..96141c9ba89 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -105,6 +105,15 @@ Bug Fixes sparse variables sharing the same dims (:issue:`4007`). By `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 `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/computation/computation.py b/xarray/computation/computation.py index db2125cdd08..0024511f7c0 100644 --- a/xarray/computation/computation.py +++ b/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 diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 9b0a0dcd775..c282193af2c 100644 --- a/xarray/core/dataset.py +++ b/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") - Size: 52B + Size: 40B Dimensions: (y: 3) Coordinates: * y (y) int64 24B -1 0 1 Data variables: int >> ds.idxmax(dim="x") - Size: 52B + Size: 40B Dimensions: (y: 3) Coordinates: * y (y) int64 24B -1 0 1 Data variables: int 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