Internal

PiperOrigin-RevId: 836794431

Author: Grain Team

grain/_src/python/dataset/transformations/BUILD

@@ -142,6 +142,7 @@ py_test(
         "//grain/_src/python/dataset",
         "//grain/_src/python/dataset:base",
         "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
         "@pypi//numpy:pkg",
     ],
 )

grain/_src/python/dataset/transformations/mix.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import bisect
+import collections
 from collections.abc import Sequence
 import sys
 from typing import Any, Mapping, TypeVar
@@ -132,10 +133,38 @@ def __getitem__(self, index):
       if sys.version_info >= (3, 11):
         e.add_note(
             f"Exception caught while processing dataset @ {dataset_index=},"
-            f" {index_in_dataset=}"
         )
       raise e
 
+  def _getitems(self, indices: Sequence[int]) -> Sequence[T | None]:
+    """Returns a sequence of elements corresponding to the given indices."""
+
+    with self._stats.record_self_time(num_elements=len(indices)):
+      parent_and_key_pairs = (self._selection_map[i] for i in indices)
+
+      # Group indices by parent.
+      # keys_by_parent will store {parent_idx: [key1, key2, ...]}
+      # where key is the index in the parent dataset.
+      keys_by_parent = collections.defaultdict(list)
+      # original_indices_by_parent will store {parent_idx: [original_idx1, ...]}
+      # where original_idx is the position in the input `indices` list.
+      original_indices_by_parent = collections.defaultdict(list)
+
+      for i, (parent_idx, key) in enumerate(parent_and_key_pairs):
+        keys_by_parent[parent_idx].append(key)
+        original_indices_by_parent[parent_idx].append(i)
+
+    mixed_elements = [None] * len(indices)
+    for parent_idx, keys in keys_by_parent.items():
+      # Call _getitems for each parent.
+      parent_elements = self.parents[parent_idx]._getitems(keys)  # pylint: disable=protected-access
+      original_indices = original_indices_by_parent[parent_idx]
+      # Assign elements to their proper index in the final mixed dataset.
+      for i, element in enumerate(parent_elements):
+        mixed_elements[original_indices[i]] = element
+
+    return self._stats.record_output_spec_for_batch(mixed_elements)
+
 
 class _MixedDatasetIterator(dataset.DatasetIterator[T]):
   """Iterator that mixes elements from iterators based on given proportions.

grain/_src/python/dataset/transformations/mix_test.py

@@ -17,6 +17,7 @@
 from typing import Callable, Tuple
 
 from absl.testing import absltest
+from absl.testing import parameterized
 from grain._src.python.dataset import base
 from grain._src.python.dataset import dataset
 from grain._src.python.dataset.transformations import mix
@@ -154,7 +155,7 @@ def _subset_and_shuffle_dataset(index):
     self.assertEqual(expected_dataset, unrolled_dataset)
 
 
-class MixedMapDatasetTest(absltest.TestCase):
+class MixedMapDatasetTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
@@ -169,7 +170,7 @@ def test_len(self):
     # Equal proportions.
     ds = mix.MixedMapDataset([ds1, ds2, ds3])
     self.assertLen(ds, 15)
-    # Heigher weight for second dataset.
+    # Higher weight for second dataset.
     ds = mix.MixedMapDataset([ds1, ds2, ds3], proportions=[1, 2, 1])
     self.assertLen(ds, 5 + 10 + 5)
 
@@ -333,6 +334,99 @@ def _subset_and_shuffle_dataset(index):
 
     self.assertEqual(list(ds), expected_dataset)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="equal_proportions",
+          proportions=[1, 1],
+          expected_values=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+      ),
+      dict(
+          testcase_name="float_proportions",
+          proportions=[0.75, 0.25],
+          expected_values=[0, 2, 4, 1, 6, 8, 0, 3, 2, 4],
+      ),
+      dict(
+          testcase_name="integer_proportions",
+          proportions=[1, 2],
+          expected_values=[0, 1, 3, 2, 5, 7, 4, 9, 1, 6],
+      ),
+  )
+  def test_getitems_mixing(self, proportions, expected_values):
+    mixed_ds = mix.MixedMapDataset(
+        parents=[self.even_ds, self.odd_ds],
+        proportions=proportions,
+    )
+    # Request a batch of indices that results in interleaved elements
+    indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    actual_values = mixed_ds._getitems(indices)
+    self.assertEqual(expected_values, actual_values)
+
+  def test_getitems_interleaved_map(self):
+    def _inteleaved_dataset(index):
+      if index > 9:
+        raise IndexError("index our of range")
+      ds = index % 2
+      ds_index = index // 2
+      return (ds, ds_index)
+
+    interleaved_map = ExplicitSelectionMap(10, _inteleaved_dataset)
+
+    ds = mix.MixedMapDataset(
+        parents=[self.even_ds, self.odd_ds], selection_map=interleaved_map
+    )
+
+    expected_dataset = list(range(10))
+    self.assertEqual(ds._getitems(list(range(10))), expected_dataset)
+
+  def test_getitems_sequential_map(self):
+    def _sequential_dataset(index):
+      if index > 9:
+        raise IndexError("index our of range")
+      if index < 5:
+        ds = 0
+      else:
+        ds = 1
+      ds_index = index % 5
+      return (ds, ds_index)
+
+    sequential_map = ExplicitSelectionMap(10, _sequential_dataset)
+
+    ds = mix.MixedMapDataset(
+        parents=[self.even_ds, self.odd_ds], selection_map=sequential_map
+    )
+
+    expected_dataset = list(range(0, 10, 2)) + list(range(1, 10, 2))
+    self.assertEqual(ds._getitems(list(range(10))), expected_dataset)
+
+  def test_getitems_subset_and_shuffle_map(self):
+    first_epoch = [0, 1, 2, 3, 4]
+    second_epoch = [1, 0, 3, 2, 4]
+
+    expected_dataset = first_epoch + second_epoch
+
+    def _subset_and_shuffle_dataset(index):
+      if index > 9:
+        raise IndexError("index our of range")
+      if index < 5:
+        ds = index % 2
+        ds_index = index // 2
+      else:
+        mapped_index = second_epoch[index - 5]
+        ds = mapped_index % 2
+        ds_index = mapped_index // 2
+      return (ds, ds_index)
+
+    subset_and_shuffle_map = ExplicitSelectionMap(
+        10, _subset_and_shuffle_dataset
+    )
+
+    ds = mix.MixedMapDataset(
+        parents=[self.even_ds, self.odd_ds],
+        selection_map=subset_and_shuffle_map,
+    )
+
+    self.assertEqual(ds._getitems(list(range(10))), expected_dataset)
+
 
 class MixedIterDatasetTest(absltest.TestCase):
 
@@ -627,6 +721,17 @@ def test_cannot_concatenate_infinite_datasets(self):
     ):
       _ = mix._ConcatSelectionMap([zeros, ones])
 
+  def test_getitems_concatenate_finite_datasets(self):
+    evens = dataset.MapDataset.range(0, 10, 2)
+    odds = dataset.MapDataset.range(1, 10, 2)
+    ds = mix.ConcatenateMapDataset([evens, odds])
+    self.assertLen(evens, 5)
+    self.assertLen(odds, 5)
+    self.assertLen(ds, 10)
+
+    expected_values = [0, 2, 4, 6, 8, 1, 3, 5, 7, 9]
+    self.assertEqual(ds._getitems(list(range(10))), expected_values)
+
 
 if __name__ == "__main__":
   absltest.main()