Add .prune() method to remove empty nodes from DataTree after filtering operations

[aladinor]

Is your feature request related to a problem?

Propose adding a .prune() method to DataTree that removes empty nodes (nodes with no data variables or only empty data variables) from the tree structure. This would be particularly useful after filtering operations that can leave some nodes empty.

When performing selection operations on DataTrees, such as sel(), some nodes may end up empty if their coordinate ranges don't overlap with the selection criteria. Currently, these empty nodes remain in the tree structure, which can be confusing and problematic for downstream operations.

Reproducible Example

import xarray as xr
import numpy as np
import pandas as pd

# Create DataTree with different time ranges in each node
ds1 = xr.Dataset(
    {"foo": ("time", np.random.rand(5))},
    coords={"time": pd.date_range("2023-01-01", periods=5, freq="D")}
)
ds2 = xr.Dataset(
    {"var": ("time", np.random.rand(5))},
    coords={"time": pd.date_range("2023-01-04", periods=5, freq="D")}
)

dtree = xr.DataTree.from_dict({"a": ds1, "b": ds2})
print("Original tree:")
print(dtree)

# Filter - node 'b' becomes empty since its time range doesn't overlap
filtered = dtree.sel(time=slice("2023-01-01", "2023-01-03"))
print("\nAfter filtering:")
print(filtered)

Current Output:
Original tree:
<xarray.DataTree>
Group: /
├── Group: /a
│       Dimensions:  (time: 5)
│       Data variables: foo
└── Group: /b
        Dimensions:  (time: 5)
        Data variables: var

After filtering:
<xarray.DataTree>
Group: /
├── Group: /a
│       Dimensions:  (time: 3)
│       Data variables: foo
└── Group: /b          # <- Empty node remains
        Dimensions:  (time: 0)
        Data variables: var (empty)

Describe the solution you'd like

This .prune() method will allow us to remove the empty nodes.

# Proposed usage
pruned = filtered.prune()  # Remove empty nodes
print("\nAfter pruning:")
print(pruned)  # Should only show node 'a'

 Desired Output after .prune():
  <xarray.DataTree>
  Group: /
  └── Group: /a
          Dimensions:  (time: 3)
          Data variables: foo

Describe alternatives you've considered

Proposed API

  def prune(self, recursive=True):
      """
      Remove empty nodes from the DataTree.
      
      Parameters
      ----------
      recursive : bool, default True
          If True, recursively remove empty parent nodes that become 
          empty after pruning their children.
      
      Returns
      -------
      DataTree
          A new DataTree with empty nodes removed.
      """

Use Cases

1. Post-filtering cleanup: Remove empty nodes after sel(), isel(), where() operations
2. Data pipeline optimization: Clean up intermediate results in processing workflows
3. Tree structure simplification: Remove nodes that have become empty due to data processing
4. Memory optimization: Remove unnecessary empty nodes to reduce the memory footprint

Definition of "Empty Node"

A node would be considered empty if:

• It has no data variables, OR
• All its data variables have zero-sized dimensions, OR
• All its data variables contain only NaN/null values (optional behavior)

Implementation Considerations

• Should preserve coordinate information in remaining nodes
• Should handle parent nodes that become empty after children are pruned
• Should work with both named and unnamed nodes
• Should maintain tree structure integrity

Additional context

CC @TomNicholas

[shoyer]

• All its data variables contain only NaN/null values (optional behavior)

This is the only part that feels questionable to me, because checking for NaN values can be quite expensive depending on how large the data tree is. Rather than adding a keyword argument, the alternative here might explicitly use dropna(), e.g., something like tree.dropna().prune().

[TomNicholas]

I agree that we should not check values.

A node would be considered empty if:

Ideally we would make this definition consistent with the current .is_empty property. Then then your .prune implementation can just be

class DataTree:
    def prune(self) -> DataTree:
        nonempty_nodes = {
            path: node.dataset
            for path, node in self.subtree_with_keys
            if not node.is_empty
        }
        return DataTree.from_dict(nonempty_nodes, name=self.name)

Unfortunately the current definition of "empty" used in .is_empty doesn't include the zero-length dimensions part, which you need:

xarray/xarray/core/datatree.py

Line 686 in 0751b72

"""False if node contains any data or attrs. Does not look at children."""

The attrs bit of that definition is also a problem for your use case.

So we either have to make a breaking change to .is_empty or invent a new property with a different definition of emptiness.

EDIT: Or I guess there could be kwargs to .prune...

[TomNicholas]

In fact pruning is arguably just filtering where the filterfunc is is_empty.

xarray/xarray/core/datatree.py

Line 1374 in 0751b72

def filter(self: DataTree, filterfunc: Callable[[DataTree], bool]) -> DataTree:

So the implementation of prune could just call filter. And even without adding the .prune method there is already a pretty short and clear way to do what you want today @aladinor . But a .prune method would be an improvement.

[aladinor]

Thanks @shoyer and @TomNicholas for your feedback! I completely agree with both points:

1. Avoiding NaN checking - You're right about the performance implications. The tree.dropna().prune() pattern is much cleaner.
2. Building on existing infrastructure- @TomNicholas, your suggestion to use .filter() is brilliant and much more maintainable.

I will submit a pr with a suggested .prune implementation