Add set_xindex and drop_indexes methods

doc/api.rst

@@ -107,6 +107,7 @@ Dataset contents
    Dataset.swap_dims
    Dataset.expand_dims
    Dataset.drop_vars
+   Dataset.drop_indexes
    Dataset.drop_duplicates
    Dataset.drop_dims
    Dataset.set_coords
@@ -146,6 +147,7 @@ Indexing
    Dataset.reindex_like
    Dataset.set_index
    Dataset.reset_index
+   Dataset.set_xindex
    Dataset.reorder_levels
    Dataset.query
 
@@ -298,6 +300,7 @@ DataArray contents
    DataArray.swap_dims
    DataArray.expand_dims
    DataArray.drop_vars
+   DataArray.drop_indexes
    DataArray.drop_duplicates
    DataArray.reset_coords
    DataArray.copy
@@ -330,6 +333,7 @@ Indexing
    DataArray.reindex_like
    DataArray.set_index
    DataArray.reset_index
+   DataArray.set_xindex
    DataArray.reorder_levels
    DataArray.query
 
@@ -1080,6 +1084,7 @@ Advanced API
    Variable
    IndexVariable
    as_variable
+   Index
    Context
    register_dataset_accessor
    register_dataarray_accessor

doc/whats-new.rst

@@ -22,6 +22,10 @@ v2022.07.0 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Add :py:meth:`Dataset.set_xindex` and :py:meth:`Dataset.drop_indexes` and
+  their DataArray counterpart for setting and dropping pandas or custom indexes
+  given a set of arbitrary coordinates. (:pull:`6971`)
+  By `Benoît Bovy <https://github.com/benbovy>`_ and `Justus Magin <https://github.com/keewis>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/__init__.py

@@ -30,6 +30,7 @@
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
 from .core.extensions import register_dataarray_accessor, register_dataset_accessor
+from .core.indexes import Index
 from .core.merge import Context, MergeError, merge
 from .core.options import get_options, set_options
 from .core.parallel import map_blocks
@@ -99,6 +100,7 @@
     "Coordinate",
     "DataArray",
     "Dataset",
+    "Index",
     "IndexVariable",
     "Variable",
     # Exceptions

xarray/core/dataarray.py

@@ -2169,6 +2169,11 @@ def set_index(
         """Set DataArray (multi-)indexes using one or more existing
         coordinates.
 
+        This legacy method is limited to pandas (multi-)indexes and
+        1-dimensional "dimension" coordinates. See
+        :py:meth:`~DataArray.set_xindex` for setting a pandas or a custom
+        Xarray-compatible index from one or more arbitrary coordinates.
+
         Parameters
         ----------
         indexes : {dim: index, ...}
@@ -2213,6 +2218,7 @@ def set_index(
         See Also
         --------
         DataArray.reset_index
+        DataArray.set_xindex
         """
         ds = self._to_temp_dataset().set_index(indexes, append=append, **indexes_kwargs)
         return self._from_temp_dataset(ds)
@@ -2226,6 +2232,12 @@ def reset_index(
     ) -> DataArray:
         """Reset the specified index(es) or multi-index level(s).
 
+        This legacy method is specific to pandas (multi-)indexes and
+        1-dimensional "dimension" coordinates. See the more generic
+        :py:meth:`~DataArray.drop_indexes` and :py:meth:`~DataArray.set_xindex`
+        method to respectively drop and set pandas or custom indexes for
+        arbitrary coordinates.
+
         Parameters
         ----------
         dims_or_levels : Hashable or sequence of Hashable
@@ -2244,10 +2256,40 @@ def reset_index(
         See Also
         --------
         DataArray.set_index
+        DataArray.set_xindex
+        DataArray.drop_indexes
         """
         ds = self._to_temp_dataset().reset_index(dims_or_levels, drop=drop)
         return self._from_temp_dataset(ds)
 
+    def set_xindex(
+        self,
+        coord_names: Hashable | Sequence[Hashable],
+        index_cls: type[Index],
+        **options,
+    ) -> DataArray:
+        """Set a new, Xarray-compatible index from one or more existing
+        coordinate(s).
+
+        Parameters
+        ----------
+        coord_names : str or list
+            Name(s) of the coordinate(s) used to build the index.
+            If several names are given, their order matters.
+        index_cls : subclass of :class:`~xarray.Index`
+            The type of index to create.
+        **options
+            Options passed to the index constructor.
+
+        Returns
+        -------
+        obj : Dataset
+            Another dataarray, with this dataarray's data and with a new index.
+
+        """
+        ds = self._to_temp_dataset().set_xindex(coord_names, index_cls, **options)
+        return self._from_temp_dataset(ds)
+
     def reorder_levels(
         self: T_DataArray,
         dim_order: Mapping[Any, Sequence[int | Hashable]] | None = None,
@@ -2558,6 +2600,31 @@ def drop_vars(
         ds = self._to_temp_dataset().drop_vars(names, errors=errors)
         return self._from_temp_dataset(ds)
 
+    def drop_indexes(
+        self,
+        coord_names: Hashable | Iterable[Hashable],
+        *,
+        errors: ErrorOptions = "raise",
+    ) -> DataArray:
+        """Drop the indexes assigned to the given coordinates.
+
+        Parameters
+        ----------
+        coord_names : hashable or iterable of hashable
+            Name(s) of the coordinate(s) for which to drop the index.
+        errors : {"raise", "ignore"}, default: "raise"
+            If 'raise', raises a ValueError error if any of the coordinates
+            passed have no index or are not in the dataset.
+            If 'ignore', no error is raised.
+
+        Returns
+        -------
+        dropped : DataArray
+            A new dataarray with dropped indexes.
+        """
+        ds = self._to_temp_dataset().drop_indexes(coord_names, errors=errors)
+        return self._from_temp_dataset(ds)
+
     def drop(
         self: T_DataArray,
         labels: Mapping[Any, Any] | None = None,

xarray/core/dataset.py

@@ -3942,6 +3942,11 @@ def set_index(
         """Set Dataset (multi-)indexes using one or more existing coordinates
         or variables.
 
+        This legacy method is limited to pandas (multi-)indexes and
+        1-dimensional "dimension" coordinates. See
+        :py:meth:`~Dataset.set_xindex` for setting a pandas or a custom
+        Xarray-compatible index from one or more arbitrary coordinates.
+
         Parameters
         ----------
         indexes : {dim: index, ...}
@@ -3989,6 +3994,7 @@ def set_index(
         See Also
         --------
         Dataset.reset_index
+        Dataset.set_xindex
         Dataset.swap_dims
         """
         dim_coords = either_dict_or_kwargs(indexes, indexes_kwargs, "set_index")
@@ -4031,7 +4037,7 @@ def set_index(
                         f"dimension mismatch: try setting an index for dimension {dim!r} with "
                         f"variable {var_name!r} that has dimensions {var.dims}"
                     )
-                idx = PandasIndex.from_variables({dim: var})
+                idx = PandasIndex.from_variables({dim: var}, {})
                 idx_vars = idx.create_variables({var_name: var})
             else:
                 if append:
@@ -4080,6 +4086,12 @@ def reset_index(
     ) -> T_Dataset:
         """Reset the specified index(es) or multi-index level(s).
 
+        This legacy method is specific to pandas (multi-)indexes and
+        1-dimensional "dimension" coordinates. See the more generic
+        :py:meth:`~Dataset.drop_indexes` and :py:meth:`~Dataset.set_xindex`
+        method to respectively drop and set pandas or custom indexes for
+        arbitrary coordinates.
+
         Parameters
         ----------
         dims_or_levels : Hashable or Sequence of Hashable
@@ -4097,6 +4109,8 @@ def reset_index(
         See Also
         --------
         Dataset.set_index
+        Dataset.set_xindex
+        Dataset.drop_indexes
         """
         if isinstance(dims_or_levels, str) or not isinstance(dims_or_levels, Sequence):
             dims_or_levels = [dims_or_levels]
@@ -4149,6 +4163,101 @@ def reset_index(
 
         return self._replace(variables, coord_names=coord_names, indexes=indexes)
 
+    def set_xindex(
+        self,
+        coord_names: Hashable | Sequence[Hashable],
+        index_cls: type[Index],
+        **options,
+    ) -> Dataset:
+        """Set a new, Xarray-compatible index from one or more existing
+        coordinate(s).
+
+        Parameters
+        ----------
+        coord_names : str or list
+            Name(s) of the coordinate(s) used to build the index.
+            If several names are given, their order matters.
+        index_cls : subclass of :class:`~xarray.Index`
+            The type of index to create.
+        **options
+            Options passed to the index constructor.
+
+        Returns
+        -------
+        obj : Dataset
+            Another dataset, with this dataset's data and with a new index.
+
+        """
+        if not issubclass(index_cls, Index):
+            raise TypeError(f"{index_cls} is not a subclass of xarray.Index")
+
+        # the Sequence check is required for mypy
+        if is_scalar(coord_names) or not isinstance(coord_names, Sequence):
+            coord_names = [coord_names]
+
+        invalid_coords = set(coord_names) - self._coord_names
+
+        if invalid_coords:
+            raise ValueError(f"those coordinates don't exist: {invalid_coords}")
+
+        # 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)
+
+        new_coord_vars = index.create_variables(coord_vars)
+
+        # special case for setting a pandas multi-index from level coordinates
+        # TODO: remove it once we depreciate pandas multi-index dimension (tuple
+        # elements) coordinate
+        if isinstance(index, PandasMultiIndex):
+            coord_names = [index.dim] + list(coord_names)
+
+        variables: dict[Hashable, Variable]
+        indexes: dict[Hashable, Index]
+
+        if len(coord_names) == 1:
+            variables = self._variables.copy()
+            indexes = self._indexes.copy()
+
+            name = list(coord_names).pop()
+            if name in new_coord_vars:
+                variables[name] = new_coord_vars[name]
+            indexes[name] = index
+        else:
+            # reorder variables and indexes so that coordinates having the same
+            # index are next to each other
+            variables = {}
+            for name, var in self._variables.items():
+                if name not in coord_names:
+                    variables[name] = var
+
+            indexes = {}
+            for name, idx in self._indexes.items():
+                if name not in coord_names:
+                    indexes[name] = idx
+
+            for name in coord_names:
+                try:
+                    variables[name] = new_coord_vars[name]
+                except KeyError:
+                    variables[name] = self._variables[name]
+                indexes[name] = index
+
+        return self._replace(
+            variables=variables,
+            coord_names=self._coord_names | set(coord_names),
+            indexes=indexes,
+        )
+
     def reorder_levels(
         self: T_Dataset,
         dim_order: Mapping[Any, Sequence[int | Hashable]] | None = None,
@@ -4870,6 +4979,59 @@ def drop_vars(
             variables, coord_names=coord_names, indexes=indexes
         )
 
+    def drop_indexes(
+        self: T_Dataset,
+        coord_names: Hashable | Iterable[Hashable],
+        *,
+        errors: ErrorOptions = "raise",
+    ) -> T_Dataset:
+        """Drop the indexes assigned to the given coordinates.
+
+        Parameters
+        ----------
+        coord_names : hashable or iterable of hashable
+            Name(s) of the coordinate(s) for which to drop the index.
+        errors : {"raise", "ignore"}, default: "raise"
+            If 'raise', raises a ValueError error if any of the coordinates
+            passed have no index or are not in the dataset.
+            If 'ignore', no error is raised.
+
+        Returns
+        -------
+        dropped : Dataset
+            A new dataset with dropped indexes.
+
+        """
+        # the Iterable check is required for mypy
+        if is_scalar(coord_names) or not isinstance(coord_names, Iterable):
+            coord_names = {coord_names}
+        else:
+            coord_names = set(coord_names)
+
+        if errors == "raise":
+            invalid_coords = coord_names - self._coord_names
+            if invalid_coords:
+                raise ValueError(f"those coordinates don't exist: {invalid_coords}")
+
+            unindexed_coords = set(coord_names) - set(self._indexes)
+            if unindexed_coords:
+                raise ValueError(
+                    f"those coordinates do not have an index: {unindexed_coords}"
+                )
+
+        assert_no_index_corrupted(self.xindexes, coord_names, action="remove index(es)")
+
+        variables = {}
+        for name, var in self._variables.items():
+            if name in coord_names:
+                variables[name] = var.to_base_variable()
+            else:
+                variables[name] = var
+
+        indexes = {k: v for k, v in self._indexes.items() if k not in coord_names}
+
+        return self._replace(variables=variables, indexes=indexes)
+
     def drop(
         self: T_Dataset,
         labels=None,
@@ -7778,7 +7940,7 @@ def pad(
                 # reset default index of dimension coordinates
                 if (name,) == var.dims:
                     dim_var = {name: variables[name]}
-                    index = PandasIndex.from_variables(dim_var)
+                    index = PandasIndex.from_variables(dim_var, {})
                     index_vars = index.create_variables(dim_var)
                     indexes[name] = index
                     variables[name] = index_vars[name]