Internal

PiperOrigin-RevId: 799280753

Author: Grain Team

grain/_src/python/BUILD

@@ -25,7 +25,7 @@ py_test(
     srcs = [
         "data_sources_test.py",
     ],
-    args = ["--test_srcdir=grain/_src/python"],
+    args = ["--test_srcdir=grain/_src/python/testdata"],
     data = [
         "//grain/_src/python/testdata:digits.array_record-00000-of-00002",
         "//grain/_src/python/testdata:digits.array_record-00001-of-00002",
@@ -180,7 +180,7 @@ py_test(
     srcs = [
         "data_loader_test.py",
     ],
-    args = ["--test_srcdir=grain/_src/python"],
+    args = ["--test_srcdir=grain/_src/python/testdata"],
     data = [
         "//grain/_src/python/testdata:digits.array_record-00000-of-00002",
         "//grain/_src/python/testdata:digits.array_record-00001-of-00002",

grain/_src/python/data_loader_test.py

@@ -165,7 +165,7 @@ class DataLoaderTest(absl_parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.testdata_dir = pathlib.Path(FLAGS.test_srcdir) / "testdata"
+    self.testdata_dir = pathlib.Path(FLAGS.test_srcdir)
     self.read_options = (
         options.ReadOptions(num_threads=self.num_threads_per_worker)
         if (self.num_threads_per_worker is not None)

grain/_src/python/data_sources.py

@@ -122,6 +122,12 @@ def __getitem__(self, record_key: SupportsIndex) -> bytes:
     _bytes_read_counter.IncrementBy(len(data), "ArrayRecordDataSource")
     return data
 
+  @dataset_stats.trace_input_pipeline(stage_category=dataset_stats.IPL_CAT_READ)
+  def getitems(self, record_keys: Sequence[SupportsIndex]) -> Sequence[Any]:
+    # TODO: Modify the implementation to use underlying
+    # ArrayRecordReader to read multiple records at once.
+    return [self.__getitem__(record_key) for record_key in record_keys]
+
   @property
   def paths(self) -> ArrayRecordDataSourcePaths:
     return self._paths
@@ -156,6 +162,25 @@ def __getitem__(self, record_key: SupportsIndex) -> T:
     """
 
 
+@typing.runtime_checkable
+class RandomAccessDataSourceWithBatchedRead(
+    RandomAccessDataSource[T], Protocol, Generic[T]
+):
+  """Interface for datasources that support batched reads."""
+
+  def getitems(self, record_keys: Sequence[SupportsIndex]) -> Sequence[T]:
+    """Returns the values for the given record_keys.
+
+    This method must be threadsafe and deterministic.
+
+    Arguments:
+      record_keys: A sequence of integers in [0, len(self)-1].
+
+    Returns:
+      The sequence of corresponding records.
+    """
+
+
 class RangeDataSource:
   """Range data source, similar to python range() function."""
 

grain/_src/python/data_sources_test.py

@@ -45,6 +45,14 @@ def setUp(self):
     super().setUp()
     self.testdata_dir = pathlib.Path(FLAGS.test_srcdir)
 
+  def _to_full_paths(
+      self, paths: str | Sequence[str]
+  ) -> pathlib.Path | Sequence[pathlib.Path]:
+    if isinstance(paths, epath.PathLike):
+      return self.testdata_dir / paths
+    else:
+      return [(self.testdata_dir / path) for path in paths]
+
 
 class RangeDataSourceTest(DataSourceTest):
 
@@ -137,15 +145,88 @@ def test_str(self):
 
 class ArrayRecordDataSourceTest(DataSourceTest):
 
-  def test_array_record_data_implements_random_access(self):
+  def test_array_record_data_implements_random_access_with_batched_read(self):
     assert issubclass(
         data_sources.ArrayRecordDataSource, data_sources.RandomAccessDataSource
     )
+    assert issubclass(
+        data_sources.ArrayRecordDataSource,
+        data_sources.RandomAccessDataSourceWithBatchedRead,
+    )
 
   def test_array_record_source_empty_sequence(self):
     with self.assertRaises(ValueError):
       data_sources.ArrayRecordDataSource([])
 
+  @parameterized.parameters(
+      ("digits.array_record-00000-of-00002", False),
+      (
+          (
+              "digits.array_record-00000-of-00002",
+              "digits.array_record-00001-of-00002",
+          ),
+          True,
+      ),
+      (["digits.array_record-00000-of-00002"], False),
+  )
+  def test_array_record_data_source_sequential_get(
+      self, paths: str | Sequence[str], expect_both_shards: bool
+  ):
+    ar_ds = data_sources.ArrayRecordDataSource(paths=self._to_full_paths(paths))
+    expected_data = [b"0", b"1", b"2", b"3", b"4"]
+    if expect_both_shards:
+      expected_data += [b"5", b"6", b"7", b"8", b"9"]
+    actual_data = [ar_ds[i] for i in range(len(ar_ds))]
+    self.assertEqual(expected_data, actual_data)
+
+  @parameterized.parameters(
+      ("digits.array_record-00000-of-00002", False),
+      (
+          (
+              "digits.array_record-00000-of-00002",
+              "digits.array_record-00001-of-00002",
+          ),
+          True,
+      ),
+      (["digits.array_record-00000-of-00002"], False),
+  )
+  def test_array_record_data_source_get_all_items(
+      self, paths: str | Sequence[str], expect_both_shards: bool
+  ):
+    ar_ds = data_sources.ArrayRecordDataSource(paths=self._to_full_paths(paths))
+    data = [b"0", b"1", b"2", b"3", b"4"]
+    if expect_both_shards:
+      data += [b"5", b"6", b"7", b"8", b"9"]
+
+    indices_to_read = list(range(len(ar_ds)))
+    expected_data = [data[idx] for idx in indices_to_read]
+    actual_data = ar_ds.getitems(indices_to_read)
+    self.assertEqual(expected_data, actual_data)
+
+  @parameterized.parameters(
+      ("digits.array_record-00000-of-00002", False),
+      (
+          (
+              "digits.array_record-00000-of-00002",
+              "digits.array_record-00001-of-00002",
+          ),
+          True,
+      ),
+      (["digits.array_record-00000-of-00002"], False),
+  )
+  def test_array_record_data_source_get_multiple_items(
+      self, paths: str | Sequence[str], expect_both_shards: bool
+  ):
+    ar_ds = data_sources.ArrayRecordDataSource(paths=self._to_full_paths(paths))
+    data = [b"0", b"1", b"2", b"3", b"4"]
+    if expect_both_shards:
+      data += [b"5", b"6", b"7", b"8", b"9"]
+
+    indices_to_read = list(range(len(ar_ds)))[slice(0, None, 2)]
+    expected_data = [data[idx] for idx in indices_to_read]
+    actual_data = ar_ds.getitems(indices_to_read)
+    self.assertEqual(expected_data, actual_data)
+
 
 if __name__ == "__main__":
   absltest.main()