wip refactor set_index

Author: benbovy

xarray/core/dataset.py

@@ -194,90 +194,6 @@ def calculate_dimensions(variables: Mapping[Any, Variable]) -> Dict[Hashable, in
     return dims
 
 
-def merge_indexes(
-    indexes: Mapping[Any, Union[Hashable, Sequence[Hashable]]],
-    variables: Mapping[Any, Variable],
-    coord_names: Set[Hashable],
-    append: bool = False,
-) -> Tuple[Dict[Hashable, Variable], Set[Hashable]]:
-    """Merge variables into multi-indexes.
-
-    Not public API. Used in Dataset and DataArray set_index
-    methods.
-    """
-    vars_to_replace: Dict[Hashable, Variable] = {}
-    vars_to_remove: List[Hashable] = []
-    dims_to_replace: Dict[Hashable, Hashable] = {}
-    error_msg = "{} is not the name of an existing variable."
-
-    for dim, var_names in indexes.items():
-        if isinstance(var_names, str) or not isinstance(var_names, Sequence):
-            var_names = [var_names]
-
-        names: List[Hashable] = []
-        codes: List[List[int]] = []
-        levels: List[List[int]] = []
-        current_index_variable = variables.get(dim)
-
-        for n in var_names:
-            try:
-                var = variables[n]
-            except KeyError:
-                raise ValueError(error_msg.format(n))
-            if (
-                current_index_variable is not None
-                and var.dims != current_index_variable.dims
-            ):
-                raise ValueError(
-                    f"dimension mismatch between {dim!r} {current_index_variable.dims} and {n!r} {var.dims}"
-                )
-
-        if current_index_variable is not None and append:
-            current_index = current_index_variable.to_index()
-            if isinstance(current_index, pd.MultiIndex):
-                names.extend(current_index.names)
-                codes.extend(current_index.codes)
-                levels.extend(current_index.levels)
-            else:
-                names.append(f"{dim}_level_0")
-                cat = pd.Categorical(current_index.values, ordered=True)
-                codes.append(cat.codes)
-                levels.append(cat.categories)
-
-        if not len(names) and len(var_names) == 1:
-            idx = pd.Index(variables[var_names[0]].values)
-
-        else:  # MultiIndex
-            for n in var_names:
-                try:
-                    var = variables[n]
-                except KeyError:
-                    raise ValueError(error_msg.format(n))
-                names.append(n)
-                cat = pd.Categorical(var.values, ordered=True)
-                codes.append(cat.codes)
-                levels.append(cat.categories)
-
-            idx = pd.MultiIndex(levels, codes, names=names)
-            for n in names:
-                dims_to_replace[n] = dim
-
-        vars_to_replace[dim] = IndexVariable(dim, idx)
-        vars_to_remove.extend(var_names)
-
-    new_variables = {k: v for k, v in variables.items() if k not in vars_to_remove}
-    new_variables.update(vars_to_replace)
-
-    # update dimensions if necessary, GH: 3512
-    for k, v in new_variables.items():
-        if any(d in dims_to_replace for d in v.dims):
-            new_dims = [dims_to_replace.get(d, d) for d in v.dims]
-            new_variables[k] = v._replace(dims=new_dims)
-    new_coord_names = coord_names | set(vars_to_replace)
-    new_coord_names -= set(vars_to_remove)
-    return new_variables, new_coord_names
-
-
 def split_indexes(
     dims_or_levels: Union[Hashable, Sequence[Hashable]],
     variables: Mapping[Any, Variable],
@@ -3307,7 +3223,7 @@ def _rename_dims(self, name_dict):
 
     def _rename_indexes(self, name_dict, dims_dict):
         if self._indexes is None:
-            return None, {}
+            return {}, {}
 
         indexes = {}
         variables = {}
@@ -3751,11 +3667,90 @@ def set_index(
         Dataset.reset_index
         Dataset.swap_dims
         """
-        indexes = either_dict_or_kwargs(indexes, indexes_kwargs, "set_index")
-        variables, coord_names = merge_indexes(
-            indexes, self._variables, self._coord_names, append=append
+        dim_coords = either_dict_or_kwargs(indexes, indexes_kwargs, "set_index")
+
+        new_indexes: Dict[Hashable, Index] = {}
+        new_variables: Dict[Hashable, IndexVariable] = {}
+        maybe_drop_indexes: List[Hashable] = []
+        drop_variables: List[Hashable] = []
+        replace_dims: Dict[Hashable, Hashable] = {}
+
+        index_coord_names = {
+            k: coord_names
+            for _, coord_names in group_coords_by_index(self.xindexes)
+            for k in coord_names
+        }
+
+        for dim, _var_names in dim_coords.items():
+            if isinstance(_var_names, str) or not isinstance(_var_names, Sequence):
+                var_names = [_var_names]
+            else:
+                var_names = list(_var_names)
+
+            invalid_vars = set(var_names) - set(self._variables)
+            if invalid_vars:
+                raise ValueError(
+                    ", ".join([str(v) for v in invalid_vars])
+                    + " variable(s) do not exist"
+                )
+
+            current_coord_names = index_coord_names.get(dim, [])
+
+            # drop any pre-existing index involved
+            maybe_drop_indexes.extend(current_coord_names + var_names)
+            for k in var_names:
+                maybe_drop_indexes.extend(index_coord_names.get(k, []))
+
+            drop_variables.extend(var_names)
+
+            if len(var_names) == 1 and (not append or dim not in self.xindexes):
+                var_name = var_names[0]
+                var = self._variables[var_name]
+                if var.dims != (dim,):
+                    raise ValueError(
+                        f"dimension mismatch: try setting an index for dimension {dim!r} with "
+                        f"variable {var_name!r} that has dimensions {var.dims}"
+                    )
+                idx, idx_vars = PandasIndex.from_variables({dim: var})
+            else:
+                if append:
+                    current_variables = {
+                        k: self._variables[k] for k in current_coord_names
+                    }
+                else:
+                    current_variables = {}
+                idx, idx_vars = PandasMultiIndex.from_variables_maybe_expand(
+                    dim,
+                    current_variables,
+                    {k: self._variables[k] for k in var_names},
+                )
+                for n in idx.index.names:
+                    replace_dims[n] = dim
+
+            new_indexes.update({k: idx for k in idx_vars})
+            new_variables.update(idx_vars)
+
+        indexes_: Dict[Any, Index] = {
+            k: v for k, v in self.xindexes.items() if k not in maybe_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)
+
+        # update dimensions if necessary, GH: 3512
+        for k, v in variables.items():
+            if any(d in replace_dims for d in v.dims):
+                new_dims = [replace_dims.get(d, d) for d in v.dims]
+                variables[k] = v._replace(dims=new_dims)
+
+        coord_names = set(new_variables) | self._coord_names
+
+        return self._replace_with_new_dims(
+            variables, coord_names=coord_names, indexes=indexes_
         )
-        return self._replace_vars_and_dims(variables, coord_names=coord_names)
 
     def reset_index(
         self,

xarray/core/indexes.py

@@ -312,32 +312,56 @@ def __getitem__(self, indexer: Any):
         return self._replace(self.index[indexer])
 
 
-def _create_variables_from_multiindex(index, dim, level_meta=None):
-    from .variable import IndexVariable
+def _check_dim_compat(variables: Mapping[Any, "Variable"]) -> Hashable:
+    """Check that all multi-index variable candidates share the same (single) dimension
+    and return the name of that dimension.
+
+    """
+    if any([var.ndim != 1 for var in variables.values()]):
+        raise ValueError("PandasMultiIndex only accepts 1-dimensional variables")
 
-    if level_meta is None:
-        level_meta = {}
+    dims = set([var.dims for var in variables.values()])
 
-    variables = {}
+    if len(dims) > 1:
+        raise ValueError(
+            "unmatched dimensions for variables "
+            + ", ".join([f"{k!r} {v.dims}" for k, v in variables.items()])
+        )
 
-    dim_coord_adapter = PandasMultiIndexingAdapter(index)
-    variables[dim] = IndexVariable(dim, dim_coord_adapter, fastpath=True)
+    return next(iter(dims))[0]
 
-    for level in index.names:
-        meta = level_meta.get(level, {})
+
+def _create_variables_from_multiindex(index, dim, var_meta=None):
+    from .variable import IndexVariable
+
+    if var_meta is None:
+        var_meta = {}
+
+    def create_variable(name):
+        if name == dim:
+            level = None
+        else:
+            level = name
+        meta = var_meta.get(name, {})
         data = PandasMultiIndexingAdapter(index, dtype=meta.get("dtype"), level=level)
-        variables[level] = IndexVariable(
+        return IndexVariable(
             dim,
             data,
             attrs=meta.get("attrs"),
             encoding=meta.get("encoding"),
             fastpath=True,
         )
 
+    variables = {}
+    variables[dim] = create_variable(dim)
+    for level in index.names:
+        variables[level] = create_variable(level)
+
     return variables
 
 
 class PandasMultiIndex(PandasIndex):
+    """Wrap a pandas.MultiIndex as an xarray compatible index."""
 
     level_coords_dtype: Dict[str, Any]
 
@@ -358,51 +382,101 @@ def _replace(self, index, dim=None, level_coords_dtype=None) -> "PandasMultiInde
         return type(self)(index, dim, level_coords_dtype)
 
     @classmethod
-    def from_variables(cls, variables: Mapping[Any, "Variable"]):
-        if any([var.ndim != 1 for var in variables.values()]):
-            raise ValueError("PandasMultiIndex only accepts 1-dimensional variables")
-
-        dims = set([var.dims for var in variables.values()])
-        if len(dims) != 1:
-            raise ValueError(
-                "unmatched dimensions for variables "
-                + ",".join([str(k) for k in variables])
-            )
+    def from_variables(
+        cls, variables: Mapping[Any, "Variable"]
+    ) -> Tuple["PandasMultiIndex", IndexVars]:
+        dim = _check_dim_compat(variables)
 
-        dim = next(iter(dims))[0]
         index = pd.MultiIndex.from_arrays(
             [var.values for var in variables.values()], names=variables.keys()
         )
         level_coords_dtype = {name: var.dtype for name, var in variables.items()}
         obj = cls(index, dim, level_coords_dtype=level_coords_dtype)
 
-        level_meta = {
+        var_meta = {
             name: {"dtype": var.dtype, "attrs": var.attrs, "encoding": var.encoding}
             for name, var in variables.items()
         }
-        index_vars = _create_variables_from_multiindex(
-            index, dim, level_meta=level_meta
-        )
+        index_vars = _create_variables_from_multiindex(index, dim, var_meta=var_meta)
+
+        return obj, index_vars
+
+    @classmethod
+    def from_variables_maybe_expand(
+        cls,
+        dim: Hashable,
+        current_variables: Mapping[Any, "Variable"],
+        variables: Mapping[Any, "Variable"],
+    ) -> Tuple["PandasMultiIndex", IndexVars]:
+        """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.
+        """
+        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})
+
+        def add_level_var(name, var):
+            var_meta[name] = {
+                "dtype": var.dtype,
+                "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
+            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
+            var = next(iter(current_variables.values()))
+            new_var_name = f"{dim}_level_0"
+            names.append(new_var_name)
+            cat = pd.Categorical(var.values, ordered=True)
+            codes.append(cat.codes)
+            levels.append(cat.categories)
+            add_level_var(new_var_name, var)
+
+        for name, var in variables.items():
+            names.append(name)
+            cat = pd.Categorical(var.values, ordered=True)
+            codes.append(cat.codes)
+            levels.append(cat.categories)
+            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
 
     @classmethod
     def from_pandas_index(
         cls, index: pd.MultiIndex, dim: Hashable
     ) -> Tuple["PandasMultiIndex", IndexVars]:
-        level_meta = {}
+        var_meta = {}
         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}"
                 )
-            level_meta[name] = {"dtype": idx.dtype}
+            var_meta[name] = {"dtype": idx.dtype}
 
-        index = index.rename(level_meta.keys())
-        index_vars = _create_variables_from_multiindex(
-            index, dim, level_meta=level_meta
-        )
+        index = index.rename(var_meta.keys())
+        index_vars = _create_variables_from_multiindex(index, dim, var_meta=var_meta)
         return cls(index, dim), index_vars
 
     def query(self, labels, method=None, tolerance=None) -> QueryResult: