Remove children from xarray_beam.Key

We're not going to use this for DataTree, after all.

PiperOrigin-RevId: 820455141

Author: shoyer

xarray_beam/_src/core.py

@@ -29,20 +29,15 @@
 
 
 class Key:
-  """Key for keeping track of chunks of a distributed Dataset or DataTree.
+  """Key for keeping track of chunks of a distributed Dataset.
 
-  Key object in Xarray-Beam include three components:
+  Key object in Xarray-Beam include two components:
 
   - "offsets": an immutable dict indicating integer offsets (total number of
-    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.
+    array elements) from the origin along each dimension for this chunk.
   - "vars": either an frozenset or None, indicating the subset of Dataset
     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
@@ -81,11 +76,6 @@ class Key:
     >>> key.replace(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
@@ -94,32 +84,25 @@ 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,
-      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
-    if children is _DEFAULT:
-      children = self.children
-    return type(self)(offsets, vars, children)
+    return type(self)(offsets, vars)
 
   def with_offsets(self, **offsets: int | None) -> Key:
     """Replace some offsets with new values.
@@ -145,20 +128,17 @@ def __repr__(self) -> str:
       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, self.children))
+    return hash((self.offsets, self.vars))
 
   def __eq__(self, other) -> bool:
     if not isinstance(other, Key):
       return NotImplemented
     return (
         self.offsets == other.offsets
         and self.vars == other.vars
-        and self.children == other.children
     )
 
   def __ne__(self, other) -> bool:
@@ -169,10 +149,7 @@ 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)
-    children_state = (
-        None if self.children is None else sorted(self.children.items())
-    )
-    return offsets_state, vars_state, children_state
+    return offsets_state, vars_state
 
   def __setstate__(self, state):
     self.__init__(*state)

xarray_beam/_src/core_test.py

@@ -37,20 +37,11 @@ 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')
@@ -82,16 +73,6 @@ 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})
 
@@ -125,10 +106,6 @@ def test_repr(self):
     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):
     first = {xbeam.Key({'x': 0, 'y': 10}): 1}
     second = {xbeam.Key({'x': 0, 'y': 10}): 1}
@@ -144,11 +121,6 @@ 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),
@@ -175,22 +147,9 @@ 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})}
+        {'x': 0, 'y': 10}, vars={'foo'}
     )
     unpickled = pickle.loads(pickle.dumps(key))
     self.assertEqual(key, unpickled)