diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9e6fb954a8e..646be2955cd 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -62,6 +62,9 @@ Bug fixes - XFAIL several tests which are expected to fail on ARM systems due to a ``datetime`` issue in NumPy (:issue:`2334`). By `Graham Inggs `_. +- Fix KeyError that arises when using .sel method with float values + different from coords float type (:issue:`3137`). + By `Hasan Ahmad `_. - Fixed bug in ``combine_by_coords()`` causing a `ValueError` if the input had an unused dimension with coordinates which were not monotonic (:issue:`3150`). By `Tom Nicholas `_. diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index a6edb0a9562..c5c3cadf7a2 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -11,7 +11,7 @@ from . import duck_array_ops, nputils, utils from .npcompat import DTypeLike from .pycompat import dask_array_type, integer_types -from .utils import is_dict_like +from .utils import is_dict_like, maybe_cast_to_coords_dtype def expanded_indexer(key, ndim): @@ -269,6 +269,8 @@ def remap_label_indexers(data_obj, indexers, method=None, tolerance=None): ) pos_indexers[dim] = label else: + coords_dtype = data_obj.coords[dim].dtype + label = maybe_cast_to_coords_dtype(label, coords_dtype) idxr, new_idx = convert_label_indexer(index, label, dim, method, tolerance) pos_indexers[dim] = idxr if new_idx is not None: diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 4541049d7e1..66a7adfc95e 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -85,6 +85,12 @@ def _maybe_cast_to_cftimeindex(index: pd.Index) -> pd.Index: return index +def maybe_cast_to_coords_dtype(label, coords_dtype): + if coords_dtype.kind == "f" and not isinstance(label, slice): + label = np.asarray(label, dtype=coords_dtype) + return label + + def safe_cast_to_index(array: Any) -> pd.Index: """Given an array, safely cast it to a pandas.Index. diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 8b63b650dc2..71734f09cfd 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -928,6 +928,33 @@ def test_sel_dataarray_datetime(self): result = array.sel(delta=slice(array.delta[0], array.delta[-1])) assert_equal(result, array) + def test_sel_float(self): + data_values = np.arange(4) + + # case coords are float32 and label is list of floats + float_values = [0.0, 0.111, 0.222, 0.333] + coord_values = np.asarray(float_values, dtype="float32") + array = DataArray(data_values, [("float32_coord", coord_values)]) + expected = DataArray(data_values[1:3], [("float32_coord", coord_values[1:3])]) + actual = array.sel(float32_coord=float_values[1:3]) + # case coords are float16 and label is list of floats + coord_values_16 = np.asarray(float_values, dtype="float16") + expected_16 = DataArray( + data_values[1:3], [("float16_coord", coord_values_16[1:3])] + ) + array_16 = DataArray(data_values, [("float16_coord", coord_values_16)]) + actual_16 = array_16.sel(float16_coord=float_values[1:3]) + + # case coord, label are scalars + expected_scalar = DataArray( + data_values[2], coords={"float32_coord": coord_values[2]} + ) + actual_scalar = array.sel(float32_coord=float_values[2]) + + assert_equal(expected, actual) + assert_equal(expected_scalar, actual_scalar) + assert_equal(expected_16, actual_16) + def test_sel_no_index(self): array = DataArray(np.arange(10), dims="x") assert_identical(array[0], array.sel(x=0))