Multi-coordinate indexes are dropped when not all associated coordinates survive `_replace_maybe_drop_dims`

What is your issue?

I've been experimenting with a custom index for unstructured topologies, and my impression is that _replace_maybe_drop_dims is overly eager in removing coordinates and indexes. My understanding might be a limited, but I think the desired behavior here should be to consider the indexes more broadly: A custom index spanning multiple dimensions should be retained as long as at least one of its associated coordinates remains, but is currently dropped entirely if any coordinate is filtered out.

I'll try to demonstrate with a minimal example.

Minimal example

Here is a minimal example that triggers the behavior, namely a reduction over an additional dimension (called "three" here):

import numpy as np
import xarray as xr
from xarray import Index, Variable


VERTICES = np.array(
    [
        [0.0, 0.0],  # 0
        [1.0, 0.0],  # 1
        [2.0, 0.0],  # 2
        [0.0, 1.0],  # 3
        [1.0, 1.0],  # 4
        [2.0, 1.0],  # 5
        [1.0, 2.0],  # 6
    ]
)
FACES = np.array(
    [
        [0, 1, 4, 3],
        [1, 2, 5, 4],
        [3, 4, 6, -1],
        [4, 5, 6, -1],
    ]
)

class FaceNodeIndex(Index):
    """
    Minimal example of a multi-dimension index spanning face and node dimensions,
    analogous to a UGRID topology index.
    """

    def __init__(self, vertices, faces):
        self._vertices = vertices  # (n_nodes, 2)
        self._faces = faces        # (n_faces, max_vertices_per_face)
        # Compute face centroids, ignoring fill values (-1)
        masked = np.where(faces[:, :, np.newaxis] == -1, np.nan, vertices[np.where(faces == -1, 0, faces)])
        self._face_x = np.nanmean(masked[:, :, 0], axis=1)
        self._face_y = np.nanmean(masked[:, :, 1], axis=1)

    def should_add_coord_to_array(self, name, var, dims):
        return True

    @classmethod
    def from_variables(cls, variables, options):
        vertices = np.column_stack([v.values for v in variables.values()])
        faces = options["faces"]
        return cls(vertices, faces)

    def create_variables(self, variables=None):
        return {
            "node_x": Variable("nodes", self._vertices[:, 0]),
            "node_y": Variable("nodes", self._vertices[:, 1]),
            "face_x": Variable("faces", self._face_x),
            "face_y": Variable("faces", self._face_y),
        }

# Build the index
index = FaceNodeIndex(VERTICES, FACES)

# Register all four coordinates with the same index instance
coords = xr.Coordinates(
    coords=index.create_variables(),
    indexes={k: index for k in ["node_x", "node_y", "face_x", "face_y"]},
)

ds = xr.Dataset(
    {
        "node_data": xr.Variable(("nodes", "three"), np.random.rand(len(VERTICES), 3)),
        "face_data": xr.Variable(("faces", "three"), np.random.rand(len(FACES), 3)),
    }
).assign_coords(coords)

The resulting .xindexes on the dataset appear as expected:

Indexes:
  ┌ node_x   FaceNodeIndex
  │ node_y
  │ face_x
  └ face_y

And these are preserved on variable selection via getitem:

nodeda = ds["node_data"]
faceda = ds["face_data"]

def has_custom_index(obj):
       return isinstance(next(iter(obj.xindexes.values())), FaceNodeIndex)

assert has_custom_index(nodeda)
assert has_custom_index(faceda)

Now if I reduce over the additional dimension:

reduced_ds = ds.mean("three")
reduced_node = nodeda.mean("three")
reduced_face = faceda.mean("three")

assert has_custom_index(reduced_ds)
assert has_custom_index(reduced_node)

We see that the custom index is preserved in the Dataset (because the dataset contains both face and node dimension), while it is omitted for the DataArrays: for the face data, the index is dropped because it does not contain the node dimension; for the node data it is dropped because it does not contain the face dimension.

Expected behavior

I would expect the index to be preserved, because at least one dimension that is associated with the index persists. (contrast with getitem on the dataset, which does not cause the custom index to be dropped.)

I did a little bit of experimenting, and I could get the desired behavior by modifying two functions:

https://github.com/pydata/xarray/blob/37f2d49b5cfbf5ca7e24dbad6347f99f5a24a368/xarray/core/indexes.py#L2215

https://github.com/pydata/xarray/blob/37f2d49b5cfbf5ca7e24dbad6347f99f5a24a368/xarray/core/dataarray.py#L521

filter_indexes_from_coords is very eager to delete:

    for idx_coord_names in index_coord_names.values():
        if not idx_coord_names <= filtered_coord_names:
            for k in idx_coord_names:
                del filtered_indexes[k]

I.e. if one coordinate isn't present, it is sufficient to delete the index. I think the logic should be changed, if any coordinate is present in the index, the index should be preserved.
Similarly, e.g., the final branch in _replace_maybe_drop_dims checks:

             allowed_dims = set(variable.dims)
            coords = {
                k: v for k, v in self._coords.items() if set(v.dims) <= allowed_dims
            }

Which I think should be something along the lines of:

            coords = {
                k: v for k, v in self._coords.items()
                if set(v.dims) <= allowed_dims or k in indexes
            }

Similarly for the elif branch.

Start with filter_indexes_from_coords in xarray/core/indexes.py and _replace_maybe_drop_dims in xarray/core/dataarray.py, using the minimal FaceNodeIndex example as the reproduction. Trace how multi-coordinate indexes are filtered during reductions over "three". Done means the custom index remains on the reduced DataArrays when at least one associated dimension and coordinate remain, while still being removed when none remain.