Add BagzDataSource for BagZ file format support

grain/_src/python/data_sources.py

@@ -33,6 +33,7 @@
 from absl import logging
 from etils import epath
 from grain._src.core import monitoring
+from grain._src.python.dataset import base as dataset_base
 from grain._src.python.dataset import stats as dataset_stats
 
 
@@ -50,6 +51,13 @@ def __init__(self, *args, **kwargs):
   )
 # pylint: enable=g-import-not-at-top, g-importing-member, g-bad-import-order
 
+try:
+  import bagz as _bagz
+
+  _HAS_BAGZ = True
+except ImportError:
+  _HAS_BAGZ = False
+
 _api_usage_counter = monitoring.Counter(
     "/grain/python/data_sources/api",
     monitoring.Metadata(description="API initialization counter."),
@@ -105,6 +113,63 @@ def paths(self) -> ArrayRecordDataSourcePaths:
     return self._paths
 
 
+BagzDataSourcePaths = Union[
+    PathLikeOrFileInstruction, Sequence[PathLikeOrFileInstruction]
+]
+
+
+class BagzDataSource(dataset_base.RandomAccessDataSource):
+  """Data source for Bagz files."""
+
+  def __init__(self, paths: BagzDataSourcePaths):
+    """Creates a new BagzDataSource object.
+
+    Args:
+      paths: A single path/FileInstruction or list of paths/FileInstructions.
+    """
+    if not _HAS_BAGZ:
+      raise RuntimeError(
+          "bagz is not installed. Please install bagz to use BagzDataSource."
+      )
+    if isinstance(paths, (list, tuple)):
+      if not paths:
+        raise ValueError("paths cannot be an empty sequence.")
+      self._path = ",".join(sorted(str(epath.Path(p)) for p in paths))
+    else:
+      self._path = str(epath.Path(paths))
+    self._len = None
+    self._reader = None
+    _api_usage_counter.Increment("BagzDataSource")
+
+  @property
+  def reader(self):
+    if self._reader is None:
+      self._reader = _bagz.Reader(self._path)
+    return self._reader
+
+  def __len__(self) -> int:
+    if self._len is None:
+      self._len = len(_bagz.Reader(self._path))
+    return self._len
+
+  @dataset_stats.trace_input_pipeline(stage_category=dataset_stats.IPL_CAT_READ)
+  def __getitem__(self, record_key: SupportsIndex) -> bytes:
+    record_key = record_key.__index__()
+    return self.reader[record_key]
+
+  def __repr__(self) -> str:
+    return f"BagzDataSource(path={self._path!r})"
+
+  def __getstate__(self):
+    state = self.__dict__.copy()
+    state["_reader"] = None
+    return state
+
+  def __setstate__(self, state):
+    self.__dict__.update(state)
+    self._reader = None
+
+
 class RangeDataSource:
   """Range data source, similar to python range() function."""
 

grain/_src/python/data_sources_test.py

@@ -152,5 +152,56 @@ def test_array_record_source_empty_sequence(self):
       data_sources.ArrayRecordDataSource([])
 
 
+try:
+  import bagz as _bagz
+
+  _HAS_BAGZ = True
+except ImportError:
+  _HAS_BAGZ = False
+
+
+@absltest.skipIf(not _HAS_BAGZ, "bagz is not installed or supported.")
+class BagzDataSourceTest(DataSourceTest):
+
+  def setUp(self):
+    super().setUp()
+    import os
+    self.temp_dir = self.create_tempdir("bagz_test")
+    self.bagz_path = os.path.join(self.temp_dir.full_path, "test.bagz")
+    self.num_records = 10
+    w = _bagz.Writer(self.bagz_path)
+    for i in range(self.num_records):
+      w.write(f"record_{i}".encode("utf-8"))
+    w.close()
+
+  def test_bagz_implements_random_access(self):
+    self.assertTrue(
+        issubclass(
+            data_sources.BagzDataSource, dataset_base.RandomAccessDataSource
+        )
+    )
+
+  def test_bagz_source_empty_sequence(self):
+    with self.assertRaises(ValueError):
+      data_sources.BagzDataSource([])
+
+  def test_bagz_source_random_access(self):
+    source = data_sources.BagzDataSource([self.bagz_path])
+    self.assertLen(source, self.num_records)
+    self.assertEqual(source[0], b"record_0")
+    self.assertEqual(
+        source[self.num_records - 1],
+        f"record_{self.num_records - 1}".encode("utf-8"),
+    )
+
+  def test_bagz_source_pickling(self):
+    source = data_sources.BagzDataSource(self.bagz_path)
+    _ = source[0]  # instantiate reader
+    pickled = pickle.dumps(source)
+    unpickled = pickle.loads(pickled)
+    self.assertLen(unpickled, self.num_records)
+    self.assertEqual(unpickled[5], b"record_5")
+
+
 if __name__ == "__main__":
   absltest.main()

grain/python/__init__.py

@@ -52,6 +52,7 @@
 )
 from grain._src.python.data_sources import (
     ArrayRecordDataSource,
+    BagzDataSource,
     SharedMemoryDataSource as InMemoryDataSource,
     RangeDataSource,
 )

grain/sources.py

@@ -25,6 +25,7 @@
 #    all supported format dependencies.
 from grain._src.python.data_sources import (
     ArrayRecordDataSource,
+    BagzDataSource,
     SharedMemoryDataSource,
     RangeDataSource,
 )