wip refactor set_index

- Refactor reset_index
- Improve creating new xarray indexes from pandas indexes with proper
  propagation of variable metadata (dtype, attrs, encoding)

Author: benbovy

xarray/core/dataarray.py

@@ -46,7 +46,7 @@
 from .common import AbstractArray, DataWithCoords
 from .computation import unify_chunks
 from .coordinates import DataArrayCoordinates, assert_coordinate_consistent
-from .dataset import Dataset, split_indexes
+from .dataset import Dataset
 from .formatting import format_item
 from .indexes import Index, Indexes, default_indexes, propagate_indexes
 from .indexing import is_fancy_indexer, map_index_queries
@@ -2016,10 +2016,8 @@ def reset_index(
         --------
         DataArray.set_index
         """
-        coords, _ = split_indexes(
-            dims_or_levels, self._coords, set(), self._level_coords, drop=drop
-        )
-        return self._replace(coords=coords)
+        ds = self._to_temp_dataset().reset_index(dims_or_levels, drop=drop)
+        return self._from_temp_dataset(ds)
 
     def reorder_levels(
         self,

xarray/core/dataset.py

@@ -14,7 +14,6 @@
     Any,
     Callable,
     Collection,
-    DefaultDict,
     Dict,
     Hashable,
     Iterable,
@@ -194,64 +193,6 @@ def calculate_dimensions(variables: Mapping[Any, Variable]) -> Dict[Hashable, in
     return dims
 
 
-def split_indexes(
-    dims_or_levels: Union[Hashable, Sequence[Hashable]],
-    variables: Mapping[Any, Variable],
-    coord_names: Set[Hashable],
-    level_coords: Mapping[Any, Hashable],
-    drop: bool = False,
-) -> Tuple[Dict[Hashable, Variable], Set[Hashable]]:
-    """Extract (multi-)indexes (levels) as variables.
-
-    Not public API. Used in Dataset and DataArray reset_index
-    methods.
-    """
-    if isinstance(dims_or_levels, str) or not isinstance(dims_or_levels, Sequence):
-        dims_or_levels = [dims_or_levels]
-
-    dim_levels: DefaultDict[Any, List[Hashable]] = defaultdict(list)
-    dims = []
-    for k in dims_or_levels:
-        if k in level_coords:
-            dim_levels[level_coords[k]].append(k)
-        else:
-            dims.append(k)
-
-    vars_to_replace = {}
-    vars_to_create: Dict[Hashable, Variable] = {}
-    vars_to_remove = []
-
-    for d in dims:
-        index = variables[d].to_index()
-        if isinstance(index, pd.MultiIndex):
-            dim_levels[d] = index.names
-        else:
-            vars_to_remove.append(d)
-            if not drop:
-                vars_to_create[str(d) + "_"] = Variable(d, index, variables[d].attrs)
-
-    for d, levs in dim_levels.items():
-        index = variables[d].to_index()
-        if len(levs) == index.nlevels:
-            vars_to_remove.append(d)
-        else:
-            vars_to_replace[d] = IndexVariable(d, index.droplevel(levs))
-
-        if not drop:
-            for lev in levs:
-                idx = index.get_level_values(lev)
-                vars_to_create[idx.name] = Variable(d, idx, variables[d].attrs)
-
-    new_variables = dict(variables)
-    for v in set(vars_to_remove):
-        del new_variables[v]
-    new_variables.update(vars_to_replace)
-    new_variables.update(vars_to_create)
-    new_coord_names = (coord_names | set(vars_to_create)) - set(vars_to_remove)
-
-    return new_variables, new_coord_names
-
-
 def _assert_empty(args: tuple, msg: str = "%s") -> None:
     if args:
         raise ValueError(msg % args)
@@ -3777,14 +3718,61 @@ def reset_index(
         --------
         Dataset.set_index
         """
-        variables, coord_names = split_indexes(
-            dims_or_levels,
-            self._variables,
-            self._coord_names,
-            cast(Mapping[Hashable, Hashable], self._level_coords),
-            drop=drop,
-        )
-        return self._replace_vars_and_dims(variables, coord_names=coord_names)
+        if isinstance(dims_or_levels, str) or not isinstance(dims_or_levels, Sequence):
+            dims_or_levels = [dims_or_levels]
+
+        invalid_coords = set(dims_or_levels) - set(self.xindexes)
+        if invalid_coords:
+            raise ValueError(
+                f"{tuple(invalid_coords)} are not coordinates with an index"
+            )
+
+        drop_indexes: List[Hashable] = []
+        drop_variables: List[Hashable] = []
+        replaced_indexes: List[PandasMultiIndex] = []
+        new_indexes: Dict[Hashable, Index] = {}
+        new_variables: Dict[Hashable, IndexVariable] = {}
+
+        index_coord_names = {
+            k: coord_names
+            for _, coord_names in group_coords_by_index(self.xindexes)
+            for k in coord_names
+        }
+
+        for name in dims_or_levels:
+            index = self.xindexes[name]
+            drop_indexes += [k for k in index_coord_names[name]]
+
+            if isinstance(index, PandasMultiIndex) and name not in self.dims:
+                # special case for pd.MultiIndex (name is an index level):
+                # replace by a new index with dropped level(s) instead of just drop the index
+                # TODO: eventually extend Index API to allow this for custom multi-indexes?
+                if index not in replaced_indexes:
+                    level_names = index.index.names
+                    level_vars = {
+                        k: self._variables[k]
+                        for k in level_names
+                        if k not in dims_or_levels
+                    }
+                    idx, idx_vars = index.keep_levels(level_vars)
+                    new_indexes.update({k: idx for k in idx_vars})
+                    new_variables.update(idx_vars)
+                replaced_indexes.append(index)
+
+            if drop:
+                drop_variables.append(name)
+
+        indexes = {k: v for k, v in self.xindexes.items() if k not in drop_indexes}
+        indexes.update(new_indexes)
+
+        variables = {
+            k: v for k, v in self._variables.items() if k not in drop_variables
+        }
+        variables.update(new_variables)
+
+        coord_names = set(new_variables) | self._coord_names
+
+        return self._replace(variables, coord_names=coord_names, indexes=indexes)
 
     def reorder_levels(
         self,

xarray/core/indexes.py

@@ -11,6 +11,7 @@
     Optional,
     Tuple,
     Union,
+    cast,
 )
 
 import numpy as np
@@ -212,7 +213,10 @@ def from_variables(
 
     @classmethod
     def from_pandas_index(
-        cls, index: pd.Index, dim: Hashable
+        cls,
+        index: pd.Index,
+        dim: Hashable,
+        var_meta: Optional[Dict[Any, Dict]] = None,
     ) -> Tuple["PandasIndex", IndexVars]:
         from .variable import IndexVariable
 
@@ -223,10 +227,21 @@ def from_pandas_index(
         else:
             name = index.name
 
-        data = PandasIndexingAdapter(index)
-        index_var = IndexVariable(dim, data, fastpath=True)
+        if var_meta is None:
+            var_meta = {name: {}}
+
+        data = PandasIndexingAdapter(index, dtype=var_meta[name].get("dtype"))
+        index_var = IndexVariable(
+            dim,
+            data,
+            fastpath=True,
+            attrs=var_meta[name].get("attrs"),
+            encoding=var_meta[name].get("encoding"),
+        )
 
-        return cls(index, dim), {name: index_var}
+        return cls(index, dim, coord_dtype=var_meta[name].get("dtype")), {
+            name: index_var
+        }
 
     def to_pandas_index(self) -> pd.Index:
         return self.index
@@ -297,13 +312,11 @@ def rename(self, name_dict, dims_dict):
             return self, {}
 
         new_name = name_dict.get(self.index.name, self.index.name)
-        pd_idx = self.index.rename(new_name)
+        index = self.index.rename(new_name)
         new_dim = dims_dict.get(self.dim, self.dim)
+        var_meta = {new_name: {"dtype": self.coord_dtype}}
 
-        index, index_vars = self.from_pandas_index(pd_idx, dim=new_dim)
-        index.coord_dtype = self.coord_dtype
-
-        return index, index_vars
+        return self.from_pandas_index(index, dim=new_dim, var_meta=var_meta)
 
     def copy(self, deep=True):
         return self._replace(self.index.copy(deep=deep))
@@ -411,13 +424,12 @@ def from_variables_maybe_expand(
         """Create a new multi-index maybe by expanding an existing one with
         new variables as index levels.
 
-        the index might be created along a new dimension.
+        The index and its corresponding coordinates may be created along a new dimension.
         """
         names: List[Hashable] = []
         codes: List[List[int]] = []
         levels: List[List[int]] = []
         var_meta: Dict[str, Dict] = {}
-        level_coords_dtype: Dict[Hashable, Any] = {}
 
         _check_dim_compat({**current_variables, **variables})
 
@@ -427,20 +439,21 @@ def add_level_var(name, var):
                 "attrs": var.attrs,
                 "encoding": var.encoding,
             }
-            level_coords_dtype[name] = var.dtype
 
         if len(current_variables) > 1:
-            current_index: pd.MultiIndex = next(
-                iter(current_variables.values())
-            )._data.array
+            # expand from an existing multi-index
+            data = cast(
+                PandasMultiIndexingAdapter, next(iter(current_variables.values()))._data
+            )
+            current_index = data.array
             names.extend(current_index.names)
             codes.extend(current_index.codes)
             levels.extend(current_index.levels)
             for name in current_index.names:
                 add_level_var(name, current_variables[name])
 
         elif len(current_variables) == 1:
-            # one 1D variable (no multi-index): convert it to an index level
+            # expand from one 1D variable (no multi-index): convert it to an index level
             var = next(iter(current_variables.values()))
             new_var_name = f"{dim}_level_0"
             names.append(new_var_name)
@@ -457,27 +470,63 @@ def add_level_var(name, var):
             add_level_var(name, var)
 
         index = pd.MultiIndex(levels, codes, names=names)
-        obj = cls(index, dim, level_coords_dtype=level_coords_dtype)
-        index_vars = _create_variables_from_multiindex(index, dim, var_meta=var_meta)
 
-        return obj, index_vars
+        return cls.from_pandas_index(index, dim, var_meta=var_meta)
+
+    def keep_levels(
+        self, level_variables: Mapping[Any, "Variable"]
+    ) -> Tuple[Union["PandasMultiIndex", PandasIndex], IndexVars]:
+        """Keep only the provided levels and return a new multi-index with its
+        corresponding coordinates.
+
+        """
+        var_meta: Dict[str, Dict] = {}
+
+        for name, var in level_variables.items():
+            var_meta[name] = {
+                "dtype": var.dtype,
+                "attrs": var.attrs,
+                "encoding": var.encoding,
+            }
+
+        index = self.index.droplevel(
+            [k for k in self.index.names if k not in level_variables]
+        )
+
+        if isinstance(index, pd.MultiIndex):
+            return self.from_pandas_index(index, self.dim, var_meta=var_meta)
+        else:
+            return PandasIndex.from_pandas_index(index, self.dim, var_meta=var_meta)
 
     @classmethod
     def from_pandas_index(
-        cls, index: pd.MultiIndex, dim: Hashable
+        cls,
+        index: pd.MultiIndex,
+        dim: Hashable,
+        var_meta: Optional[Dict[Any, Dict]] = None,
     ) -> Tuple["PandasMultiIndex", IndexVars]:
-        var_meta = {}
+
+        names = []
+        idx_dtypes = {}
         for i, idx in enumerate(index.levels):
             name = idx.name or f"{dim}_level_{i}"
             if name == dim:
                 raise ValueError(
                     f"conflicting multi-index level name {name!r} with dimension {dim!r}"
                 )
-            var_meta[name] = {"dtype": idx.dtype}
+            names.append(name)
+            idx_dtypes[name] = idx.dtype
+
+        if var_meta is None:
+            var_meta = {k: {} for k in names}
+        for name, dtype in idx_dtypes.items():
+            var_meta[name]["dtype"] = var_meta[name].get("dtype", dtype)
+
+        level_coords_dtype = {k: var_meta[k]["dtype"] for k in names}
 
-        index = index.rename(var_meta.keys())
+        index = index.rename(names)
         index_vars = _create_variables_from_multiindex(index, dim, var_meta=var_meta)
-        return cls(index, dim), index_vars
+        return cls(index, dim, level_coords_dtype=level_coords_dtype), index_vars
 
     def query(self, labels, method=None, tolerance=None) -> QueryResult:
         if method is not None or tolerance is not None:
@@ -570,15 +619,19 @@ def query(self, labels, method=None, tolerance=None) -> QueryResult:
                         raise KeyError(f"not all values found in index {coord_name!r}")
 
         if new_index is not None:
+            # variable(s) attrs and encoding metadata are propagated
+            # when replacing the indexes in the resulting xarray object
+            var_meta = {k: {"dtype": v} for k, v in self.level_coords_dtype.items()}
+
             if isinstance(new_index, pd.MultiIndex):
                 new_index, new_vars = PandasMultiIndex.from_pandas_index(
-                    new_index, self.dim
+                    new_index, self.dim, var_meta=var_meta
                 )
                 dims_dict = {}
                 drop_coords = set(self.index.names) - set(new_index.index.names)
             else:
                 new_index, new_vars = PandasIndex.from_pandas_index(
-                    new_index, new_index.name
+                    new_index, new_index.name, var_meta=var_meta
                 )
                 dims_dict = {self.dim: new_index.index.name}
                 drop_coords = set(self.index.names) - {new_index.index.name} | {
@@ -602,15 +655,14 @@ def rename(self, name_dict, dims_dict):
 
         # pandas 1.3.0: could simply do `self.index.rename(names_dict)`
         new_names = [name_dict.get(k, k) for k in self.index.names]
-        pd_idx = self.index.rename(new_names)
-        new_dim = dims_dict.get(self.dim, self.dim)
+        index = self.index.rename(new_names)
 
-        index, index_vars = self.from_pandas_index(pd_idx, new_dim)
-        index.level_coords_dtype = {
-            k: v for k, v in zip(new_names, self.level_coords_dtype.values())
+        new_dim = dims_dict.get(self.dim, self.dim)
+        var_meta = {
+            k: {"dtype": v} for k, v in zip(new_names, self.level_coords_dtype.values())
         }
 
-        return index, index_vars
+        return self.from_pandas_index(index, new_dim, var_meta=var_meta)
 
 
 def remove_unused_levels_categories(index: pd.Index) -> pd.Index: