Refactor Dataset internals to store data variables and coordinate variables as separate dicts

[TomNicholas]

There's a lot of other discussion in #9063, but I wanted to pull out this suggestion for independent discussion:

• Our list of internal attributes on a DataTree node is still not just those on a Dataset plus the inherited coordinates

This is indeed a bit of a con, but in my mind the right fix is probably to adjust the Dataset data model to using dictionaries of data_variables and coord_variables, rather than the current solution of a dict of variables and a set of coord_names. Using a separate dictionary for coord_variables would also be more aligned with how DataArray is implemented. The internal Dataset data model is a hold-over from the very early days of Xarray, before we had a notion of coordinate variables that are not indexes.

Originally posted by @shoyer in #9063 (comment)

[TomNicholas]

In fact could the coord_variables dict just be an actual xr.Coordinates object?

(idea from #9204 (comment))

[shoyer]

In fact could the coord_variables dict just be an actual xr.Coordinates object?

I think Coordinates should remain a thin wrapper object. It needs access to both coordinate Variable objects and associated indexes.

[headtr1ck]

I was wondering about that already quite some time, haha.

Does this mean that handling of coordinates can be handled in the parent class DataWithCoords?

[TomNicholas]

Does this mean that handling of coordinates can be handled in the parent class DataWithCoords?

You mean have all of Dataset, DataArray, and DataTree inherit from DataWithCoords, and add things like the .coords property to DataWithCoords? Relates to #9204 (comment).

[benbovy]

The refactor of Dataset internals may (should?) be an opportunity to also refactor coordinates and indexes proxies. The current situation is quite messy and fragile:

• Coordinates is a virtual "stand-alone" container of coordinate variables and their indexes that actually wraps a Dataset object (kind of weird but it was to avoid a bigger refactor)
• DatasetCoordinates and DataArrayCoordinates (returned by Dataset.coords and DataArray.coords respectively) both inherit from Coordinates but are container proxies, which also implement in-place mutation of coordinate variables via access to their underlying Dataset or DataArray.
• Indexes (returned by Dataset.indexes / Dataset.xindexes) is a container proxy of indexes but also of their associated coordinate variables.

A cleaner model might be (from #7368 (comment)):

• Have one true stand-alone Coordinates container that stores data (coordinate variables, indexes and dimension sizes) and that is encapsulated in Dataset and DataArray (instead of having _coord_names, _coords, _indexes, etc. internal dictionaries)
• Remove DatasetCoordinates and DataArrayCoordinates, which are not needed as dataset.coords and dataarray.coords would directly return their encapsulated Coordinates instance
• Make Dataset.xindexes an alias of Dataset.coords.xindexes for convenience
• Get rid of Indexes? Dataset.coords.xindexes could probably just return a frozen mapping instead and specific methods like Indexes.group_by_index() (unfortunately public API) could be progressively moved into Coordinates.

This would represent a much bigger refactor, though. One difficulty with this model is also that updating coordinates cannot be done totally independently of data variables (alignment).

[TomNicholas]

A cleaner model might be (from #7368 (comment)):

• Have one true stand-alone Coordinates container that stores data (coordinate variables, indexes and dimension sizes) and that is encapsulated in Dataset and DataArray (instead of having _coord_names, _coords, _indexes, etc. internal dictionaries)

This sounds great, but bear in mind that that plan was conceived pre-datatree-in-xarray. The DataTree class also has these attributes, but to implement the "coordinate inheritance" feature they often actually use ChainMap to dynamically construct them from coordinates stored on parent tree nodes. Any refactor design should play nicely with this too. See this part of the code:

xarray/xarray/core/datatree.py

Line 564 in 4c2540d

def _indexes(self) -> ChainMap[Hashable, Index]:

[shoyer]

I am very much aligned with this revised data model! Coordinates on DataTree will still need to use a chainmap, but hopefully this is a simple variation of the base Coordinates class. Maybe even the same object, just initialized with the chainmap mappings instead of dicts.

…

On Tue, Mar 11, 2025 at 2:50 PM Tom Nicholas ***@***.***> wrote: A cleaner model might be (from #7368 (comment) <#7368 (comment)>): - Have one true stand-alone Coordinates container that stores data (coordinate variables, indexes and dimension sizes) and that is encapsulated in Dataset and DataArray (instead of having _coord_names, _coords, _indexes, etc. internal dictionaries) This sounds great, but bear in mind that that plan was conceived pre-datatree-in-xarray. The DataTree class also has these attributes, but to implement the "coordinate inheritance" feature they often actually use ChainMap to dynamically construct them from coordinates stored on parent tree nodes. Any refactor design should play nicely with this too. See this part of the code: https://github.com/pydata/xarray/blob/4c2540de54af8d9da4456715d7dcac6b16f330d8/xarray/core/datatree.py#L564 — Reply to this email directly, view it on GitHub <#9203 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVUYP47YXDW6UEVLQXD2T4V6VAVCNFSM6AAAAABYZHMU2CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOMJVGM4TMOBUGM> . You are receiving this because you were mentioned.Message ID: ***@***.***> [image: TomNicholas]*TomNicholas* left a comment (pydata/xarray#9203) <#9203 (comment)> A cleaner model might be (from #7368 (comment) <#7368 (comment)>): - Have one true stand-alone Coordinates container that stores data (coordinate variables, indexes and dimension sizes) and that is encapsulated in Dataset and DataArray (instead of having _coord_names, _coords, _indexes, etc. internal dictionaries) This sounds great, but bear in mind that that plan was conceived pre-datatree-in-xarray. The DataTree class also has these attributes, but to implement the "coordinate inheritance" feature they often actually use ChainMap to dynamically construct them from coordinates stored on parent tree nodes. Any refactor design should play nicely with this too. See this part of the code: https://github.com/pydata/xarray/blob/4c2540de54af8d9da4456715d7dcac6b16f330d8/xarray/core/datatree.py#L564 — Reply to this email directly, view it on GitHub <#9203 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVUYP47YXDW6UEVLQXD2T4V6VAVCNFSM6AAAAABYZHMU2CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOMJVGM4TMOBUGM> . You are receiving this because you were mentioned.Message ID: ***@***.***>