diff --git a/grain/_src/python/data_sources.py b/grain/_src/python/data_sources.py index f885daeb8..2b5de615f 100644 --- a/grain/_src/python/data_sources.py +++ b/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.""" diff --git a/grain/_src/python/data_sources_test.py b/grain/_src/python/data_sources_test.py index a5cd2dfe8..708a356d3 100644 --- a/grain/_src/python/data_sources_test.py +++ b/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() diff --git a/grain/python/__init__.py b/grain/python/__init__.py index ed97769ec..58cf0f5fe 100644 --- a/grain/python/__init__.py +++ b/grain/python/__init__.py @@ -52,6 +52,7 @@ ) from grain._src.python.data_sources import ( ArrayRecordDataSource, + BagzDataSource, SharedMemoryDataSource as InMemoryDataSource, RangeDataSource, ) diff --git a/grain/sources.py b/grain/sources.py index b4089e7ff..41d923605 100644 --- a/grain/sources.py +++ b/grain/sources.py @@ -25,6 +25,7 @@ # all supported format dependencies. from grain._src.python.data_sources import ( ArrayRecordDataSource, + BagzDataSource, SharedMemoryDataSource, RangeDataSource, )