Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions doc/internals/how-to-create-custom-index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -224,12 +224,10 @@ 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
da_raster = da.set_xindex(["x", "y"], RasterIndex)
# Xarray creates default indexes for the 'x' and 'y' coordinates
# Use drop_existing=True to replace them with a custom index
da_raster = da.set_xindex(["x", "y"], RasterIndex, drop_existing=True)

# RasterIndex now takes care of label-based selection
selected = da_raster.sel(x=10, y=slice(20, 50))
4 changes: 4 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@ v2025.12.1 (unreleased)
New Features
~~~~~~~~~~~~

- Added ``drop_existing`` parameter to :py:meth:`Dataset.set_xindex` and
:py:meth:`DataArray.set_xindex` to allow replacing existing indexes without
needing to call :py:meth:`drop_indexes` first (:pull:`11008`).
By `Ian Hunt-Isaak <https://github.com/ianhi>`_.

Breaking Changes
~~~~~~~~~~~~~~~~
Expand Down
8 changes: 7 additions & 1 deletion xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -2866,6 +2866,7 @@ def set_xindex(
self,
coord_names: str | Sequence[Hashable],
index_cls: type[Index] | None = None,
drop_existing: bool = False,
**options,
) -> Self:
"""Set a new, Xarray-compatible index from one or more existing
Expand All @@ -2879,6 +2880,9 @@ def set_xindex(
index_cls : subclass of :class:`~xarray.indexes.Index`
The type of index to create. By default, try setting
a pandas (multi-)index from the supplied coordinates.
drop_existing : bool
Whether to drop indexes on any existing coord_names if one
is present.
**options
Options passed to the index constructor.

Expand All @@ -2888,7 +2892,9 @@ def set_xindex(
Another dataarray, with this dataarray's data and with a new index.

"""
ds = self._to_temp_dataset().set_xindex(coord_names, index_cls, **options)
ds = self._to_temp_dataset().set_xindex(
coord_names, index_cls, drop_existing, **options
)
return self._from_temp_dataset(ds)

def reorder_levels(
Expand Down
13 changes: 10 additions & 3 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4956,6 +4956,7 @@ def set_xindex(
self,
coord_names: str | Sequence[Hashable],
index_cls: type[Index] | None = None,
drop_existing: bool = False,
**options,
) -> Self:
"""Set a new, Xarray-compatible index from one or more existing
Expand All @@ -4970,6 +4971,9 @@ def set_xindex(
The type of index to create. By default, try setting
a ``PandasIndex`` if ``len(coord_names) == 1``,
otherwise a ``PandasMultiIndex``.
drop_existing : bool
Whether to drop indexes on any existing coord_names if one
is present
**options
Options passed to the index constructor.

Expand Down Expand Up @@ -5010,9 +5014,12 @@ def set_xindex(
indexed_coords = set(coord_names) & set(self._indexes)

if indexed_coords:
raise ValueError(
f"those coordinates already have an index: {indexed_coords}"
)
if drop_existing:
self.drop_indexes(indexed_coords)
else:
raise ValueError(
f"those coordinates already have an index: {indexed_coords}"
)

coord_vars = {name: self._variables[name] for name in coord_names}

Expand Down
6 changes: 6 additions & 0 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -2418,6 +2418,12 @@ 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:
# Basic test that drop_existing parameter is passed through to Dataset
da = DataArray([1, 2, 3, 4], coords={"x": ("x", [0, 1, 2, 3])}, dims="x")
result = da.set_xindex("x", PandasIndex, drop_existing=True)
assert "x" in result.xindexes

def test_dataset_getitem(self) -> None:
dv = self.ds["foo"]
assert_identical(dv, self.dv)
Expand Down
22 changes: 22 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4000,6 +4000,28 @@ class NotAnIndex: ...
with pytest.raises(ValueError, match="those coordinates already have an index"):
ds2.set_xindex("x", PandasIndex)

def test_set_xindex_drop_existing(self) -> None:
# Test that drop_existing=True allows replacing an existing index
# (the default drop_existing=False raising ValueError is tested in test_set_xindex)
ds = Dataset(coords={"x": ("x", [0, 1, 2, 3])})

# With drop_existing=True, it should succeed
result = ds.set_xindex("x", PandasIndex, drop_existing=True)
assert "x" in result.xindexes
assert isinstance(result.xindexes["x"], PandasIndex)

# Test that drop_existing=True replaces with a custom index
class CustomIndex(PandasIndex):
pass

result_custom = ds.set_xindex("x", CustomIndex, drop_existing=True)
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"])})

Expand Down
Loading