diff --git a/doc/internals/how-to-create-custom-index.rst b/doc/internals/how-to-create-custom-index.rst index 351694fc62d..2002915ac84 100644 --- a/doc/internals/how-to-create-custom-index.rst +++ b/doc/internals/how-to-create-custom-index.rst @@ -224,11 +224,9 @@ custom index to a Dataset or DataArray, e.g., using the ``RasterIndex`` above: dims=("y", "x"), ) - # Xarray create default indexes for the 'x' and 'y' coordinates - # we first need to explicitly drop it - da = da.drop_indexes(["x", "y"]) - # Build a RasterIndex from the 'x' and 'y' coordinates + # Xarray creates default indexes for the 'x' and 'y' coordinates + # this will automatically drop those indexes da_raster = da.set_xindex(["x", "y"], RasterIndex) # RasterIndex now takes care of label-based selection diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 7e3badc7143..1c463e885fc 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -14,6 +14,10 @@ v2025.12.1 (unreleased) New Features ~~~~~~~~~~~~ +- :py:meth:`Dataset.set_xindex` and :py:meth:`DataArray.set_xindex` + automatically replace any existing index being set instead of erroring + or needing needing to call :py:meth:`drop_indexes` first (:pull:`11008`). + By `Ian Hunt-Isaak `_. Breaking Changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 71d427d3db9..80f7cb6d011 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -2869,7 +2869,7 @@ def set_xindex( **options, ) -> Self: """Set a new, Xarray-compatible index from one or more existing - coordinate(s). + coordinate(s). Existing index(es) on the coord(s) will be replaced. Parameters ---------- diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index bce048048da..10b7070736b 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -4959,7 +4959,7 @@ def set_xindex( **options, ) -> Self: """Set a new, Xarray-compatible index from one or more existing - coordinate(s). + coordinate(s). Existing index(es) on the coord(s) will be replaced. Parameters ---------- @@ -5005,15 +5005,6 @@ def set_xindex( ) raise ValueError("\n".join(msg)) - # we could be more clever here (e.g., drop-in index replacement if index - # coordinates do not conflict), but let's not allow this for now - indexed_coords = set(coord_names) & set(self._indexes) - - if indexed_coords: - raise ValueError( - f"those coordinates already have an index: {indexed_coords}" - ) - coord_vars = {name: self._variables[name] for name in coord_names} index = index_cls.from_variables(coord_vars, options=options) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 7ae36421e14..df9d29843ff 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -2418,6 +2418,11 @@ def from_variables(cls, variables, options): assert "foo" in indexed.xindexes assert indexed.xindexes["foo"].opt == 1 # type: ignore[attr-defined] + def test_set_xindex_drop_existing(self) -> None: + da = DataArray([1, 2, 3, 4], coords={"x": ("x", [0, 1, 2, 3])}, dims="x") + result = da.set_xindex("x", PandasIndex) + assert "x" in result.xindexes + def test_dataset_getitem(self) -> None: dv = self.ds["foo"] assert_identical(dv, self.dv) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 6dce32aeb5c..17d3e25b642 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -3995,10 +3995,23 @@ class NotAnIndex: ... with pytest.raises(ValueError, match="those variables are data variables"): ds.set_xindex("data_var", PandasIndex) - ds2 = Dataset(coords={"x": ("x", [0, 1, 2, 3])}) + ds = Dataset(coords={"x": ("x", [0, 1, 2, 3])}) - with pytest.raises(ValueError, match="those coordinates already have an index"): - ds2.set_xindex("x", PandasIndex) + # With drop_existing=True, it should succeed + result = ds.set_xindex("x", PandasIndex) + assert "x" in result.xindexes + assert isinstance(result.xindexes["x"], PandasIndex) + + class CustomIndex(PandasIndex): + pass + + result_custom = ds.set_xindex("x", CustomIndex) + assert "x" in result_custom.xindexes + assert isinstance(result_custom.xindexes["x"], CustomIndex) + + # Verify the result is equivalent to drop_indexes + set_xindex + expected = ds.drop_indexes("x").set_xindex("x", CustomIndex) + assert_identical(result_custom, expected) def test_set_xindex_options(self) -> None: ds = Dataset(coords={"foo": ("x", ["a", "a", "b", "b"])})