No public description

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,13 +530,39 @@ 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"):
       return reader.read_record(position)
     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 = [