Add indices to xarray_beam.Key

`indices` is an alternative data model to `offsets` that should work better for `DataTree` objects with different chunk sizes, For example, we could use indices (but not offsets) to support sub-groups with different resolutions along a "time" dimensions (e.g., hourly vs 6-hourly vs daily).

PiperOrigin-RevId: 822731332

Author: shoyer

xarray_beam/_src/core.py

@@ -58,14 +58,25 @@ def inc_timer_msec(namespace: str | type[Any], name: str) -> Iterator[None]:
 class Key:
   """Key for keeping track of chunks of a distributed Dataset.
 
-  Key object in Xarray-Beam include two components:
+  Key objects in Xarray-Beam include two components:
 
-  - "offsets": an immutable dict indicating integer offsets (total number of
+  - `offsets`: an immutable dict indicating integer offsets (total number of
     array elements) from the origin along each dimension for this chunk.
-  - "vars": either an frozenset or None, indicating the subset of Dataset
+  - `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.
 
+  Alternatively, `indices` may be specified instead of `offsets`. This is a
+  newer data model that is not yet fully supported:
+
+  - `indices`: an immutable dict indicating integer chunk indices from the
+    origin along each dimension for this chunk.
+
+  `offsets` and `indices` are mutually exclusive: only one of them may be used
+  for any given `Key`. For example, if there are chunks of size 100 along the
+  'x' dimension, then ``offsets={'x': 400}`` would correspond to
+  ``indices={'x': 4}``.
+
   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
   immutable and hashable, which makes them usable as keys in Python
@@ -102,6 +113,15 @@ class Key:
 
     >>> key.replace(vars=None)
     Key(offsets={'x': 10})
+
+  You can use `indices` instead of `offsets` to refer to chunks by index::
+
+    >>> key = xarray_beam.Key(indices={'x': 4}, vars={'bar'})
+    >>> key
+    Key(indices={'x': 4}, vars={'bar'})
+    >>> key.with_indices(x=5)
+    Key(indices={'x': 5}, vars={'bar'})
+
   """
 
   # pylint: disable=redefined-builtin
@@ -110,25 +130,34 @@ def __init__(
       self,
       offsets: Mapping[str, int] | None = None,
       vars: Set[str] | None = None,
+      indices: Mapping[str, int] | None = None,
   ):
+    if offsets and indices:
+      raise ValueError("offsets and indices are mutually exclusive")
     if offsets is None:
       offsets = {}
+    if indices is None:
+      indices = {}
     if isinstance(vars, str):
       raise TypeError(f"vars must be a set or None, but is {vars!r}")
     self.offsets = immutabledict.immutabledict(offsets)
+    self.indices = immutabledict.immutabledict(indices)
     self.vars = None if vars is None else frozenset(vars)
 
   def replace(
       self,
       offsets: Mapping[str, int] | object = _DEFAULT,
       vars: Set[str] | None | object = _DEFAULT,
+      indices: Mapping[str, int] | 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 indices is _DEFAULT:
+      indices = self.indices
+    return type(self)(offsets, vars, indices)
 
   def with_offsets(self, **offsets: int | None) -> Key:
     """Replace some offsets with new values.
@@ -140,6 +169,8 @@ def with_offsets(self, **offsets: int | None) -> Key:
     Returns:
       New Key with the specified offsets.
     """
+    if self.indices:
+      raise ValueError("cannot call with_offsets on a Key with indices")
     new_offsets = dict(self.offsets)
     for k, v in offsets.items():
       if v is None:
@@ -148,31 +179,58 @@ def with_offsets(self, **offsets: int | None) -> Key:
         new_offsets[k] = v
     return self.replace(offsets=new_offsets)
 
+  def with_indices(self, **indices: int | None) -> Key:
+    """Replace some indices with new values.
+
+    Args:
+      **indices: indices to override (for integer values) or remove, with
+        values of ``None``.
+
+    Returns:
+      New Key with the specified indices.
+    """
+    if self.offsets:
+      raise ValueError("cannot call with_indices on a Key with offsets")
+    new_indices = dict(self.indices)
+    for k, v in indices.items():
+      if v is None:
+        del new_indices[k]
+      else:
+        new_indices[k] = v
+    return self.replace(indices=new_indices)
+
   def __repr__(self) -> str:
     components = []
     if self.offsets:
       components.append(f"offsets={dict(self.offsets)}")
+    if self.indices:
+      components.append(f"indices={dict(self.indices)}")
     if self.vars is not None:
       components.append(f"vars={set(self.vars)}")
     return f"{type(self).__name__}({', '.join(components)})"
 
   def __hash__(self) -> int:
-    return hash((self.offsets, self.vars))
+    return hash((self.offsets, self.vars, self.indices))
 
   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.indices == other.indices
+        and self.vars == other.vars
+    )
 
   def __ne__(self, other) -> bool:
     return not self == other
 
   # Beam uses these methods (also used for pickling) for "deterministic
   # encoding" of groupby keys
   def __getstate__(self):
-    offsets_state = sorted(self.offsets.items())
+    offsets_state = sorted(self.offsets.items()) if self.offsets else None
     vars_state = None if self.vars is None else sorted(self.vars)
-    return offsets_state, vars_state
+    indices_state = sorted(self.indices.items()) if self.indices else None
+    return offsets_state, vars_state, indices_state
 
   def __setstate__(self, state):
     self.__init__(*state)

xarray_beam/_src/core_test.py

@@ -47,6 +47,17 @@ def test_constructor(self):
     with self.assertRaisesRegex(TypeError, 'vars must be a set or None'):
       xbeam.Key(vars='foo')
 
+    key = xbeam.Key(indices={'x': 0, 'y': 1})
+    self.assertIsInstance(key.indices, immutabledict.immutabledict)
+    self.assertEqual(dict(key.indices), {'x': 0, 'y': 1})
+    self.assertEqual(dict(key.offsets), {})
+    self.assertIsNone(key.vars)
+
+    with self.assertRaisesRegex(
+        ValueError, 'offsets and indices are mutually exclusive'
+    ):
+      xbeam.Key(offsets={'x': 0}, indices={'x': 0})
+
   def test_replace(self):
     key = xbeam.Key({'x': 0}, {'foo'})
 
@@ -74,6 +85,33 @@ def test_replace(self):
     actual = key.replace({'y': 1}, {'bar'})
     self.assertEqual(expected, actual)
 
+  def test_replace_with_indices(self):
+    key_i = xbeam.Key(indices={'x': 0}, vars={'foo'})
+
+    expected = xbeam.Key(indices={'x': 1}, vars={'foo'})
+    actual = key_i.replace(indices={'x': 1})
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'y': 1}, vars={'foo'})
+    actual = key_i.replace(indices={'y': 1})
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'x': 0})
+    actual = key_i.replace(vars=None)
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'x': 0}, vars={'bar'})
+    actual = key_i.replace(vars={'bar'})
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'y': 1}, vars={'foo'})
+    actual = key_i.replace(indices={'y': 1}, vars={'foo'})
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'y': 1}, vars={'bar'})
+    actual = key_i.replace(indices={'y': 1}, vars={'bar'})
+    self.assertEqual(expected, actual)
+
   def test_with_offsets(self):
     key = xbeam.Key({'x': 0})
 
@@ -98,6 +136,42 @@ def test_with_offsets(self):
     actual = key2.with_offsets(x=1)
     self.assertEqual(expected, actual)
 
+    key_i = xbeam.Key(indices={'x': 0})
+    with self.assertRaisesRegex(
+        ValueError, 'cannot call with_offsets on a Key with indices'
+    ):
+      key_i.with_offsets(x=1)
+
+  def test_with_indices(self):
+    key = xbeam.Key(indices={'x': 0})
+
+    expected = xbeam.Key(indices={'x': 1})
+    actual = key.with_indices(x=1)
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'x': 0, 'y': 1})
+    actual = key.with_indices(y=1)
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={})
+    actual = key.with_indices(x=None)
+    self.assertEqual(expected, actual)
+
+    expected = xbeam.Key(indices={'y': 1, 'z': 2})
+    actual = key.with_indices(x=None, y=1, z=2)
+    self.assertEqual(expected, actual)
+
+    key2 = xbeam.Key(indices={'x': 0}, vars={'foo'})
+    expected = xbeam.Key(indices={'x': 1}, vars={'foo'})
+    actual = key2.with_indices(x=1)
+    self.assertEqual(expected, actual)
+
+    key_o = xbeam.Key(offsets={'x': 0})
+    with self.assertRaisesRegex(
+        ValueError, 'cannot call with_indices on a Key with offsets'
+    ):
+      key_o.with_indices(x=1)
+
   def test_repr(self):
     key = xbeam.Key({'x': 0, 'y': 10})
     expected = "Key(offsets={'x': 0, 'y': 10})"
@@ -107,6 +181,10 @@ def test_repr(self):
     expected = "Key(vars={'foo'})"
     self.assertEqual(repr(key), expected)
 
+    key = xbeam.Key(indices={'x': 0, 'y': 1})
+    expected = "Key(indices={'x': 0, 'y': 1})"
+    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}
@@ -115,13 +193,25 @@ def test_dict_key(self):
   def test_equality(self):
     key = xbeam.Key({'x': 0, 'y': 10})
     self.assertEqual(key, key)
-    self.assertNotEqual(key, None)
 
     key2 = xbeam.Key({'x': 0, 'y': 10}, {'bar'})
     self.assertEqual(key2, key2)
     self.assertNotEqual(key, key2)
     self.assertNotEqual(key2, key)
 
+    key_i = xbeam.Key(indices={'x': 0, 'y': 1})
+    self.assertEqual(key_i, key_i)
+    self.assertNotEqual(key_i, key)
+    self.assertNotEqual(key, key_i)
+
+    key_o = xbeam.Key(offsets={'x': 0, 'y': 1})
+    self.assertNotEqual(key_i, key_o)
+
+    key_i2 = xbeam.Key(indices={'x': 0, 'y': 1}, vars={'bar'})
+    self.assertEqual(key_i2, key_i2)
+    self.assertNotEqual(key_i, key_i2)
+    self.assertNotEqual(key_i2, key_i)
+
   def test_offsets_as_beam_key(self):
     inputs = [
         (xbeam.Key({'x': 0, 'y': 1}), 1),
@@ -135,6 +225,19 @@ def test_offsets_as_beam_key(self):
     actual = inputs | beam.GroupByKey()
     self.assertEqual(actual, expected)
 
+  def test_indices_as_beam_key(self):
+    inputs = [
+        (xbeam.Key(indices={'x': 0, 'y': 1}), 1),
+        (xbeam.Key(indices={'x': 0, 'y': 2}), 2),
+        (xbeam.Key(indices={'y': 1, 'x': 0}), 3),
+    ]
+    expected = [
+        (xbeam.Key(indices={'x': 0, 'y': 1}), [1, 3]),
+        (xbeam.Key(indices={'x': 0, 'y': 2}), [2]),
+    ]
+    actual = inputs | beam.GroupByKey()
+    self.assertEqual(actual, expected)
+
   def test_vars_as_beam_key(self):
     inputs = [
         (xbeam.Key(vars={'foo'}), 1),
@@ -155,6 +258,10 @@ def test_pickle(self):
     unpickled = pickle.loads(pickle.dumps(key))
     self.assertEqual(key, unpickled)
 
+    key_i = xbeam.Key(indices={'x': 0, 'y': 1}, vars={'foo'})
+    unpickled_i = pickle.loads(pickle.dumps(key_i))
+    self.assertEqual(key_i, unpickled_i)
+
 
 class TestOffsetsToSlices(test_util.TestCase):