perf: make compaction bin splitting linear#7763
Open
HaroldBenoit wants to merge 1 commit into
Open
Conversation
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
📝 WalkthroughWalkthroughChangesThe compaction planner rewrites candidate bin splitting to track row counts incrementally with Compaction bin splitting
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/dataset/optimize.rs`:
- Around line 2396-2409: Update the rstest attributes in
test_candidate_bin_split_large_and_small_edges to use readable case names with
the #[case::{name}(...)] syntax, assigning distinct descriptive names to each
row_counts scenario while preserving the existing arguments and assertions.
- Around line 1421-1486: Update the reset logic in split_for_size so
current_fragments, current_candidacy, and current_row_counts retain or receive a
capacity estimate after each cut instead of being left with zero capacity by
mem::take. Use the remaining fragment count or another safe upper-bound
estimate, while preserving the existing bin contents, cut behavior, and final
indices handling.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 57d002e0-046d-4914-ba0c-98ed9047cb67
📒 Files selected for processing (1)
rust/lance/src/dataset/optimize.rs
Comment on lines
+1421
to
1486
| fn split_for_size(self, min_num_rows: usize) -> Vec<Self> { | ||
| debug_assert_eq!(self.fragments.len(), self.candidacy.len()); | ||
| debug_assert_eq!(self.fragments.len(), self.row_counts.len()); | ||
| debug_assert_eq!(self.fragments.len(), self.pos_range.len()); | ||
|
|
||
| loop { | ||
| let mut bin_len = 0; | ||
| let mut bin_row_count = 0; | ||
| while bin_row_count < min_num_rows && bin_len < self.row_counts.len() { | ||
| bin_row_count += self.row_counts[bin_len]; | ||
| bin_len += 1; | ||
| } | ||
| if self.row_counts.is_empty() || min_num_rows == 0 { | ||
| return vec![self]; | ||
| } | ||
|
|
||
| // If there's enough remaining to make another worthwhile bin, then | ||
| // push what we have as a bin. | ||
| if self.row_counts[bin_len..].iter().sum::<usize>() >= min_num_rows { | ||
| let Self { | ||
| fragments, | ||
| pos_range, | ||
| candidacy, | ||
| row_counts, | ||
| indices, | ||
| } = self; | ||
|
|
||
| // A CandidateBin cannot contain enough entries for this sum to overflow u128. | ||
| let mut remaining_row_count = row_counts | ||
| .iter() | ||
| .map(|&row_count| row_count as u128) | ||
| .sum::<u128>(); | ||
| let min_num_rows = min_num_rows as u128; | ||
| let mut bins = Vec::new(); | ||
| let mut current_fragments = Vec::new(); | ||
| let mut current_candidacy = Vec::new(); | ||
| let mut current_row_counts = Vec::new(); | ||
| let mut current_row_count = 0_u128; | ||
| let mut current_pos = pos_range.start; | ||
|
|
||
| for ((fragment, candidacy), row_count) in | ||
| fragments.into_iter().zip(candidacy).zip(row_counts) | ||
| { | ||
| let row_count = row_count as u128; | ||
| remaining_row_count -= row_count; | ||
| current_row_count += row_count; | ||
| current_fragments.push(fragment); | ||
| current_candidacy.push(candidacy); | ||
| current_row_counts.push(row_count as usize); | ||
|
|
||
| if current_row_count >= min_num_rows && remaining_row_count >= min_num_rows { | ||
| let next_pos = current_pos + current_fragments.len(); | ||
| bins.push(Self { | ||
| fragments: self.fragments.drain(0..bin_len).collect(), | ||
| pos_range: self.pos_range.start..(self.pos_range.start + bin_len), | ||
| candidacy: self.candidacy.drain(0..bin_len).collect(), | ||
| row_counts: self.row_counts.drain(0..bin_len).collect(), | ||
| fragments: std::mem::take(&mut current_fragments), | ||
| pos_range: current_pos..next_pos, | ||
| candidacy: std::mem::take(&mut current_candidacy), | ||
| row_counts: std::mem::take(&mut current_row_counts), | ||
| // By the time we are splitting for size we are done considering indices | ||
| indices: Vec::new(), | ||
| }); | ||
| self.pos_range.start += bin_len; | ||
| } else { | ||
| // Otherwise, just push the remaining fragments into the last bin | ||
| bins.push(self); | ||
| break; | ||
| current_pos = next_pos; | ||
| current_row_count = 0; | ||
| } | ||
| } | ||
|
|
||
| // Any underfilled tail stays attached to the current group. | ||
| bins.push(Self { | ||
| fragments: current_fragments, | ||
| pos_range: current_pos..pos_range.end, | ||
| candidacy: current_candidacy, | ||
| row_counts: current_row_counts, | ||
| indices, | ||
| }); | ||
|
|
||
| bins | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Correctness verified — matches prior semantics with linear complexity.
Traced the cut condition (current_row_count >= min_num_rows && remaining_row_count >= min_num_rows) against split_for_size_reference for exact-multiple, boundary-sum, tail-merge, and all-in-one-bin cases; all produce identical output, and this is now backed by the randomized equivalence test.
One nit: current_fragments, current_candidacy, and current_row_counts are reset via mem::take (→ capacity 0) after every cut and regrow from scratch for each new bin. For workloads like the benchmarked 1M-fragment case (~30K+ bins), this causes repeated reallocation as each bin regrows to its final size.
♻️ Suggested capacity hint on reset
- bins.push(Self {
- fragments: std::mem::take(&mut current_fragments),
+ bins.push(Self {
+ fragments: std::mem::replace(&mut current_fragments, Vec::with_capacity(4)),
pos_range: current_pos..next_pos,
- candidacy: std::mem::take(&mut current_candidacy),
+ candidacy: std::mem::replace(&mut current_candidacy, Vec::with_capacity(4)),
row_counts: std::mem::take(&mut current_row_counts),As per coding guidelines, "Use Vec::with_capacity() when size is known or estimable, and prefer over-estimating capacity to multiple reallocations."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| fn split_for_size(self, min_num_rows: usize) -> Vec<Self> { | |
| debug_assert_eq!(self.fragments.len(), self.candidacy.len()); | |
| debug_assert_eq!(self.fragments.len(), self.row_counts.len()); | |
| debug_assert_eq!(self.fragments.len(), self.pos_range.len()); | |
| loop { | |
| let mut bin_len = 0; | |
| let mut bin_row_count = 0; | |
| while bin_row_count < min_num_rows && bin_len < self.row_counts.len() { | |
| bin_row_count += self.row_counts[bin_len]; | |
| bin_len += 1; | |
| } | |
| if self.row_counts.is_empty() || min_num_rows == 0 { | |
| return vec![self]; | |
| } | |
| // If there's enough remaining to make another worthwhile bin, then | |
| // push what we have as a bin. | |
| if self.row_counts[bin_len..].iter().sum::<usize>() >= min_num_rows { | |
| let Self { | |
| fragments, | |
| pos_range, | |
| candidacy, | |
| row_counts, | |
| indices, | |
| } = self; | |
| // A CandidateBin cannot contain enough entries for this sum to overflow u128. | |
| let mut remaining_row_count = row_counts | |
| .iter() | |
| .map(|&row_count| row_count as u128) | |
| .sum::<u128>(); | |
| let min_num_rows = min_num_rows as u128; | |
| let mut bins = Vec::new(); | |
| let mut current_fragments = Vec::new(); | |
| let mut current_candidacy = Vec::new(); | |
| let mut current_row_counts = Vec::new(); | |
| let mut current_row_count = 0_u128; | |
| let mut current_pos = pos_range.start; | |
| for ((fragment, candidacy), row_count) in | |
| fragments.into_iter().zip(candidacy).zip(row_counts) | |
| { | |
| let row_count = row_count as u128; | |
| remaining_row_count -= row_count; | |
| current_row_count += row_count; | |
| current_fragments.push(fragment); | |
| current_candidacy.push(candidacy); | |
| current_row_counts.push(row_count as usize); | |
| if current_row_count >= min_num_rows && remaining_row_count >= min_num_rows { | |
| let next_pos = current_pos + current_fragments.len(); | |
| bins.push(Self { | |
| fragments: self.fragments.drain(0..bin_len).collect(), | |
| pos_range: self.pos_range.start..(self.pos_range.start + bin_len), | |
| candidacy: self.candidacy.drain(0..bin_len).collect(), | |
| row_counts: self.row_counts.drain(0..bin_len).collect(), | |
| fragments: std::mem::take(&mut current_fragments), | |
| pos_range: current_pos..next_pos, | |
| candidacy: std::mem::take(&mut current_candidacy), | |
| row_counts: std::mem::take(&mut current_row_counts), | |
| // By the time we are splitting for size we are done considering indices | |
| indices: Vec::new(), | |
| }); | |
| self.pos_range.start += bin_len; | |
| } else { | |
| // Otherwise, just push the remaining fragments into the last bin | |
| bins.push(self); | |
| break; | |
| current_pos = next_pos; | |
| current_row_count = 0; | |
| } | |
| } | |
| // Any underfilled tail stays attached to the current group. | |
| bins.push(Self { | |
| fragments: current_fragments, | |
| pos_range: current_pos..pos_range.end, | |
| candidacy: current_candidacy, | |
| row_counts: current_row_counts, | |
| indices, | |
| }); | |
| bins | |
| } | |
| fn split_for_size(self, min_num_rows: usize) -> Vec<Self> { | |
| debug_assert_eq!(self.fragments.len(), self.candidacy.len()); | |
| debug_assert_eq!(self.fragments.len(), self.row_counts.len()); | |
| debug_assert_eq!(self.fragments.len(), self.pos_range.len()); | |
| if self.row_counts.is_empty() || min_num_rows == 0 { | |
| return vec![self]; | |
| } | |
| let Self { | |
| fragments, | |
| pos_range, | |
| candidacy, | |
| row_counts, | |
| indices, | |
| } = self; | |
| // A CandidateBin cannot contain enough entries for this sum to overflow u128. | |
| let mut remaining_row_count = row_counts | |
| .iter() | |
| .map(|&row_count| row_count as u128) | |
| .sum::<u128>(); | |
| let min_num_rows = min_num_rows as u128; | |
| let mut bins = Vec::new(); | |
| let mut current_fragments = Vec::new(); | |
| let mut current_candidacy = Vec::new(); | |
| let mut current_row_counts = Vec::new(); | |
| let mut current_row_count = 0_u128; | |
| let mut current_pos = pos_range.start; | |
| for ((fragment, candidacy), row_count) in | |
| fragments.into_iter().zip(candidacy).zip(row_counts) | |
| { | |
| let row_count = row_count as u128; | |
| remaining_row_count -= row_count; | |
| current_row_count += row_count; | |
| current_fragments.push(fragment); | |
| current_candidacy.push(candidacy); | |
| current_row_counts.push(row_count as usize); | |
| if current_row_count >= min_num_rows && remaining_row_count >= min_num_rows { | |
| let next_pos = current_pos + current_fragments.len(); | |
| bins.push(Self { | |
| fragments: std::mem::replace(&mut current_fragments, Vec::with_capacity(4)), | |
| pos_range: current_pos..next_pos, | |
| candidacy: std::mem::replace(&mut current_candidacy, Vec::with_capacity(4)), | |
| row_counts: std::mem::take(&mut current_row_counts), | |
| // By the time we are splitting for size we are done considering indices | |
| indices: Vec::new(), | |
| }); | |
| current_pos = next_pos; | |
| current_row_count = 0; | |
| } | |
| } | |
| // Any underfilled tail stays attached to the current group. | |
| bins.push(Self { | |
| fragments: current_fragments, | |
| pos_range: current_pos..pos_range.end, | |
| candidacy: current_candidacy, | |
| row_counts: current_row_counts, | |
| indices, | |
| }); | |
| bins | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/optimize.rs` around lines 1421 - 1486, Update the
reset logic in split_for_size so current_fragments, current_candidacy, and
current_row_counts retain or receive a capacity estimate after each cut instead
of being left with zero capacity by mem::take. Use the remaining fragment count
or another safe upper-bound estimate, while preserving the existing bin
contents, cut behavior, and final indices handling.
Source: Coding guidelines
Comment on lines
+2396
to
+2409
| #[rstest] | ||
| #[case(vec![1_000], 500)] | ||
| #[case(vec![10], 500)] | ||
| #[case(vec![1_000, 10], 500)] | ||
| #[case(vec![10, 1_000], 500)] | ||
| fn test_candidate_bin_split_large_and_small_edges( | ||
| #[case] row_counts: Vec<usize>, | ||
| #[case] min_num_rows: usize, | ||
| ) { | ||
| let original = candidate_bin(row_counts); | ||
| let split = original.clone().split_for_size(min_num_rows); | ||
|
|
||
| assert_eq!(split, vec![original]); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Name the rstest cases.
#[case(vec![1_000], 500)] etc. lack readable names. As per coding guidelines, "use #[case::{name}(...)] for readable Rust case names."
♻️ Suggested fix
- #[case(vec![1_000], 500)]
- #[case(vec![10], 500)]
- #[case(vec![1_000, 10], 500)]
- #[case(vec![10, 1_000], 500)]
+ #[case::large_only(vec![1_000], 500)]
+ #[case::small_only(vec![10], 500)]
+ #[case::large_then_small(vec![1_000, 10], 500)]
+ #[case::small_then_large(vec![10, 1_000], 500)]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| #[rstest] | |
| #[case(vec![1_000], 500)] | |
| #[case(vec![10], 500)] | |
| #[case(vec![1_000, 10], 500)] | |
| #[case(vec![10, 1_000], 500)] | |
| fn test_candidate_bin_split_large_and_small_edges( | |
| #[case] row_counts: Vec<usize>, | |
| #[case] min_num_rows: usize, | |
| ) { | |
| let original = candidate_bin(row_counts); | |
| let split = original.clone().split_for_size(min_num_rows); | |
| assert_eq!(split, vec![original]); | |
| } | |
| #[rstest] | |
| #[case::large_only(vec![1_000], 500)] | |
| #[case::small_only(vec![10], 500)] | |
| #[case::large_then_small(vec![1_000, 10], 500)] | |
| #[case::small_then_large(vec![10, 1_000], 500)] | |
| fn test_candidate_bin_split_large_and_small_edges( | |
| #[case] row_counts: Vec<usize>, | |
| #[case] min_num_rows: usize, | |
| ) { | |
| let original = candidate_bin(row_counts); | |
| let split = original.clone().split_for_size(min_num_rows); | |
| assert_eq!(split, vec![original]); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/optimize.rs` around lines 2396 - 2409, Update the
rstest attributes in test_candidate_bin_split_large_and_small_edges to use
readable case names with the #[case::{name}(...)] syntax, assigning distinct
descriptive names to each row_counts scenario while preserving the existing
arguments and assertions.
Source: Coding guidelines
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
1 participant
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.
Summary
CandidateBin::split_for_sizewith a single consuming passpos_range, and index-clearing behaviorPerformance
Release timings for uniform one-row fragments grouped at 30 rows:
Testing
cargo test -p lance candidate_bin_split --lib --quietcargo test -p lance test_candidate_bin --lib --quietcargo test -p lance benchmark_candidate_bin_split_scalability --lib --release -- --ignored --nocapturecargo fmt --all -- --checkcargo clippy -p lance --lib --tests --quiet -- -D warningsSummary by CodeRabbit
Bug Fixes
Tests