Skip to content

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
mainfrom
test_993008074
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_993008074

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Oct 4, 2026 •

Copy link
Copy Markdown

Enable adaptive 4 MiB GCS readahead and thread-sticky reader pooling 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.

@copybara-service
copybara-service Bot force-pushed the test_993008074 branch 2 times, most recently from 5b9ab1f to 0fff3c6 Compare October 5, 2026 00:41
…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

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant