wip: deeper refactoring label-based sel

Created QueryResult and MergedQueryResults classes for convenience.

Author: benbovy

xarray/core/alignment.py

@@ -69,8 +69,8 @@ def _override_indexes(objects, all_indexes, exclude):
             dim: all_indexes[dim][0] for dim in obj.xindexes if dim not in exclude
         }
 
-        # TODO: benbovy - explicit indexes: not refactored yet (dirty fix)
-        objects[idx + 1] = obj._overwrite_indexes(new_indexes, {}, [])
+        # TODO: benbovy - explicit indexes: not refactored yet!
+        objects[idx + 1] = obj._overwrite_indexes(new_indexes)
 
     return objects
 

xarray/core/coordinates.py

@@ -398,9 +398,7 @@ def remap_label_indexers(
     method: str = None,
     tolerance=None,
     **indexers_kwargs: Any,
-) -> Tuple[
-    dict, dict, dict, list
-]:  # TODO more precise return type after annotations in indexing
+) -> Any:
     """Remap indexers from obj.coords.
     If indexer is an instance of DataArray and it has coordinate, then this coordinate
     will be attached to pos_indexers.
@@ -421,23 +419,21 @@ def remap_label_indexers(
         for k, v in indexers.items()
     }
 
-    (
-        pos_indexers,
-        new_indexes,
-        new_variables,
-        drop_variables,
-    ) = indexing.remap_label_indexers(
+    query_results = indexing.remap_label_indexers(
         obj, v_indexers, method=method, tolerance=tolerance
     )
 
     # attach indexer's coordinate to pos_indexers
     for k, v in indexers.items():
+        dim_indexer = query_results.dim_indexers.get(k, None)
         if isinstance(v, Variable):
-            pos_indexers[k] = Variable(v.dims, pos_indexers[k])
+            query_results.dim_indexers[k] = Variable(v.dims, dim_indexer)
         elif isinstance(v, DataArray):
             # drop coordinates found in indexers since .sel() already
             # ensures alignments
             coords = {k: var for k, var in v._coords.items() if k not in indexers}
-            pos_indexers[k] = DataArray(pos_indexers[k], coords=coords, dims=v.dims)
+            query_results.dim_indexers[k] = DataArray(
+                dim_indexer, coords=coords, dims=v.dims
+            )
 
-    return pos_indexers, new_indexes, new_variables, drop_variables
+    return query_results

xarray/core/dataarray.py

@@ -463,37 +463,33 @@ def _replace_maybe_drop_dims(
     def _overwrite_indexes(
         self,
         indexes: Mapping[Hashable, Index],
-        coords: Mapping[Hashable, Variable],
-        drop_coords: List[Hashable],
+        coords: Optional[Mapping[Hashable, Variable]] = None,
+        drop_coords: Optional[List[Hashable]] = None,
+        rename_dims: Optional[Mapping[Hashable, Hashable]] = None,
     ) -> "DataArray":
         """Maybe replace indexes and their corresponding coordinates."""
         if not indexes:
             return self
 
-        assert indexes.keys() == coords.keys()
+        if coords is None:
+            coords = {}
+        if drop_coords is None:
+            drop_coords = []
 
         new_variable = self.variable.copy()
         new_coords = self._coords.copy()
         new_indexes = dict(self.xindexes)
-        dims_dict = {}
 
         for name in indexes:
-            # new coordinate variables may have renamed dimensions (e.g., level
-            # name of a multi-index converted to a single index)
-            old_vs_new_dims = zip(self._coords[name].dims, coords[name].dims)
-            for old_dim, new_dim in old_vs_new_dims:
-                if old_dim != new_dim:
-                    dims_dict[old_dim] = new_dim
-
             new_coords[name] = coords[name]
             new_indexes[name] = indexes[name]
 
         for name in drop_coords:
             new_coords.pop(name)
             new_indexes.pop(name)
 
-        if dims_dict:
-            new_variable.dims = [dims_dict.get(d, d) for d in new_variable.dims]
+        if rename_dims:
+            new_variable.dims = [rename_dims.get(d, d) for d in new_variable.dims]
 
         return self._replace(
             variable=new_variable, coords=new_coords, indexes=new_indexes

xarray/core/dataset.py

@@ -1166,30 +1166,24 @@ def _replace_vars_and_dims(
     def _overwrite_indexes(
         self,
         indexes: Mapping[Hashable, Index],
-        variables: Mapping[Hashable, Variable],
-        drop_variables: List[Hashable],
+        variables: Optional[Mapping[Hashable, Variable]] = None,
+        drop_variables: Optional[List[Hashable]] = None,
+        rename_dims: Optional[Mapping[Hashable, Hashable]] = None,
     ) -> "Dataset":
         """Maybe replace indexes and their corresponding index variables."""
         if not indexes:
             return self
 
-        assert indexes.keys() == variables.keys()
+        if variables is None:
+            variables = {}
+        if drop_variables is None:
+            drop_variables = []
 
         new_variables = self._variables.copy()
         new_coord_names = self._coord_names.copy()
         new_indexes = dict(self.xindexes)
-        dims_dict = {}
 
         for name in indexes:
-            # new coordinate variables may have renamed dimensions (e.g., level
-            # name of a multi-index converted to a single index)
-            # TODO: instead of infer renamed dimensions from the coordinates,
-            # should we require explicitly providing it from Index.query?
-            old_vs_new_dims = zip(self._variables[name].dims, variables[name].dims)
-            for old_dim, new_dim in old_vs_new_dims:
-                if old_dim != new_dim:
-                    dims_dict[old_dim] = new_dim
-
             new_variables[name] = variables[name]
             new_indexes[name] = indexes[name]
 
@@ -1202,10 +1196,10 @@ def _overwrite_indexes(
             variables=new_variables, coord_names=new_coord_names, indexes=new_indexes
         )
 
-        if dims_dict:
+        if rename_dims:
             # skip rename indexes: they should already have the right name(s)
-            dims = replaced._rename_dims(dims_dict)
-            new_variables, new_coord_names = replaced._rename_vars({}, dims_dict)
+            dims = replaced._rename_dims(rename_dims)
+            new_variables, new_coord_names = replaced._rename_vars({}, rename_dims)
             return replaced._replace(
                 variables=new_variables, coord_names=new_coord_names, dims=dims
             )
@@ -2480,12 +2474,12 @@ def sel(
         DataArray.sel
         """
         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel")
-        pos_indexers, new_indexes, new_variables, drop_variables = remap_label_indexers(
+        query_results = remap_label_indexers(
             self, indexers=indexers, method=method, tolerance=tolerance
         )
 
-        result = self.isel(indexers=pos_indexers, drop=drop)
-        return result._overwrite_indexes(new_indexes, new_variables, drop_variables)
+        result = self.isel(indexers=query_results.dim_indexers, drop=drop)
+        return result._overwrite_indexes(*query_results.to_tuple()[1:])
 
     def head(
         self,

xarray/core/indexes.py

@@ -16,7 +16,7 @@
 import pandas as pd
 
 from . import formatting, utils
-from .indexing import PandasIndexingAdapter, PandasMultiIndexingAdapter
+from .indexing import PandasIndexingAdapter, PandasMultiIndexingAdapter, QueryResult
 from .utils import is_dict_like, is_scalar
 
 if TYPE_CHECKING:
@@ -47,7 +47,7 @@ def to_pandas_index(self) -> pd.Index:
 
     def query(
         self, labels: Dict[Hashable, Any], **kwargs
-    ) -> Tuple[Mapping[str, Any], Optional[IndexWithVars]]:  # pragma: no cover
+    ) -> QueryResult:  # pragma: no cover
         raise NotImplementedError()
 
     def equals(self, other):  # pragma: no cover
@@ -243,7 +243,7 @@ def query(self, labels, method=None, tolerance=None):
                 if np.any(indexer < 0):
                     raise KeyError(f"not all values found in index {coord_name!r}")
 
-        return {self.dim: indexer}, None
+        return QueryResult({self.dim: indexer})
 
     def equals(self, other):
         return self.index.equals(other.index)
@@ -425,13 +425,27 @@ def query(self, labels, method=None, tolerance=None):
                 new_index, new_vars = PandasMultiIndex.from_pandas_index(
                     new_index, self.dim
                 )
+                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
                 )
-            return {self.dim: indexer}, (new_index, new_vars)
+                dims_dict = {self.dim: new_index.index.name}
+                drop_coords = set(self.index.names) - {new_index.index.name} | {
+                    self.dim
+                }
+
+            return QueryResult(
+                {self.dim: indexer},
+                index=new_index,
+                index_vars=new_vars,
+                drop_coords=list(drop_coords),
+                rename_dims=dims_dict,
+            )
+
         else:
-            return {self.dim: indexer}, None
+            return QueryResult({self.dim: indexer})
 
 
 def remove_unused_levels_categories(index: pd.Index) -> pd.Index: