[Experimental] Refactor Dataset to store variables in a manifest

[TomNicholas]

This PR is part of an experiment to see how to integrate a DataTree into xarray.

What is does is refactor Dataset to store variables in a DataManifest class, which is also capable of maintaining a ledger of child tree nodes. The point of this is to prevent name collisions between stored variables and child datatree nodes, as first mentioned in https://github.com/TomNicholas/datatree/issues/38 and explained further in https://github.com/TomNicholas/datatree/issues/2.

("Manifest" in the old sense, of a noun meaning "a document giving comprehensive details of a ship and its cargo and other contents")

• Would eventually close https://github.com/TomNicholas/datatree/issues/38
• Tests added
• Passes pre-commit run --all-files
• User visible changes (including notable bug fixes) are documented in whats-new.rst

[TomNicholas]

What I wanted was for self._variables to always point to self._manifest.variables, such that updating the former by definition updates the latter. But I'm not really sure if that kind of pointer-like behaviour is possible.

[TomNicholas]

@shoyer What do you think about the way I've extended Dataset by adding a ._manifest attribute?

I could have instead only altered ._variables to point to the manifest. but then I end up with confusing code such as

classDataset:
@propertydefvariables(self):
returnself._variables.variables

(I was pleasantly surprised by the very small number of places I needed to change the code to accommodate the manifest - the way that practically everything goes through ._replace() makes it pretty easy to keep all the tests passing.)

[TomNicholas]

The idea is to use this when someone calls dt.ds - so that the DataTree returns an actual Dataset, but one whose contents are linked to the wrapping tree node.

[shoyer]

How about making custom Mapping for use as Dataset._variables directly, which directly is a mapping of dataset variables? You could still be storing the underlying variables in a different way.

[TomNicholas]

It just seems confusing to have an object referred to as `._variables` which actually is just as much a container of child `DataTree` nodes as it is of variables... Also the `DataTree` class needs to wrap this "manifest" too (otherwise I'm storing children in multiple places at once), and so I want it to make sense in that context too.

…

On Tue, 9 Nov 2021, 13:46 Stephan Hoyer, ***@***.***> wrote: How about making custom Mapping for use as Dataset._variables directly, which directly is a mapping of dataset variables? You could still be storing the underlying variables in a different way. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#5961 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AISNPI4LEDDTEGYBWJ7VK6LULFT75ANCNFSM5HVPT5AA> .

[shoyer]

From a xarray.Dataset perspective, Dataset._variables just needs to be a MutableMapping of xarray.Variable objects. And in most cases (when not using DataTree), _variables would still be a plain dictionary, which means adding DataTree support would have no performance implications for normal Dataset objects.

My tentative suggestion would be to use a mixed dictionary with either xarray.Variable or nested dictionaries as entries for the data in DataTree. Then you could make a proxy mapping object that only exposes Variable objects for use in Dataset.

[TomNicholas]

From a xarray.Dataset perspective, Dataset._variables just needs to be a MutableMapping of xarray.Variable objects. And in most cases (when not using DataTree), _variables would still be a plain dictionary, which means adding DataTree support would have no performance implications for normal Dataset objects.

That sounds nice, and might not require any changes to Dataset at all!

My tentative suggestion would be to use a mixed dictionary with either xarray.Variable or nested dictionaries as entries for the data in DataTree.

I think it's a lot easier to have a dict of DataTree objects rather than a nested dict of data, as then each node just points to its child nodes instead of having a node which knows about all the data in the whole tree (if that's what you meant).

How about making custom Mapping for use as Dataset._variables directly, which directly is a mapping of dataset variables?

So this is my understanding of what you're suggesting - I'm just not sure if it solves all the requirements:

classDataManifest(MutableMapping):
""" Acts like a dict of keys to variables, but  prevents setting variables to same key as any  children """def__init__(self, variables={}, children={}):
# check for collisions hereself._variables= {}
self._children= {}
def__getitem__(self, key):
# only expose the variables so this acts like a normal dict of variablesreturnself._variables[key]
def__setitem__(self, key, var):
ifkeyinself._children:
raiseKeyError(
"key already in use to denote a child""node in wrapping DataTree node"
)
self.__dict__[key] =varclassDataset:
self._variables=Mapping[Any, Variable]
# in standard case just use dict of vars as before# Use ._construct_direct as the constructor# as it allows for setting ._variables directly# therefore no changes to Dataset required!classDataTree:
def__init__(self, name, data, parent, children):
self._childrenself._variablesself._coord_namesself._dims
...
@propertydefds(self):
manifest=DataManifest(variables, children)
returnDataset._from_treenode(
variables=manifest,
coord_names=self._coord_names,
dims=self._dims,
...
)
@ds.setterdefds(self, ds):
# check for collisions between ds.data_vars and self.children
...
----------------ds=Dataset({'a': 0})
subtree1=Datatree('group1')
dt=Datatree('root', data=ds, children=[subtree])
wrapped_ds=dt.dswrapped_ds['group1'] =1# raises KeyError - good!subtree2=Datatree('b')
dt.ds['b'] =2# this will happily add a variable to the datasetdt.add_child(subtree2) # want to ensure this raises a KeyError as it conflicts with the new variable, but with this design I'm not sure if it will...

EDIT: Actually maybe this would work? So long as in DataTree we have

classDataTree:
self._variables=manifestself._children=manifest.children

Then adding a new child node would also update the manifest, meaning that the linked dataset should know about it too...

[TomNicholas]

Update: I tried making a custom mapping class (code in drop-down below), then swapping out ._variables = dict(variables) for ._variables = DataManifest(variables=variables) in the Dataset constructors to see if that would break anything, as a first step towards that kind of integration. (At some point it would be good to be able to automatically run the test_dataset.py tests again for the manifest case.)

It kind of works?

From a xarray.Dataset perspective, Dataset._variables just needs to be a MutableMapping of xarray.Variable objects.

It's not quite as simple as this - you need a .copy method (fine), a repr (okay), and there are several places inside Dataset and DataArray that explicitly check that the type of ._variables is a dict.

To get tests to pass I can either relax those type constraints (which leads to >2/3 of test_dataset.py passing immediately) or maybe try making DataManifest inherit from dict so that it passes isinstance(ds._variables, dict)? This probably deserves a new PR...

(EDIT: Though maybe inheriting from dict is more trouble than it's worth)

Code for custom mapping class

fromcollections.abcimportMutableMappingfromtypingimportDict, Hashable, Mapping, Iterator, Sequencefromxarray.core.variableimportVariable#from xarray.tree.datatree import DataTreeclassDataTree:
"""Purely for type hinting purposes for now (and to avoid a circular import)"""
...
classDataManifest(MutableMapping):
""" Stores variables like a dict, but also stores children alongside in a hidden manner, to check against. Acts like a dict of keys to variables, but prevents setting variables to same key as any children. It prevents name collisions by acting as a common record of stored items for both the DataTree instance and its wrapped Dataset instance. """def__init__(
self,
variables: Dict[Hashable, Variable] = {},
children: Dict[Hashable, DataTree] = {},
):
ifvariablesandchildren:
keys_in_both=set(variables.keys()) &set(children.keys())
ifkeys_in_both:
raiseKeyError(
f"The keys {keys_in_both} exist in both the variables and child nodes"
)
self._variables=variablesself._children=children@propertydefchildren(self) ->Dict[Hashable, DataTree]:
"""Stores list of the node's children"""returnself._children@children.setterdefchildren(self, children: Dict[Hashable, DataTree]):
forkey, childinchildren.items():
ifkeyinself.keys():
raiseKeyError("Cannot add child under key {key} because a variable is already stored under that key")
ifnotisinstance(child, DataTree):
raiseTypeErrorself._children=childrendef__getitem__(self, key: Hashable) ->Variable:
"""Forward to the variables here so the manifest acts like a normal dict of variables"""returnself._variables[key]
def__setitem__(self, key: Hashable, value: Variable):
"""Allow adding new variables, but first check if they conflict with children"""ifkeyinself._children:
raiseKeyError(
f"key {key} already in use to denote a child""node in wrapping DataTree node"
)
ifisinstance(value, Variable):
self._variables[key] =valueelse:
raiseTypeError(f"Cannot store object of type {type(value)}")
def__delitem__(self, key: Hashable):
"""Forward to the variables here so the manifest acts like a normal dict of variables"""ifkeyinself._variables:
delself._variables[key]
elifkeyinself.children:
# TODO might be better not to del children here?delself._children[key]
else:
raiseKeyError(f"Cannot remove item because nothing is stored under {key}")
def__contains__(self, item: object) ->bool:
"""Forward to the variables here so the manifest acts like a normal dict of variables"""returniteminself._variablesdef__iter__(self) ->Iterator:
"""Forward to the variables here so the manifest acts like a normal dict of variables"""returniter(self._variables)
def__len__(self) ->int:
"""Forward to the variables here so the manifest acts like a normal dict of variables"""returnlen(self._variables)
defcopy(self) ->"DataManifest":
"""Required for consistency with dict"""returnDataManifest(variables=self._variables.copy(), children=self._children.copy())
# TODO __repr__

[TomNicholas]

Question: Does this change to ._variables need to propagate down to DatasetCoordinates and DataArrayCoordinates? I'm not sure what the intended behaviour is if the user alters .coords directly.

1. It seems it is possible to alter a ds via ds.coords[new_coord_name] = new_coord:

In [34]: ds=xr.Dataset({'a': 0})
In [35]: ds.coords['c'] =2In [36]: dsOut[36]: <xarray.Dataset>Dimensions: ()
Coordinates:
cint642Datavariables:
aint640

(That's a bit weird given that the docstring of DatasetCoordinates describes it as an "immutable dictionary" 😕)

2. It also seems it is possible to similarly alter a da via da.coords[new_coord_name] = new_coord:

In [30]: da=xr.DataArray(0)
In [31]: da.coords['c'] =1In [32]: daOut[32]: <xarray.DataArray ()>array(0)
Coordinates:
cint641

3. However is it meant to be possible to alter a ds via ds[var].coords[new_coord_name] = new_coord? Because that currently silently fails to update:

In [37]: ds=xr.Dataset({'a': 0})
In [38]: ds['a'].coords['c'] =2In [39]: dsOut[39]: <xarray.Dataset>Dimensions: ()
Datavariables:
aint640In [40]: ds['a']
Out[40]: <xarray.DataArray'a' ()>array(0)
In [41]: ds['a'].coordsOut[41]: Coordinates:
*empty*

Bizarrely this does change though:

In [42]: coords=ds['a'].coordsIn [43]: coords['c'] =2In [44]: coordsOut[44]: Coordinates:
cint642

If altering .coords is intended behaviour though that means that my DataManifest also has to be accessible from DataArrayCoordinates too, so that the wrapping DataTree node can know about any changes to dt.ds[var].coords.

[TomNicholas]

Closing for same reasons as in #6086 (comment)