Internal

PiperOrigin-RevId: 834497160

Author: Grain Team

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

@@ -44,6 +44,14 @@ def __getitem__(self, index):
       return self.slice(index)
     return tuple(p[index] for p in self._parents)
 
+  def _getitems(self, indices: Sequence[int]):
+    # p._getitems(indices) returns a list of elements of the requested indices.
+    # We get a list of lists that we need to zip.
+    parent_elements = [
+        p._getitems(indices) for p in self.parents  # pylint: disable=protected-access
+    ]
+    return list(zip(*parent_elements))
+
   def __str__(self) -> str:
     return f"ZipMapDataset(parents={self._parents}"
 

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

@@ -55,6 +55,19 @@ def test_getitem(self, ds_idx_list):
     for i in range(20):
       self.assertEqual(ds[i], tuple(i + ds_idx for ds_idx in ds_idx_list))
 
+  @parameterized.parameters(
+      {"ds_idx_list": x}
+      for x in list(itertools.combinations(range(3), 3))
+      + list(itertools.combinations(range(3), 2))
+      + list(itertools.combinations(range(3), 1))
+  )
+  def test_getitems(self, ds_idx_list):
+    ds = zip_ds.ZipMapDataset(parents=[self.ds_list[i] for i in ds_idx_list])
+    indices = [0, 5, 19]
+    expected_elements = [ds[i] for i in indices]
+    actual_elements = ds._getitems(indices)
+    self.assertEqual(expected_elements, actual_elements)
+
 
 class ZipIterDatasetTest(parameterized.TestCase):