Repository navigation
Enable adaptive 4 MiB GCS readahead and thread-sticky reader pooling in ArrayRecordDataSource for gs:// shards. - #213
Open
copybara-service[bot] wants to merge 1 commit into
Open
copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
copybara-service
Bot
force-pushed
the
test_993008074
branch
2 times, most recently
from
October 5, 2026 00:41
5b9ab1f to
0fff3c6
Compare
…in `ArrayRecordDataSource` for `gs://` shards. When `ArrayRecordDataSource` reads from `gs://` paths: 1. `_create_reader` previously passed `readahead_buffer_size:0` and `_read_record` issued stateless `reader.read([position])[0]` point reads, resulting in one HTTP Range GET per record on GCS. 2. Unconditional stateful readahead (`seek` + `read`) improves sequential and small-shard reads, but causes read amplification on sparse random reads across large shards (`shard_size > readahead_buffer_size`) and readahead buffer thrashing when concurrent threads interleave single-record `__getitem__` calls on a single pooled reader. 3. The 0-arg `ArrayRecordReader.read()` PyBind11 binding previously held the Python GIL during `ReadRecord`, serializing multi-threaded stateful readahead fetches across Python threads. 4. In `ArrayRecordReaderBase::ParallelReadRecordsWithIndices` (`array_record_reader.cc`), `buf_size` (the byte span of the current coalesced group of chunks) was computed as `state_->ChunkEndOffset(i) - chunk_indices_per_buffer.back()[0]`. Because `state_->ChunkEndOffset(i)` is a file byte offset while `chunk_indices_per_buffer.back()[0]` is an integer chunk index (`0, 1, 2, ...`), subtracting a chunk index from a cumulative byte offset caused `buf_size` to evaluate to approximately `ChunkEndOffset(i)` (distance from the start of the file rather than from the start of the first chunk in the group). Consequently, for any chunk located past the buffer-size threshold into a shard, `buf_size` was always huge and adjacent chunks in `reader.read(indices)` could never be coalesced into a single byte-range read. This change: - Routes `gs://` paths in `_create_reader` to `_create_gcs_reader`, leaving all non-`gs://` paths untouched. - Configures `gs://` readers with `readahead_buffer_size:4194304` (4 MiB, configurable via `--array_record_gcs_readahead_buffer_size_bytes`) and `max_parallelism:1` by default while honoring explicit `reader_options` overrides. - Wraps `gs://` readers with readahead enabled in `_StatefulArrayRecordReader`, which adaptively routes small-shard reads (`_is_likely_small_shard`), in-buffer hits, and sequential accesses to stateful `seek(position)` + `read()`, while routing non-contiguous random reads on large shards to stateless `reader.read(positions)` / `reader.read([position])` to avoid 4 MiB read amplification and buffer eviction. - Adds thread-sticky reader affinity to `_BoundedReaderPool` (`borrow(sticky=True)` in `__getitem__` with default `gs://` pool size 16 so concurrent single-record prefetch threads keep dedicated readers, and `borrow(sticky=False)` in `__getitems__` so ephemeral batch worker threads reuse a single pooled reader per shard). - Releases the Python GIL (`py::gil_scoped_release`) in the 0-arg `ArrayRecordReader.read()` binding (`array_record_module.cc`). - Fixes `buf_size` in `ArrayRecordReaderBase::ParallelReadRecordsWithIndices` (`array_record_reader.cc`) to `state_->ChunkEndOffset(i) - (*state_->chunk_offsets)[chunk_indices_per_buffer.back()[0]]`, converting the starting chunk index into its starting byte offset in the file so `buf_size` accurately measures `end_byte_of_chunk_i - start_byte_of_first_chunk_in_group`. - Adds unit tests for `gs://` readahead option defaults/overrides, adaptive readahead routing on large shards, thread-sticky vs. batch reader pooling, and a local `_MockGcsServer` test verifying sequential readahead QPS. PiperOrigin-RevId: 993008074
copybara-service
Bot
force-pushed
the
test_993008074
branch
from
October 5, 2026 04:26
0fff3c6 to
5740c0a
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Enable adaptive 4 MiB GCS readahead and thread-sticky reader pooling in
ArrayRecordDataSourceforgs://shards.When
ArrayRecordDataSourcereads fromgs://paths:_create_readerpreviously passedreadahead_buffer_size:0and_read_recordissued statelessreader.read([position])[0]point reads, resulting in one HTTP Range GET per record on GCS.seek+read) improves sequential and small-shard reads, but causes read amplification on sparse random reads across large shards (shard_size > readahead_buffer_size) and readahead buffer thrashing when concurrent threads interleave single-record__getitem__calls on a single pooled reader.ArrayRecordReader.read()PyBind11 binding previously held the Python GIL duringReadRecord, serializing multi-threaded stateful readahead fetches across Python threads.ArrayRecordReaderBase::ParallelReadRecordsWithIndices(array_record_reader.cc),buf_size(the byte span of the current coalesced group of chunks) was computed asstate_->ChunkEndOffset(i) - chunk_indices_per_buffer.back()[0]. Becausestate_->ChunkEndOffset(i)is a file byte offset whilechunk_indices_per_buffer.back()[0]is an integer chunk index (0, 1, 2, ...), subtracting a chunk index from a cumulative byte offset causedbuf_sizeto evaluate to approximatelyChunkEndOffset(i)(distance from the start of the file rather than from the start of the first chunk in the group). Consequently, for any chunk located past the buffer-size threshold into a shard,buf_sizewas always huge and adjacent chunks inreader.read(indices)could never be coalesced into a single byte-range read.This change:
gs://paths in_create_readerto_create_gcs_reader, leaving all non-gs://paths untouched.gs://readers withreadahead_buffer_size:4194304(4 MiB, configurable via--array_record_gcs_readahead_buffer_size_bytes) andmax_parallelism:1by default while honoring explicitreader_optionsoverrides.gs://readers with readahead enabled in_StatefulArrayRecordReader, which adaptively routes small-shard reads (_is_likely_small_shard), in-buffer hits, and sequential accesses to statefulseek(position)+read(), while routing non-contiguous random reads on large shards to statelessreader.read(positions)/reader.read([position])to avoid 4 MiB read amplification and buffer eviction._BoundedReaderPool(borrow(sticky=True)in__getitem__with defaultgs://pool size 16 so concurrent single-record prefetch threads keep dedicated readers, andborrow(sticky=False)in__getitems__so ephemeral batch worker threads reuse a single pooled reader per shard).py::gil_scoped_release) in the 0-argArrayRecordReader.read()binding (array_record_module.cc).buf_sizeinArrayRecordReaderBase::ParallelReadRecordsWithIndices(array_record_reader.cc) tostate_->ChunkEndOffset(i) - (*state_->chunk_offsets)[chunk_indices_per_buffer.back()[0]], converting the starting chunk index into its starting byte offset in the file sobuf_sizeaccurately measuresend_byte_of_chunk_i - start_byte_of_first_chunk_in_group.gs://readahead option defaults/overrides, adaptive readahead routing on large shards, thread-sticky vs. batch reader pooling, and a local_MockGcsServertest verifying sequential readahead QPS.