Store variables in DataTree instead of storing Dataset

[TomNicholas]

Should we prefer inheritance or composition when making the node of a datatree behave like an xarray Dataset?

Inheritance

We really want the data-containing nodes of the datatree to behave as much like xarray datasets as possible, as we will likely be calling functions/methods on them, assigning them, extracting from them and saving them as if they were actually xarray.Dataset objects. We could imagine a tree node class which directly inherits from xarray.Dataset:

class DatasetNode(xarray.Dataset, NodeMixin):
    ...

This would have all the attributes and API of a Dataset, and pass isinstance() checks, but also the attributes and methods needed to function as a node in a tree (e.g. .children, .parent). We would still need to decorate most inherited methods in order to apply them to all child nodes in the tree though.

Mostly these don't collide, except in the important case of getting/setting children of a node. xarray.Datasets already use up __getitem__ for variable selection (i.e. ds[var]) as well as the .some_variable namespace via property-like access. This means we can't immediately have an API allowing operations like dt.weather = dt.weather.mean('time') because .weather is a child of the node, not a dataset variable. (It's possible we could have both behaviours simultaneously by overwriting __getitem__, but then we might restrict the possible names of children/variables.)

I think this approach would also have the side-effect that accessor methods registered with @register_dataset_accessor would also be callable on the tree nodes.

Composition

The alternative is instead of each node being a Dataset, each node merely wraps a Dataset. This has the advantage of keeping the Data class and the Node class separate, though they would still share a large API to allow applying a method (e.g. .mean() to all child nodes in a tree.

The disadvantage is that then all the variables and dataset attributes are behind a .ds property.

This type of syntax dt.weather = dt.weather.mean('time') would then be possible (at least if we didn't allow the tree objects to have their own .attrs, else it would have to be dt['weather'] = dt['weather'].mean('time')) because we would be calling the method of a DatasetNode (rather than Dataset) and then assigning to a DatasetNode.

Selecting a particular variable from a dataset stored at a particular node would then look like dt['weather'].ds['pressure'], which has the advantage of clarifying which one is the variable, but the disadvantage of breaking up the path-like structure to get from the root down to the variable. EDIT: As there is no problem with collisions between names of groups and variables, we can actually just override __getitem__ to check in both the data variables and the children, so we can have access like dt['weather']['pressure'].

---

(There is also a possible third option described in #4)

---

For now the second approach seemed better, but I'm looking for other opinions!

[TomNicholas]

I think this really comes down to two questions:

1. Do we want isinstance(DataTree, Dataset) to return True?

2. Do we want to automatically inherit all Dataset methods, or only wrap the methods we have specifically chosen?

Pros of automatic inheritance:

• We would not have to write out the entire massive API of xarray.Dataset again within the definition of DataTree.
• Methods which rely on other Dataset methods to work should then automatically work, for example lots of methods internally rely on Dataset._construct_direct() to rebuild the result before returning. These don't need to be wrapped with map_over_subtree, so inheriting them undecorated is what we want.
• We also get all property methods defined immediately.

Pros of manual wrapping:

• There won't be any unexpected behaviour resulting from inheriting a method we hadn't thought about. Any method we didn't explicitly consider would not be defined on DataTree.
• We still have to decorate most of the Dataset methods with @map_over_subtree, and this is easier to control without mistakes if they are wrapped rather than inherited.
• We have a separate namespace for Dataset methods vs DataTree methods. This is helpful if we want to make identically-named methods on both classes easily available to the user. For example DataTree.merge(another_dt) and Dataset.merge(another_ds) clearly need different implementations, but would collide if both defined on DataTree. Calling the latter via DataTree.ds.merge(another_ds) would allow both to be accessible from the DataTree interface.

[TomNicholas]

Another, probably much better, approach to inheritance would be to make benign changes to xarray.Dataset's inheritance pattern upstream that make it easier to pick and choose the correct methods to inherit on DataTree. For example, we currently have to awkwardly create our own DatasetPropertiesMixin by wrapping properties in Dataset, but we could instead factor out those property definitions in xarray into a xarray.core.Dataset.DatasetPropertiesMixin class, then inherit something like this:

xarray.core.dataset.py:

class DatasetPropertiesMixin:
    @property
    def dims(self):
        ...  # etc.

class DatasetMethodsMixin:
    def isel(self, indexers, ...):
        ...  # etc.

class Dataset(DatasetPropertiesMixin, DatasetMethodsMixin, DataWithCoords, DatasetArithmetic):
    ...

datatree.datatree.py:

from xarray.core.dataset import DatasetPropertiesMixin, DatasetMethodsMixin

DataTreeMethodsMixin = map_all_methods_over_subtree(DatasetMethodsMixin)
DataTreeArithmetic = map_all_methods_over_subtree(DatasetArithmetic)

class Dataset(DatasetPropertiesMixin, DataTreeMethodsMixin, DataWithCoords, DataTreeArithmetic):
    ...

(where that example also has another mixin class for mapped Dataset methods too.)

This would not make isinstance(DataTree, Dataset) return True, but it would probably be the shortest and least error-prone way to get all the right methods defined on all the right classes.

[TomNicholas]

@shoyer or @max-sixty, I would appreciate either of your opinions on the above fundamental design question for DataTree, because this is practically the most core question, and at the moment I'm just talking to myself 😅

[TomNicholas]

This also came up in #26, in the context of whether an empty netCDF group is better represented by an empty Dataset or None (a question also related to #4.)

[max-sixty]

this is practically the most core question, and at the moment I'm just talking to myself

eventually software ends up at philosophy...

---

I do think it's an interesting question. I'm really not sure where we'll end up.

Without having thought as deeply as you, I would probably start with composing them. Basically for the reasons you describe — inheritance a) couples the classes together and b) stepping back, it relies on a known ontology — so mixins like DatasetMethodsMixin & DatasetPropertiesMixin might seem like a good structure now — but then for another case we're going to need a slightly different structure, and the number of Mixins propagate until they're no longer useful groupings.

Having everything behind a .ds property does seem like a worse API, but it might be a reasonable way of getting some cool examples going (I already see some good ones!). Or to reduce the amount of boilerplate to get started, you could intercept methods with __getattribute__ and decide where to route them?

And this doesn't preclude a Dataset becoming a DataTree in the future — it's much easier to go composition -> inheritance (or even integration) than the other way around.

To confirm, I haven't thought about this sufficiently, so please weigh my view appropriately!

[shoyer]

I would also start with composition. It's much more verbose, but otherwise you'll never be quite sure when DataTree methods will work properly. It should not be too painful to use a DataTree.dataset property for any missing methods.

Making issubclass(DataTree, Dataset) true is also problematic, unless a DataTree is in fact 100% substitutable for a Dataset. Even if this can be made true (which I doubt), most of us have a hard time reasoning about inheritance.

Reusing Xarray's mixins would be reasonable, but only if/when DataTree lives inside the Xarray codebase. Otherwise it would be too easy to introduce errors.

[TomNicholas]

Thanks both for your input.

I did in fact start with composition, so I'll keep it that way for now.

Reusing Xarray's mixins would be reasonable, but only if/when DataTree lives inside the Xarray codebase. Otherwise it would be too easy to introduce errors.

I'm not sure if this is quite what you meant, but I think what I'll do for now is define my own mixins internally, but make sure none of them actually inherit from any of xarray's mixins. That way I can keep the internal organisational difference between "methods on Dataset that need wrapping", "methods on Dataset that need wrapping and mapping over child nodes", "properties on Dataset that need wrapping" etc., but without actually depending on any private xarray API.

For example, this mixin for Datatree to inherit from does not itself inherit from (or even reference) any of xarray's mixins:

class MappedDatasetMethodsMixin:
    """
    Mixin to add Dataset methods like .mean(), but wrapped to map over all nodes in the subtree.
    """

    __slots__ = ()
    _wrap_then_attach_to_cls(
        cls_dict=vars(), 
        copy_methods_from=Dataset, 
        methods_to_copy=_ALL_DATASET_METHODS_TO_MAP, 
        wrap_func=map_over_subtree,
    )


class DataTree(
    TreeNode,
    DatasetPropertiesMixin,
    MappedDatasetMethodsMixin,
    MappedDataWithCoordsMixin,
    DataTreeArithmeticMixin,
):

[jhamman]

It seems that the current discussion may be hedging toward inheritance over composition to address issues like name collisions (#38) and mutability. However, it may be worth considering the possibilities of a third option that doesn't use the Dataset as a container at the node level. The concept would basically define each node as either a DataArray or a TreeNode (subtree). TreeNodes would be responsible for insuring alignment, coordinating ops among children, and would be directly convertible to Datasets with 0 or more variables.

This would solve the name collisions problem but would require much more logic here in datatree. Its possible this would be much harder than a direct inheritance implementation but I thought I would through the idea out there.

[shoyer]

The concept would basically define each node as either a DataArray or a TreeNode (subtree). TreeNodes would be responsible for insuring alignment, coordinating ops among children, and would be directly convertible to Datasets with 0 or more variables.

I like this idea a lot, but let me suggest a slight variation: store Variable rather than DataArray objects in the dict.

The core data model for xarray.Dataset includes:

1. attrs: dict of attributes
2. variables: dict of xarray.Variable
3. coord_names: set of coordinate names
4. dims: dict of dimension names to sizes
5. indexes: dict from dimension name to xarray.Index objects (currently optional and can be left as None, but will be required after the explicit indexes refactor).

My suggestion would be to make DataTree quite similar, except that variables can also hold DataTree objects.

To implement methods on DataTree, you would filter DataTree objects out of variables, and then convert the remainder directly into an xarray.Dataset.