perf(index): sample PQ training data before sub-vector copies#7930
perf(index): sample PQ training data before sub-vector copies#7930u70b3 wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughPQ training now caps sampled rows before sub-vector materialization, detects sample-size overflow, and validates the behavior with tests. KMeans benchmarks reuse cached inputs and add sampled and oversampled PQ build measurements. ChangesPQ Sampling and Benchmarking
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
rust/lance-index/src/vector/pq/utils.rs (1)
8-28: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a compiling doctest for
SubVectorKMeansData.The struct doc explains the row-major layout with a
textdiagram but has no compiling example exercisingnew()/vector(). As per coding guidelines, "Document all public APIs with examples and links to relevant structs and methods; keep examples synchronized with actual signatures."🤖 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-index/src/vector/pq/utils.rs` around lines 8 - 28, Update the public `SubVectorKMeansData` documentation with a compiling Rust doctest that constructs an instance through `new()` and exercises `vector()`, using the actual signatures and a small row-major input that demonstrates sub-vector selection. Keep the example synchronized with the current API and retain the existing layout explanation.Source: Coding guidelines
rust/lance-index/src/vector/kmeans.rs (2)
201-240: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
debug_assert!instead ofassert!for these construction invariants.
dimension > 0andvalues.len() % dimension == 0(and thelen <= self.len()check inslice_rows) are internal invariants guaranteed by every current caller, not conditions that prevent data corruption or handle fallible external input. Per guidelines these should bedebug_assert!/debug_assert_eq!rather thanassert!/assert_eq!, which will panic in release builds too.As per coding guidelines, "Prefer debug_assert! over assert! for non-safety invariants; reserve assert! for conditions preventing data corruption" and "Do not silently guard against impossible conditions; use debug_assert!, return an explicit error, or remove the check."
♻️ Proposed fix
impl<'a, T> ContiguousKMeansData<'a, T> { pub fn new(values: &'a [T], dimension: usize) -> Self { - assert!(dimension > 0, "dimension must be positive, got {dimension}"); - assert_eq!( + debug_assert!(dimension > 0, "dimension must be positive, got {dimension}"); + debug_assert_eq!( values.len() % dimension, 0, "values length {} is not a multiple of dimension {}", values.len(), dimension ); Self { values, dimension } } }fn slice_rows(&self, len: usize) -> Self { - assert!( + debug_assert!( len <= self.len(), "cannot slice {len} rows from {}", self.len() ); Self::new(&self.values[..len * self.dimension], self.dimension) }🤖 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-index/src/vector/kmeans.rs` around lines 201 - 240, Update ContiguousKMeansData::new to use debug_assert! and debug_assert_eq! for the positive dimension and divisible values-length invariants, and update slice_rows to use debug_assert! for its len <= self.len() check. Preserve the existing validation messages and slicing behavior.Source: Coding guidelines
168-213: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing doc examples on new public
KMeansData/ContiguousKMeansDataAPI.Both are
pub(required for benches to constructContiguousKMeansDatadirectly), but their doc comments have no usage example. Given this is a new core abstraction other code (and future external callers) will build on, add a short# Examplesblock showingContiguousKMeansData::new+.vector().As per coding guidelines, "Document all public APIs with examples and links to relevant structs and methods; keep examples synchronized with actual signatures."
🤖 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-index/src/vector/kmeans.rs` around lines 168 - 213, Add concise `# Examples` sections to the public `KMeansData` trait and `ContiguousKMeansData` API documentation, demonstrating construction with `ContiguousKMeansData::new` and reading a row via `.vector()`. Keep the examples valid for the actual signatures and link the trait documentation to the relevant concrete type and methods where appropriate.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@rust/lance-index/src/vector/kmeans.rs`:
- Around line 201-240: Update ContiguousKMeansData::new to use debug_assert! and
debug_assert_eq! for the positive dimension and divisible values-length
invariants, and update slice_rows to use debug_assert! for its len <= self.len()
check. Preserve the existing validation messages and slicing behavior.
- Around line 168-213: Add concise `# Examples` sections to the public
`KMeansData` trait and `ContiguousKMeansData` API documentation, demonstrating
construction with `ContiguousKMeansData::new` and reading a row via `.vector()`.
Keep the examples valid for the actual signatures and link the trait
documentation to the relevant concrete type and methods where appropriate.
In `@rust/lance-index/src/vector/pq/utils.rs`:
- Around line 8-28: Update the public `SubVectorKMeansData` documentation with a
compiling Rust doctest that constructs an instance through `new()` and exercises
`vector()`, using the actual signatures and a small row-major input that
demonstrates sub-vector selection. Keep the example synchronized with the
current API and retain the existing layout explanation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 2ab61146-2fac-4fbf-91e8-78d357e485ab
📒 Files selected for processing (4)
rust/lance-index/benches/kmeans.rsrust/lance-index/src/vector/kmeans.rsrust/lance-index/src/vector/pq/builder.rsrust/lance-index/src/vector/pq/utils.rs
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
PQ KMeans only trains on the sample_rate * num_centroids row prefix. Materialize sub-vectors from that prefix instead of copying every input row, preserving contiguous KMeans access without allocating unused scratch data. Lazy-initialize benchmark fixtures so filtered PQ benchmarks do not allocate unrelated 16 GiB training data.
d782e3f to
f23982c
Compare
u70b3
changed the title
There was a problem hiding this comment.
🧹 Nitpick comments (3)
rust/lance-index/src/vector/pq/builder.rs (2)
274-287: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert on the error variant, not just the message.
The test only checks
error.to_string().contains(...). Per coding guidelines, tests should assert on both the error variant and message content.As per coding guidelines, "Assert on both the error variant and the message content in tests; do not check only
is_err()."✅ Suggested fix
let error = params.build(&fsl, DistanceType::L2).unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); assert!( error .to_string() .contains("PQ training sample size overflows") );🤖 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-index/src/vector/pq/builder.rs` around lines 274 - 287, Update test_build_rejects_sample_size_overflow to assert that params.build returns the expected error variant, then separately verify the error message contains “PQ training sample size overflows.” Replace the message-only assertion while preserving the existing overflow setup and failure expectation.Source: Coding guidelines
112-120: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate sample-size computation vs. existing
sample_size()helper.
max_training_rowsreimplementsself.sample_rate * num_centroidsinline withchecked_mul, while the existingsample_size()method computes the same value with an unchecked*. Two sources of truth for the same quantity can silently diverge (e.g., if one is updated without the other, or ifsample_size()is used elsewhere and overflows unchecked).♻️ Suggested consolidation
- fn sample_size(&self) -> usize { - self.sample_rate * 2_usize.pow(self.num_bits as u32) - } + fn sample_size(&self, num_centroids: usize) -> Result<usize> { + self.sample_rate.checked_mul(num_centroids).ok_or_else(|| { + Error::invalid_input(format!( + "PQ training sample size overflows: sample_rate={}, num_centroids={num_centroids}", + self.sample_rate + )) + }) + }Then call
self.sample_size(num_centroids)?inbuild_from_fslinstead of duplicating the logic inline.🤖 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-index/src/vector/pq/builder.rs` around lines 112 - 120, Consolidate the training-row calculation in build_from_fsl by replacing the inline checked multiplication with the existing self.sample_size(num_centroids)? helper. Update sample_size to perform checked multiplication and return the appropriate Result/error so overflow handling remains consistent for every caller.rust/lance-index/src/vector/pq/utils.rs (1)
86-105: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test case for the
num_rows > fsl.len()error path.The new validation branch (
num_rows > fsl.len()at Lines 22-27) has no corresponding test; the currentrstestcases (all_rows,sampled_rows) both passnum_rows <= fsl.len(). A case asserting the error and its message would close this gap.As per coding guidelines, "Every bugfix and feature must have corresponding tests."
✅ Suggested additional case
#[rstest] #[case::all_rows(10)] #[case::sampled_rows(3)] fn test_divide_to_subvectors(#[case] num_rows: usize) { ... } + + #[test] + fn test_divide_to_subvectors_rejects_num_rows_overflow() { + let values = Float32Array::from_iter((0..320).map(|v| v as f32)); + let mat = FixedSizeListArray::try_new_from_values(values, 32).unwrap(); + let error = divide_to_subvectors::<Float32Type>(&mat, 4, 11).unwrap_err(); + assert!(error.to_string().contains("cannot divide 11 rows")); + }🤖 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-index/src/vector/pq/utils.rs` around lines 86 - 105, Add an rstest case to test_divide_to_subvectors that passes a num_rows value greater than the input FixedSizeListArray length, asserts the function returns an error, and verifies the error message from the num_rows > fsl.len() validation branch.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@rust/lance-index/src/vector/pq/builder.rs`:
- Around line 274-287: Update test_build_rejects_sample_size_overflow to assert
that params.build returns the expected error variant, then separately verify the
error message contains “PQ training sample size overflows.” Replace the
message-only assertion while preserving the existing overflow setup and failure
expectation.
- Around line 112-120: Consolidate the training-row calculation in
build_from_fsl by replacing the inline checked multiplication with the existing
self.sample_size(num_centroids)? helper. Update sample_size to perform checked
multiplication and return the appropriate Result/error so overflow handling
remains consistent for every caller.
In `@rust/lance-index/src/vector/pq/utils.rs`:
- Around line 86-105: Add an rstest case to test_divide_to_subvectors that
passes a num_rows value greater than the input FixedSizeListArray length,
asserts the function returns an error, and verifies the error message from the
num_rows > fsl.len() validation branch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: c7129d0c-e10e-44cd-9ea3-05e05802eef7
📒 Files selected for processing (3)
rust/lance-index/benches/kmeans.rsrust/lance-index/src/vector/pq/builder.rsrust/lance-index/src/vector/pq/utils.rs
1 participant
Summary
Related to #6928.
PQ codebook training previously materialized every input row into contiguous
per-sub-vector arrays before calling KMeans. KMeans then kept at most
sample_rate * num_centroidsrows, so oversized inputs paid to copy and retaindata that training never read.
This PR now applies that same row limit before
divide_to_subvectors. KMeanscontinues to receive contiguous arrays, but PQ only materializes the rows it can
actually pass to training.
The original version used strided KMeans views. Review and controlled
rebenchmarking found two problems with that approach:
KMeansAlgomethod signatures;The strided implementation has therefore been removed rather than hiding that
regression behind the oversized-input win.
This is intentionally a first-stage optimization for unused oversampled rows.
At or below the KMeans sample cap, the existing one-time contiguous
materialization remains. A zero-copy or bounded-scratch redesign is still
follow-up work tracked by #6928.
Changes
min(data.len(), sample_rate * num_centroids)rows.produces the same codebook as the previous materialize-then-sample order.
allocates an unrelated ~16 GiB KMeans fixture.
Validation
cargo test -p lance-index: 888 passed, 2 ignored, 0 failed.cargo test -p lance-indexdoctests: 7 passed.cargo clippy --all --tests --benches -- -D warnings.cargo fmt --all -- --check.Performance
Methodology:
mainat8f51a2117; PR:f23982c69.release-with-debugprofile.RAYON_NUM_THREADS=6, pinned to six physical cores withtaskset -c 0,2,4,6,8,10.estimates, run in main/PR then PR/main order.
End-to-end
PQBuildParams::build, 128 dimensions and 16 sub-vectors:The sampled-case confidence intervals overlap, so no sampled-path performance
change was observed. For 4M rows, peak RSS changed as follows:
The previous PR numbers were discarded. A filtered PQ benchmark still
initialized an unrelated ~16 GiB
bench_trainfixture, causing page churn, andlater attempts overlapped an external CPU-saturating workload. Lazy fixture
initialization and idle-host checks removed both sources of interference for
the results above.