Repository navigation
fix(parquet): release the bytes of a row group when the push decoder is done with it - #11235
Conversation
fe3d9b5 to
c47bbf5
Compare
|
@alamb @HippoBaro this is ready for review if one of you can take a look |
| fn contains(&self, row_group_idx: usize) -> bool { | ||
| match self { | ||
| Self::Global { row_groups, .. } => row_groups.contains(&row_group_idx), | ||
| Self::PerRowGroup(row_groups) => row_groups | ||
| .iter() | ||
| .any(|row_group| row_group.row_group_index == row_group_idx), | ||
| } | ||
| } |
There was a problem hiding this comment.
This searches the remaining queue after every completed row group. For a normal scan of N distinct groups, this mean quadratic aggregate work. Could we maintain remaining occurrence counts as the queue advances, rather than repeatedly scanning it?
b859f9f to
49e7d0e
Compare
|
Thanks for reviewing @HippoBaro ! @alamb Hippo reviewed this PR already, do you think you could lend a second set of eyes? Thanks! |
332fb86 to
e5d880d
Compare
I am sorry -- your pR descriptions are better than most, but they are still very long and I think overly detailed. It would help me to review the pRs much more if you could cut the description down to a paragraph or two (or maybe hide the unecessary detail with Then perhaps you can call out any important code / strutural changes with inline comments |
|
I'll see what I can do! |
adriangb
left a comment
There was a problem hiding this comment.
Reading guide for reviewers: these comments show where to start and which parts of the diff are mechanical.
| RowGroupBuildResult::Finished { remaining_budget } => { | ||
| self.frontier | ||
| .update_budget_after_row_group(remaining_budget); | ||
| self.release_row_group(row_group_idx); |
There was a problem hiding this comment.
Start here. This is the fix: when a row group is done, release all of its column chunks. Before, the decoder released a buffer only if its range was equal to a requested range.
| /// Release the buffered bytes of a row group that is done, unless the | ||
| /// queue reads it again. The reader of the row group holds its own | ||
| /// copies of the bytes that it reads. | ||
| fn release_row_group(&mut self, row_group_idx: Option<usize>) { |
There was a problem hiding this comment.
A row group can be queued more than one time, so its bytes are kept while it is still queued.
| /// buffers with the same start is not specified. If no buffer overlaps | ||
| /// `ranges`, the buffers do not change. | ||
| #[cfg(feature = "arrow")] | ||
| pub(crate) fn release_ranges(&mut self, ranges: &[Range<u64>]) { |
There was a problem hiding this comment.
The new primitive. It trims or splits pushed buffers into zero-copy slices. release_ranges_matches_rebuild compares it with a simple rebuild of all buffers.
| /// reader builder. | ||
| #[derive(Debug, Clone)] | ||
| enum QueuedRowGroups { | ||
| struct QueuedRowGroups { |
There was a problem hiding this comment.
This diff looks larger than it is. The old QueuedRowGroups enum is renamed to Queue, and the only addition to it is row_group_indices. The new struct adds occurrence counts so that is_queued does not scan the queue (from the review by @HippoBaro).
| // The buffers can hold bytes that this decoder does not read. For | ||
| // example, the buffers of a rebuilt decoder can hold the bytes of row | ||
| // groups that the new configuration skips. | ||
| remaining_row_groups.release_unread_bytes(); |
There was a problem hiding this comment.
The second change, and the one behavior change for callers. With a new builder the buffers are empty, so this returns early. It only has an effect with with_buffers or into_builder.
| /// Remove the buffered bytes outside the read column chunks of | ||
| /// `row_groups`. A column chunk is read if the output or a predicate | ||
| /// reads its column. Indexes that are not in the file are ignored. | ||
| pub(crate) fn release_unread_bytes(&mut self, row_groups: impl IntoIterator<Item = usize>) { |
There was a problem hiding this comment.
Kept bytes = the column chunks of the output projection and of the predicate columns, for the queued row groups.
|
Thanks @alamb, I tried to shorten the description while keeping it relevant and posted some guiding comments. To make the PR itself smaller: if you want we can move the |
|
run benchmark push_decoder |
|
🤖 Benchmark starting (GKE) | trigger Target: push_decoder Sharding factor: 1 (1 workers). Comparing claude/push-buffers-release-coalesced (e5d880d) to 6b34163 (merge-base) diff Run configurationrun benchmark push_decoder
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: push_decoder Comparing claude/push-buffers-release-coalesced (e5d880d) to 6b34163 (merge-base) diff Run configurationrun benchmark push_decoder
shards: 1
DetailsPer-runner informationpush_decoder — shard 1/1Node: gk3-benchmark-cluster-nap-1qxrz5np-3edda623-k5pw Instance: c4a-highcpu-48 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench push_decoderCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
alamb
left a comment
There was a problem hiding this comment.
All in all this seems like an important step towards better memory usage in the parquet decoder -- thank you @adriangb and @HippoBaro
I left a bunch of small wording changes, that I think would help but are not required
Also, I think @HippoBaro filed another ticket about this (and a proposed fix) back in april. I updated the PR description to reflect this
| /// Provide a preexisting [`PushBuffers`] for the built decoder to read | ||
| /// from, so bytes already fetched are not requested again. | ||
| /// | ||
| /// [`Self::build`] releases the bytes that the decoder does not read: |
There was a problem hiding this comment.
I think this makes sense -- to put the cleaning as part of the builder
| /// buffers with the same start is not specified. If no buffer overlaps | ||
| /// `ranges`, the buffers do not change. | ||
| #[cfg(feature = "arrow")] | ||
| pub(crate) fn release_ranges(&mut self, ranges: &[Range<u64>]) { |
There was a problem hiding this comment.
Can this code be simplified once we merge the change to keep the ranges sorted in
?
There was a problem hiding this comment.
I don't think much simpler, but we could make it faster using binary search. I can do that in a followup PR?
There was a problem hiding this comment.
Filed #11458 for the binary search follow-up.
Address review comments on apache#11235 Co-Authored-By: Claude Opus 5.5 <[email protected]>
|
Thanks @alamb. I will try to address the remaining feedback today. |
# Which issue does this PR close? - Part of apache#11234. # Rationale for this change `ParquetPushDecoder` keeps the pushed buffers in a `PushBuffers` list. Each lookup scans the full list. If a caller pushes one buffer per requested range (one per column chunk, or one per page), the work for one row group grows with the square of the number of buffers. The existing `push_decoder` benchmark shows this for wide files: | Benchmark (`benches/push_decoder.rs`) | `main` | This PR | Speedup, run 1 / run 2 | Control (`main` vs `main`) | |---|---:|---:|---:|---:| | `Nbuf/100000ranges` | 893 ms | 58 ms | 15.4× / 14.6× | 1.03 | | `Nbuf/10000ranges` | 13.0 ms | 5.5 ms | 2.4× / 2.3× | 1.02 | | `Nbuf/1000ranges` | 581 µs | 541 µs | 1.07× / 1.05× | 1.01 | | `1buf/*` (one pushed buffer) | | | no change (±2%) | ±6% | (`Nbuf`: one pushed buffer per column chunk. Benchmark bot runs: [run 1](apache#11237 (comment)), [run 2](apache#11237 (comment)), [control](apache#11237 (comment)), which compares `main` with itself to show the noise.) The batch mode of apache#11234 pushes one buffer per page, which makes the list much longer. # What changes are included in this PR? `PushBuffers` keeps its buffers sorted by start offset, so a lookup is a binary search plus a short backward scan. Pushed buffers can still overlap. Removing buffers keeps the maximum buffer length exact, so a large buffer that was removed does not slow down later lookups. This PR is independent of apache#11235, which also changes `push_buffers.rs`. The PR that merges second gets a small rebase. A `max_len` that is too large is still correct (lookups only scan further), so the rebase updates `max_len` in `release_ranges` for performance, not for correctness. # Are these changes tested? Yes. Unit tests for out-of-order pushes, overlapping pushes, and `clear_ranges`, and a randomized test (200 seeds) that compares pushes, clears and lookups with a full scan and checks the sort order and the maximum length. # Are there any user-facing changes? No. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 <[email protected]>
Thank you -- this is a major step forward. I can't wait to get it all integrated into the scans more deeply |
…is done with it The decoder released a pushed buffer only if its range was equal to a requested range. A caller that fetched the requested ranges with one larger request kept all bytes until the end of the scan. Add PushBuffers::release_ranges, which trims or splits buffers into zero-copy slices, and release all column chunks of a row group when the decoder is done with it, unless the scan reads the row group again. Closes apache#11233. Co-Authored-By: Claude Opus 5.5 <[email protected]>
RemainingRowGroups checked whether a finished row group is queued again by scanning the whole queue after each row group, which is quadratic in the number of row groups. QueuedRowGroups now keeps the number of times each row group is queued, and updates it as the queue advances. Claude-Session: https://claude.ai/code/session_01XvkYF2BB52ejLkuZy5WnLW Co-Authored-By: Claude Opus 5.5 <[email protected]>
buffered_bytes summed every buffered range on each call, and callers query it per batch for memory accounting. Keep a running total instead. release_ranges also returns early when no buffer overlaps the released ranges, and skips the sort when the kept parts are already sorted. Claude-Session: https://claude.ai/code/session_01XvkYF2BB52ejLkuZy5WnLW Co-Authored-By: Claude Opus 5.5 <[email protected]>
…_ranges release_ranges rebuilt both Vecs and cloned every buffer on each call, although most calls trim or remove one buffer. Now it trims the overlapping buffers in place, removes the released ones, and does not touch the others. Claude-Session: https://claude.ai/code/session_01XvkYF2BB52ejLkuZy5WnLW Co-Authored-By: Claude Opus 5.5 <[email protected]>
… read `ParquetPushDecoderBuilder::build` now releases the buffered bytes outside the read column chunks of the queued row groups. A column chunk is read if the output projection or a RowFilter predicate reads its column. Thus a decoder rebuilt via `into_builder` does not keep the bytes of row groups or columns that the new configuration skips. `PushBuffers::retain_ranges` is implemented on top of `release_ranges`, and `column_chunk_range` is the one mapping from a column chunk to its bytes. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Address review comments on apache#11235 Co-Authored-By: Claude Opus 5.5 <[email protected]>
…variant - Explain in the docs of `buffered_bytes` and `release_ranges` that the kept slices record the bytes the decoder can still read, not the allocated memory. - The buffers are now always sorted (a lookup invariant), so the tests call `assert_invariants` and check lookups, not the internal order. - Do not underflow on an inverted range in `update_max_len`. Co-Authored-By: Claude Opus 5.5 <[email protected]>
14c209b to
207057d
Compare
|
@alamb I think I've addressed the review and filed followups. Ready to merge from my perspective! |
|
Looks good -- let's do it |
|
Thank you @adriangb |
…of the buffer count (apache#11459) # Which issue does this PR close? - Closes apache#11458. Follow-up to apache#11235, which is merged. This PR has four commits: 1. apache@8087cef (new `push_decoder/prefetch` benchmark) 2. apache@c0857e7 (binary search in `release_ranges`) 3. apache@04faa64 (binary search in `clear_ranges`) 4. apache@8b98d70 (`VecDeque` and `max_len_count`) # Rationale for this change When a caller pushes many buffers before the decoder reads them (for example, read-ahead for many row groups), the push decoder pays `O(buffers)` for each row group. Thus, the full scan is quadratic. The new `push_decoder/prefetch` benchmark shows this. It pushes one buffer for each column chunk of the file, then builds all row group readers. Time for each iteration on the GKE runner: | Buffers | Before | After | Change | |--|--|--|--| | 1,000 | 1.12 ms | 0.60 ms | 1.9x faster | | 10,000 | 56.6 ms | 6.3 ms | 9x faster | | 30,000 | 478 ms | 19.2 ms | 25x faster | After this PR, the time is linear in the number of buffers. The cost: the `VecDeque` commit makes the existing `push_decoder/1buf` and `push_decoder/Nbuf` cases 1% to 6% slower, because an index into a `VecDeque` costs more than an index into a `Vec`. Those cases hold few buffers at a time, so they get no gain from this PR. # What changes are included in this PR? For each row group, the decoder calls `clear_ranges` and then `release_ranges`. Each call visited all buffers: ```diff for each row group clear_ranges(requested ranges) - rebuild all buffers, keep those that are not requested - scan all buffers for max_len + binary search for the buffers that start at each requested range + remove them, move no other buffer (VecDeque) + scan for max_len only if no buffer of that length remains release_ranges(column chunk ranges) - scan all buffers for an overlap + binary search for the buffers that can overlap ``` - **`release_ranges` (commit 2):** finds the buffers that can overlap a released range with a binary search, as `PushBuffers::find` does. The trim and split logic does not change. - **`clear_ranges` (commit 3):** finds the buffers to remove with a binary search. Before, it built new buffer lists on each call. This commit gives most of the gain. - **`VecDeque` and `max_len_count` (commit 4):** `PushBuffers` stores the buffers in a `VecDeque`, so a removal at the start or at the end moves no other buffer. A count of the buffers with length `max_len` lets a removal skip the scan that updates `max_len` while other buffers of that length remain. <details> <summary>Benchmark numbers for each commit, and edge cases</summary> `push_decoder/prefetch` on the GKE runner. Each column adds one commit: | Buffers | Benchmark only | + `release_ranges` | + `clear_ranges` | + `VecDeque` | |--|--|--|--|--| | 1,000 | 1.12 ms | 1.09 ms | 0.64 ms | 0.60 ms | | 10,000 | 56.6 ms | 53.1 ms | 11.9 ms | 6.3 ms | | 30,000 | 478 ms | 447 ms | 73.4 ms | 19.2 ms | Change of the existing `1buf` and `Nbuf` cases for each commit: | Step | Change | |--|--| | Control (same code on both sides) | −3% to +2% | | + `release_ranges` | −1% to +2% | | + `clear_ranges` | −2% to +2% | | + `VecDeque` | +1% to +6% (slower in all seven cases) | Runs: - Benchmark only compared with `release_ranges`: apache#11459 (comment) - `release_ranges` compared with `clear_ranges`: apache#11459 (comment) - `clear_ranges` compared with `VecDeque`: apache#11459 (comment) - Control: apache#11459 (comment) These runs used the commits before the rebase onto `main`. The rebase changed only the benchmark file. Time for each `release_ranges` call, from temporary timing tests that are not part of this PR (release build, Apple silicon): | Buffers | Call | Before | After commit 2 | |--|--|--|--| | 1,000 | releases nothing | 495 ns | 55 ns | | 10,000 | releases nothing | 3.5 µs | 67 ns | | 100,000 | releases nothing | 37 µs | 61 ns | | 1,000 | trims one buffer | 1.5 µs | 82 ns | | 10,000 | trims one buffer | 8.5 µs | 91 ns | | 100,000 | trims one buffer | 86 µs | 81 ns | Edge cases: - Removal in the middle: `VecDeque::drain` moves the buffers on the shorter side of the removed buffers. Thus, only a removal at the start or at the end moves none. - `max_len`: it stays exact. It is computed again from all buffers only when the last buffer of length `max_len` is removed or trimmed. - Empty buffers: `push_range` accepts empty and inverted ranges, and `release_ranges` removes all of them when it releases bytes. To keep this behavior without a visit of all buffers, `PushBuffers` counts the empty buffers. `release_ranges` looks at all buffers only if the count is not zero. - Sort order: a trim moves the start of a buffer only past released bytes, so the trimmed buffers stay in order. Only split parts can be out of order, and the sort runs only after a split. - Buffers with the same start: they are adjacent but not sorted by end. `clear_ranges` checks each of them for an exact match. </details> # Are these changes tested? Yes. The existing tests cover this, in particular `fuzz_release_ranges_matches_rebuild`, which compares `release_ranges` with a full rebuild, and `fuzz_matches_a_linear_scan`, which includes random calls to `clear_ranges`. The new test `clear_ranges_with_buffers_of_the_same_start` covers buffers that share a start and ranges that are given more than once. `assert_invariants` now also checks the count of empty buffers and the count of buffers with length `max_len`. # Are there any user-facing changes? No. All changes are in `pub(crate)` functions and private fields, so a revert is safe if a regression shows. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 <[email protected]>
Which issue does this PR close?
PushBuffers::clear_rangesis quadratic and leaks #9695PushBuffersboundary-agnostic for prefetch IO #9697Rationale for this change
ParquetPushDecoderdoes not release a pushed buffer that is larger than a requested range. A caller that merges nearby ranges into one object storage request thus holds all the bytes of the scan until the scan ends.buffered_bytes()before this PRThe reproduction is in #11233.
What changes are included in this PR?
into_builder().build()releases the bytes that the rebuilt decoder does not read, for example the bytes of row groups that a dynamic filter pruned.The inline comments show where to start and which parts of the diff are mechanical.
Details: bytes that stay buffered, internals, performance commits
These bytes stay buffered, as before, until
clear_all_ranges:buffered_bytes()counts the bytes that the decoder holds. The kept parts are zero-copy slices, so the allocator frees a pushedBytesonly after all of its slices are dropped.Internals:
PushBuffers::release_rangestrims or splits pushed buffers.retain_rangescallsrelease_rangeswith the complement. #11234 uses them again.The commits after the fix make the bookkeeping cheaper (from the review by @HippoBaro):
buffered_bytes()is O(1);release_rangestrims buffers in place.This PR includes #11236, which changed the same code.
Are these changes tested?
Yes.
test_decoder_releases_coalesced_pushis the reproduction from the issue, and it fails without the change.All new and changed tests
test_decoder_releases_coalesced_pushtest_decoder_releases_unrequested_bytes_of_row_grouptest_decoder_releases_coalesced_push_with_predicateRowFilter, which makes two requests per row grouptest_decoder_clear_all_ranges(changed)test_decoder_keeps_bytes_of_a_row_group_that_is_read_againtest_into_builder_releases_bytes_of_skipped_row_groupstest_into_builder_preserves_buffered_bytes(changed)test_into_builder_keeps_bytes_of_predicate_columnsThere are also unit tests for
release_ranges(one compares it with a rebuild of all buffers) andretain_ranges.Are there any user-facing changes?
There is no API change. There are two behavior changes:
buffered_bytes()decreases after each row group when the caller pushes buffers that are larger than the requested ranges.into_builder().build()(orbuild()withwith_buffers), the decoder holds only the bytes that it reads. A caller that narrows the projection and then widens it again must push the other columns again.🤖 Generated with Claude Code