Add children to xarray_beam.Key for DataTree support.

This change extends the `xarray_beam.Key` class to include a `children` component, allowing it to represent chunks within a nested DataTree structure. The `children` component is an immutable dictionary mapping child node names to other `Key` objects.

The structure for `Key` intentionally mirrors the structure of the DataTree structure itself, which stores variables on DataTree objects, with nested DataTree objects for descendents stored in the `children` dict.

xref #124

PiperOrigin-RevId: 811885730

Author: shoyer

xarray_beam/_src/core.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Core data model for xarray-beam."""
+from __future__ import annotations
+
 from collections.abc import Iterator, Mapping, Sequence, Set
 import itertools
 import math
@@ -27,15 +29,20 @@
 
 
 class Key:
-  """A key for keeping track of chunks of a distributed xarray.Dataset.
+  """Key for keeping track of chunks of a distributed Dataset or DataTree.
 
-  Key object in Xarray-Beam include two components:
+  Key object in Xarray-Beam include three components:
 
   - "offsets": an immutable dict indicating integer offsets (total number of
-    array elements) from the origin along each dimension for this chunk.
+    array elements) from the origin along each dimension for this chunk. For
+    DataTree chunks, offsets also apply to all child nodes, similar to how
+    DataTree dimensions are shared with child nodes.
   - "vars": either an frozenset or None, indicating the subset of Dataset
-    variables included in this chunk. None means that all variables are
-    included.
+    variables included in this chunk. The default value of None means that all
+    variables are included.
+  - "children": either an immutabledict of Key objects or None, indicating
+    subset of DataTree node descendents included in this chunk. The default
+    value of None means that all child nodes are included.
 
   Key objects are "deterministically encoded" by Beam, which makes them suitable
   for use as keys in Beam pipelines, i.e., with beam.GroupByKey. They are also
@@ -47,7 +54,7 @@ class Key:
     >>> key = xarray_beam.Key(offsets={'x': 10}, vars={'foo'})
 
     >>> key
-    xarray_beam.Key(offsets={'x': 10}, vars={'foo'})
+    Key(offsets={'x': 10}, vars={'foo'})
 
     >>> key.offsets
     immutabledict({'x': 10})
@@ -58,21 +65,27 @@ class Key:
   To replace some offsets::
 
     >>> key.with_offsets(y=0)  # insert
-    xarray_beam.Key(offsets={'x': 10, 'y': 0}, vars={'foo'})
+    Key(offsets={'x': 10, 'y': 0}, vars={'foo'})
 
     >>> key.with_offsets(x=20)  # override
-    xarray_beam.Key(offsets={'x': 20}, vars={'foo'})
+    Key(offsets={'x': 20}, vars={'foo'})
 
     >>> key.with_offsets(x=None)  # remove
-    xarray_beam.Key(offsets={}, vars={'foo'})
+    Key(offsets={}, vars={'foo'})
 
   To entirely replace offsets or variables::
 
     >>> key.replace(offsets={'y': 0})
-    xarray_beam.Key(offsets={'y': 0}, vars={'foo'})
+    Key(offsets={'y': 0}, vars={'foo'})
 
     >>> key.replace(vars=None)
-    xarray_beam.Key(offsets={'x': 10}, vars=None)
+    Key(offsets={'x': 10})
+
+  Children are defined using nested `Key` objects::
+
+    >>> xarray_beam.Key(children={'first_child': xarray_beam.Key({'x': 10})})
+    Key(children={'first_child': Key(offsets={'x': 10})})
+
   """
 
   # pylint: disable=redefined-builtin
@@ -81,26 +94,43 @@ def __init__(
       self,
       offsets: Mapping[str, int] | None = None,
       vars: Set[str] | None = None,
+      children: Mapping[str, Key] | None = None,
   ):
     if offsets is None:
       offsets = {}
     if isinstance(vars, str):
       raise TypeError(f"vars must be a set or None, but is {vars!r}")
     self.offsets = immutabledict.immutabledict(offsets)
     self.vars = None if vars is None else frozenset(vars)
+    self.children = (
+        None if children is None else immutabledict.immutabledict(children)
+    )
 
   def replace(
       self,
       offsets: Mapping[str, int] | object = _DEFAULT,
       vars: Set[str] | None | object = _DEFAULT,
-  ) -> "Key":
+      children: Mapping[str, Key] | None | object = _DEFAULT,
+  ) -> Key:
+    """Replace one or more components of this Key with new values."""
     if offsets is _DEFAULT:
       offsets = self.offsets
     if vars is _DEFAULT:
       vars = self.vars
-    return type(self)(offsets, vars)
+    if children is _DEFAULT:
+      children = self.children
+    return type(self)(offsets, vars, children)
+
+  def with_offsets(self, **offsets: int | None) -> Key:
+    """Replace some offsets with new values.
 
-  def with_offsets(self, **offsets: int | None) -> "Key":
+    Args:
+      **offsets: offsets to override (for integer values) or remove, with
+        values of ``None``.
+
+    Returns:
+      New Key with the specified offsets.
+    """
     new_offsets = dict(self.offsets)
     for k, v in offsets.items():
       if v is None:
@@ -110,17 +140,26 @@ def with_offsets(self, **offsets: int | None) -> "Key":
     return self.replace(offsets=new_offsets)
 
   def __repr__(self) -> str:
-    offsets = dict(self.offsets)
-    vars = set(self.vars) if self.vars is not None else None
-    return f"{type(self).__name__}(offsets={offsets}, vars={vars})"
+    components = []
+    if self.offsets:
+      components.append(f"offsets={dict(self.offsets)}")
+    if self.vars is not None:
+      components.append(f"vars={set(self.vars)}")
+    if self.children is not None:
+      components.append(f"children={dict(self.children)}")
+    return f"{type(self).__name__}({', '.join(components)})"
 
   def __hash__(self) -> int:
-    return hash((self.offsets, self.vars))
+    return hash((self.offsets, self.vars, self.children))
 
   def __eq__(self, other) -> bool:
     if not isinstance(other, Key):
       return NotImplemented
-    return self.offsets == other.offsets and self.vars == other.vars
+    return (
+        self.offsets == other.offsets
+        and self.vars == other.vars
+        and self.children == other.children
+    )
 
   def __ne__(self, other) -> bool:
     return not self == other
@@ -130,7 +169,10 @@ def __ne__(self, other) -> bool:
   def __getstate__(self):
     offsets_state = sorted(self.offsets.items())
     vars_state = None if self.vars is None else sorted(self.vars)
-    return (offsets_state, vars_state)
+    children_state = (
+        None if self.children is None else sorted(self.children.items())
+    )
+    return offsets_state, vars_state, children_state
 
   def __setstate__(self, state):
     self.__init__(*state)

xarray_beam/_src/core_test.py

@@ -18,6 +18,7 @@
 from absl.testing import parameterized
 import apache_beam as beam
 import immutabledict
+import pickle
 import numpy as np
 import xarray
 import xarray_beam as xbeam
@@ -36,11 +37,20 @@ def test_constructor(self):
     self.assertIsInstance(key.offsets, immutabledict.immutabledict)
     self.assertEqual(dict(key.offsets), {'x': 0, 'y': 10})
     self.assertIsNone(key.vars)
+    self.assertIsNone(key.children)
 
     key = xbeam.Key(vars={'foo'})
     self.assertEqual(dict(key.offsets), {})
     self.assertIsInstance(key.vars, frozenset)
     self.assertEqual(set(key.vars), {'foo'})
+    self.assertIsNone(key.children)
+
+    child_key = xbeam.Key({'x': 0})
+    key = xbeam.Key(children={'sub': child_key})
+    self.assertEqual(dict(key.offsets), {})
+    self.assertIsNone(key.vars)
+    self.assertIsInstance(key.children, immutabledict.immutabledict)
+    self.assertEqual(dict(key.children), {'sub': child_key})
 
     with self.assertRaisesRegex(TypeError, 'vars must be a set or None'):
       xbeam.Key(vars='foo')
@@ -72,6 +82,16 @@ def test_replace(self):
     actual = key.replace({'y': 1}, {'bar'})
     self.assertEqual(expected, actual)
 
+    child = xbeam.Key()
+    expected = xbeam.Key({'x': 0}, {'foo'}, children={'sub': child})
+    actual = key.replace(children={'sub': child})
+    self.assertEqual(expected, actual)
+
+    key2 = xbeam.Key(children={'sub': child})
+    expected = xbeam.Key()
+    actual = key2.replace(children=None)
+    self.assertEqual(expected, actual)
+
   def test_with_offsets(self):
     key = xbeam.Key({'x': 0})
 
@@ -98,11 +118,15 @@ def test_with_offsets(self):
 
   def test_repr(self):
     key = xbeam.Key({'x': 0, 'y': 10})
-    expected = "Key(offsets={'x': 0, 'y': 10}, vars=None)"
+    expected = "Key(offsets={'x': 0, 'y': 10})"
     self.assertEqual(repr(key), expected)
 
     key = xbeam.Key(vars={'foo'})
-    expected = "Key(offsets={}, vars={'foo'})"
+    expected = "Key(vars={'foo'})"
+    self.assertEqual(repr(key), expected)
+
+    key = xbeam.Key(children={'sub': xbeam.Key()})
+    expected = "Key(children={'sub': Key()})"
     self.assertEqual(repr(key), expected)
 
   def test_dict_key(self):
@@ -120,6 +144,11 @@ def test_equality(self):
     self.assertNotEqual(key, key2)
     self.assertNotEqual(key2, key)
 
+    key3 = xbeam.Key({'x': 0, 'y': 10}, children={'sub': xbeam.Key()})
+    self.assertEqual(key3, key3)
+    self.assertNotEqual(key, key3)
+    self.assertNotEqual(key3, key)
+
   def test_offsets_as_beam_key(self):
     inputs = [
         (xbeam.Key({'x': 0, 'y': 1}), 1),
@@ -146,6 +175,26 @@ def test_vars_as_beam_key(self):
     actual = inputs | beam.GroupByKey()
     self.assertEqual(actual, expected)
 
+  def test_children_as_beam_key(self):
+    inputs = [
+        (xbeam.Key(children={'sub': xbeam.Key()}), 1),
+        (xbeam.Key(children={}), 2),
+        (xbeam.Key(children={'sub': xbeam.Key()}), 3),
+    ]
+    expected = [
+        (xbeam.Key(children={'sub': xbeam.Key()}), [1, 3]),
+        (xbeam.Key(children={}), [2]),
+    ]
+    actual = inputs | beam.GroupByKey()
+    self.assertEqual(actual, expected)
+
+  def test_pickle(self):
+    key = xbeam.Key(
+        {'x': 0, 'y': 10}, vars={'foo'}, children={'sub': xbeam.Key({'z': 0})}
+    )
+    unpickled = pickle.loads(pickle.dumps(key))
+    self.assertEqual(key, unpickled)
+
 
 class TestOffsetsToSlices(test_util.TestCase):
 
@@ -517,9 +566,9 @@ def test_validate_chunk_raises_on_dask_chunked(self):
     with self.assertRaisesRegex(
         ValueError,
         re.escape(
-            "Dataset variable 'foo' corresponding to key Key(offsets={'x': 0},"
-            ' vars=None) is chunked with Dask. Datasets passed to'
-            ' validate_chunk must be fully computed (not chunked):'
+            "Dataset variable 'foo' corresponding to key Key(offsets={'x': 0})"
+            ' is chunked with Dask. Datasets passed to validate_chunk must be'
+            ' fully computed (not chunked):'
         ),
     ):
       core.validate_chunk(key, dataset)
@@ -530,7 +579,7 @@ def test_unmatched_dimension_raises_error(self):
     with self.assertRaisesRegex(
         ValueError,
         re.escape(
-            "Key offset(s) 'y' in Key(offsets={'x': 0, 'y': 0}, vars=None) not "
+            "Key offset(s) 'y' in Key(offsets={'x': 0, 'y': 0}) not "
             'found in Dataset dimensions'
         ),
     ):