Repository navigation
perf(index): sample PQ training data before sub-vector copies - #7930
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✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
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
97e593d to
9380dce
Compare
|
Rebased onto latest Could a maintainer clear the hold and take a look? Happy to address any feedback. |
a0338e7 to
f976530
Compare
|
@Xuanwo This PR optimizes PQ training by sampling before sub-vector materialization. It has now been rebased onto the latest |
f976530 to
eb70db0
Compare
![lance-gatekeeper[bot]](https://avatars.githubusercontent.com/in/4433141?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
lance-gatekeeper
Bot
added
the
K-approved
eb70db0 to
81230e3
Compare
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.
81230e3 to
a8171f8
Compare
lance-gatekeeper
Bot
removed
the
K-approved
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
lance-gatekeeper
Bot
added
the
K-approved
lance-gatekeeper
Bot
removed
the
K-approved
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The patch remains unchanged after the latest base merge. PQ/KMeans sampling and quantizer construction are unaffected, and the nearby index-metadata change integrates cleanly with ordinary vector-index creation. Keeping bounded contiguous materialization remains preferable to a strided KMeans interface because training reuses each sub-vector; checked sample sizing still rejects overflow before current dataset training data is loaded.
lance-gatekeeper
Bot
added
the
K-approved
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 retain data that training never read.This PR applies the same row limit before
divide_to_subvectors. KMeans continues to receive contiguous arrays and keeps the existing hot loops, while PQ only materializes the prefix it can actually pass to training.The original version used strided KMeans views. Review and controlled rebenchmarking found that this changed public
KMeansAlgosignatures and regressed the sampled 8-bit case. The strided implementation was removed rather than hiding that regression behind the oversized-input improvement.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.sample_size() -> usizeAPIs; add a fallibletry_sample_size()path and route all in-repository training and dataset sampling through it.InvalidInputerrors before loading training data.num_rows > FixedSizeListArray::len()validation path.Validation
cargo test -p lance-index: 1011 passed, 2 ignored, 0 failed.cargo test -p lance-indexdoctests: 9 passed.cargo check -p lance --lib.cargo clippy --all --tests --benches -- -D warnings.cargo fmt --all -- --check.Performance
Methodology:
f97653099, rebased ontomainatd20ca47d5.mainataf8e30869; PR ata0338e740. The only interveningmaincommit was the11.0.0-beta.4release metadata/lockfile version bump; it did not change PQ, KMeans, or other compiled code paths.release-with-debugprofile.RAYON_NUM_THREADS=6, pinned to six physical cores withtaskset -c 0,2,4,6,8,10.3e5208f32229626ab1696de9cde3ecf8d872eb32369d6fa67eeeb277af4dcba4) and temporarily used fixed-seed input arrays; that seed-only harness change is not part of the PR. KMeans centroid initialization still uses its existing OS-seeded path, so small changes with overlapping confidence intervals are treated as noise.End-to-end
PQBuildParams::build, 128 dimensions and 16 sub-vectors:Both sampled-case confidence intervals overlap, so no sampled-path performance change was observed. The 4M-row 4-bit intervals are clearly separated. The 4M-row 8-bit intervals overlap, so its -6.1% point-estimate change should be treated as directional rather than conclusive.
For 4M rows, peak RSS changed as follows:
The previous PR numbers were discarded because they used different random input arrays between builds and did not reject runs when host load returned during a long measurement. The refreshed results control input data, gate every run on low load, and exclude visibly contended attempts.