From a80359d5fdca944b67bde4f0ef318879c7761752 Mon Sep 17 00:00:00 2001 From: ArrayRecord Team Date: Fri, 7 Aug 2026 11:16:40 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 961026949 --- python/array_record_data_source.py | 56 ++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/python/array_record_data_source.py b/python/array_record_data_source.py index 0eb7db6..be00698 100644 --- a/python/array_record_data_source.py +++ b/python/array_record_data_source.py @@ -441,6 +441,7 @@ def __init__( ) self._read_instructions = _get_read_instructions(paths) self._paths = [ri.filename for ri in self._read_instructions] +<<<<<<< dest: 31f73915ebdd - psgogineni: Add binding annotation f... self._reader_pool_size = ( reader_pool_size or _get_flag_value(_ARRAY_RECORD_READER_POOL_SIZE) or 1 # pyrefly: ignore[bad-argument-type] ) @@ -453,6 +454,14 @@ def __init__( for ri in self._read_instructions ] +||||||| parent of source: f245e4a325e1 - oncall-updater: Extending oncall for... + # We open readers lazily when we need to read from them. + self._readers = [None] * len(self._read_instructions) +======= + self._lock = threading.Lock() + # We open readers lazily when we need to read from them. + self._readers = [None] * len(self._read_instructions) +>>>>>>> source: 15fd492d466a - huangxia: Fix race condition in Arra... self._num_records = sum( map(lambda x: x.num_records, self._read_instructions) ) @@ -467,8 +476,21 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): logging.debug("__exit__ for ArrayRecordDataSource is called.") +<<<<<<< dest: 31f73915ebdd - psgogineni: Add binding annotation f... for pool in self._shard_pools: pool.close_all() +||||||| parent of source: f245e4a325e1 - oncall-updater: Extending oncall for... + for reader in self._readers: + if reader: + reader.close() + self._readers = [None] * len(self._read_instructions) +======= + with self._lock: + for reader in self._readers: + if reader: + reader.close() + self._readers = [None] * len(self._read_instructions) +>>>>>>> source: 15fd492d466a - huangxia: Fix race condition in Arra... def __len__(self) -> int: return self._num_records @@ -508,6 +530,7 @@ def _split_keys_per_reader( positions_and_indices[reader_idx] = [(position, idx)] return positions_and_indices +<<<<<<< dest: 31f73915ebdd - psgogineni: Add binding annotation f... def _read_record(self, reader: Any, position: int) -> bytes: """Helper to read a record using the best available method.""" if hasattr(reader, "read_record"): @@ -515,6 +538,31 @@ def _read_record(self, reader: Any, position: int) -> bytes: if hasattr(reader, "read"): return reader.read([position])[0] return reader[position] +||||||| parent of source: f245e4a325e1 - oncall-updater: Extending oncall for... + def _ensure_reader_exists(self, reader_idx: int) -> None: + """Threadsafe method to create corresponding reader if it doesn't exist.""" + if self._readers[reader_idx] is not None: + return + filename = self._read_instructions[reader_idx].filename + reader = _create_reader(filename, self._reader_options_string) + _check_group_size(filename, reader) + self._readers[reader_idx] = reader +======= + def _ensure_reader_exists(self, reader_idx: int) -> None: + """Threadsafe method to create corresponding reader if it doesn't exist.""" + if reader_idx < 0 or reader_idx >= len(self._readers): + raise IndexError( + f"reader_idx {reader_idx} out of range [0, {len(self._readers)})" + ) + if self._readers[reader_idx] is not None: + return + with self._lock: + if self._readers[reader_idx] is None: + filename = self._read_instructions[reader_idx].filename + reader = _create_reader(filename, self._reader_options_string) + _check_group_size(filename, reader) + self._readers[reader_idx] = reader +>>>>>>> source: 15fd492d466a - huangxia: Fix race condition in Arra... def __getitem__(self, record_key: SupportsIndex) -> bytes: pool_idx, position = self._reader_idx_and_position(record_key) @@ -565,12 +613,20 @@ def read_records( def __getstate__(self): logging.debug("__getstate__ for ArrayRecordDataSource is called.") state = self.__dict__.copy() +<<<<<<< dest: 31f73915ebdd - psgogineni: Add binding annotation f... state.pop("_shard_pools", None) +||||||| parent of source: f245e4a325e1 - oncall-updater: Extending oncall for... + del state["_readers"] +======= + del state["_readers"] + state.pop("_lock", None) +>>>>>>> source: 15fd492d466a - huangxia: Fix race condition in Arra... return state def __setstate__(self, state): logging.debug("__setstate__ for ArrayRecordDataSource is called.") self.__dict__.update(state) + self._lock = threading.Lock() # We open readers lazily when we need to read from them. Thus, we don't # need to re-open the same files as before pickling. self._shard_pools = [