Fix multi-coordinate indexes dropped in _replace_maybe_drop_dims

[rsignell]

@benbovy, after some discussion with @Huite, he thought Claude Code might be capable of fixing this issue, so I gave it a shot. I'm only an appreciative user of the xarray codebase, so I don't know really how to test/evaluate this beyond the passing tests that Claude itself came up with. So I hope this isn't a waste of your time. With that said, here's what Claude came up with:

When a custom Index spans multiple coordinates across different dimensions (e.g. a UGRID topology index covering both nodes and faces), it was incorrectly dropped during:

• DataArray dimension reduction (e.g. .mean("extra_dim")) via _replace_maybe_drop_dims, which filtered coords by strict dim subset without consulting Index.should_add_coord_to_array()
• Dataset subsetting by variable names (e.g. ds[["node_data"]]) via _copy_listed, with the same issue

Both paths now use should_add_coord_to_array() for index-backed coordinates, consistent with how _construct_dataarray already works.

Description

Checklist

• Closes 11215
• Tests added

AI Disclosure

• This PR contains AI-generated content.
• I have tested the AI-generated content in my PR (but with Claude-generated tests)
• I take responsibility for any AI-generated content in my PR.

Tools: Claude Code

[rsignell]

@benboy, one more thing. I did ask Claude how if it used your suggestion in #11215 (comment) and it said:

[benbovy]

Thanks @rsignell! That looks good to me.

For the case of Dataset.to_dataarray(), I think it should work similarly than in DataArray._replace_maybe_drop_dims():

coords = {}
for k, v in self.coords.items():
    if k in self._indexes:
        if self._indexes[k].should_add_coord_to_array(k, v, dims):
            coords[k] = v
    elif set(v.dims) <= dims:
        coords[k] = v

[benbovy]

Suggested change

	assert isinstance(next(iter(reduced_node.xindexes.values())), MultiDimIndex), (
	"Index dropped from node DataArray after reduction"
	)
	assert isinstance(next(iter(reduced_face.xindexes.values())), MultiDimIndex), (
	"Index dropped from face DataArray after reduction"
	)
	for da in [reduced_node, reduced_face]:
	for name in ["node_x", "node_y", "face_x", "face_y"]:
	assert name in da.coords
	assert isinstance(ds.xindexes[name], MultiDimIndex)

The instance(next(iter(...), ...) checks should be enough but this suggestion is to make the expected behavior clearer to the reader (or LLM).

[benbovy]

https://github.com/pydata/xarray/pull/11286/changes#r3056276588

[rsignell]

Thanks for the review @benbovy!

Here is the Claude Code plan to address your comments -- let me know if you would like changes!

Planned changes

1. Fix Dataset.to_dataarray() (xarray/core/dataset.py ~line 7163)

The current implementation includes all coords unconditionally without calling should_add_coord_to_array, so multi-coordinate indexes can be dropped via filter_indexes_from_coords. Will fix to mirror the pattern in _copy_listed / _construct_dataarray:

needed_dims = set(broadcast_vars[0].dims)
coords: dict[Hashable, Variable] = {}
for k, v in self.coords.items():
    if k in self._indexes:
        if self._indexes[k].should_add_coord_to_array(k, v.variable, needed_dims):
            coords[k] = v.variable
    elif set(v.dims) <= needed_dims:
        coords[k] = v.variable
indexes = filter_indexes_from_coords(self._indexes, set(coords))

2. Improve test assertions in test_dataarray.py (lines 586–591)

Replace next(iter(...)) with the explicit loop you suggested (using da.xindexes[name]):

for da in [reduced_node, reduced_face]:
    for name in ["node_x", "node_y", "face_x", "face_y"]:
        assert name in da.coords
        assert isinstance(da.xindexes[name], MultiDimIndex)

3. Improve test assertions in test_dataset.py (lines 4585–4590)

Same pattern applied to the dataset test:

for subset in [node_subset, face_subset]:
    for name in ["node_x", "node_y", "face_x", "face_y"]:
        assert name in subset.coords
        assert isinstance(subset.xindexes[name], MultiDimIndex)

4. Add a new test for to_dataarray() with multi-coordinate indexes

Add test_to_dataarray_preserves_multi_coord_index in test_dataset.py using the same MultiDimIndex fixture pattern.

---

[This is Claude Code on behalf of Rich Signell]

[benbovy]

At a first glance the the proposed changes look good!

[rsignell]

@benbovy, okay, Claude thinks it addressed the comments. Fingers crossed! 😄

[benbovy]

Actually, a fix is not even necessary here since Dataset variables are broadcasted against each other and thus all coordinates are kept in the returned DataArray anyway, like mentioned in the docstrings. I added to_dataarray() to the list as I naively looked at all places where filter_indexes_from_coords is called, sorry.

(filter_indexes_from_coords is not even needed, copying self._indexes should be enough).

[rsignell]

@benbovy, I asked Claude to double-check if the issue was addressed:

[benbovy]

LGTM too, thanks!

[welcome]

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again!