fix(index): return error instead of panicking when hierarchical k-means can't reach target_k#7869
Merged
wjones127 merged 1 commit intoJul 21, 2026
Conversation
…ns can't reach target_k num_partitions large enough to trigger hierarchical k-means (>256) could panic on debug_assert_eq!(heap.len(), target_k) in debug builds, or silently build a half-empty centroid set in release builds, when the training data can't be split into that many non-empty clusters (small dataset and/or many near-duplicate vectors). Return a descriptive InvalidArgumentError instead, which propagates through train_kmeans and do_train_ivf_model as a normal Result. Fixes lance-format#7867
Contributor
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
Author
|
@coderabbitai review |
Contributor
✅ Action performedReview finished.
|
wjones127
marked this pull request as ready for review
July 20, 2026 18:49
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
BubbleCal
approved these changes
Jul 21, 2026
wjones127
added a commit
that referenced
this pull request
Jul 21, 2026
…ns can't reach target_k (#7869) ## Summary Creating an IVF index with `num_partitions` large enough to trigger hierarchical k-means (> 256) could fail badly when the training data can't be split into that many non-empty clusters (small dataset and/or many near-duplicate vectors): - Debug builds: hard panic at `debug_assert_eq!(heap.len(), target_k)` in `train_hierarchical_kmeans`. - Release builds: the assert compiles out, silently producing far fewer than `num_partitions` non-empty partitions, degrading recall with no error. This replaces the `debug_assert_eq!` with a returned `ArrowError::InvalidArgumentError` describing how many clusters could actually be formed, so both build types fail cleanly instead of panicking or silently degrading. The error propagates through `train_kmeans` → `do_train_ivf_model` as a normal `Result`. ## Test plan - Added `test_hierarchical_kmeans_too_few_distinct_vectors_errors`, which builds a heavily near-duplicated dataset with `num_partitions` (300) far exceeding the number of distinct vectors, and asserts a descriptive error is returned instead of a panic. - Verified this test panics against the old `debug_assert_eq!` code (reproducing the reported issue) and passes with the fix. - `cargo test -p lance-index --lib vector::` (318 tests) passes. - `cargo fmt --all -- --check` and `cargo clippy -p lance-index --tests --benches -- -D warnings` pass. Fixes #7867 (cherry picked from commit a3c6fce)
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
2 participants
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
Creating an IVF index with
num_partitionslarge enough to trigger hierarchical k-means (> 256) could fail badly when the training data can't be split into that many non-empty clusters (small dataset and/or many near-duplicate vectors):debug_assert_eq!(heap.len(), target_k)intrain_hierarchical_kmeans.num_partitionsnon-empty partitions, degrading recall with no error.This replaces the
debug_assert_eq!with a returnedArrowError::InvalidArgumentErrordescribing how many clusters could actually be formed, so both build types fail cleanly instead of panicking or silently degrading. The error propagates throughtrain_kmeans→do_train_ivf_modelas a normalResult.Test plan
test_hierarchical_kmeans_too_few_distinct_vectors_errors, which builds a heavily near-duplicated dataset withnum_partitions(300) far exceeding the number of distinct vectors, and asserts a descriptive error is returned instead of a panic.debug_assert_eq!code (reproducing the reported issue) and passes with the fix.cargo test -p lance-index --lib vector::(318 tests) passes.cargo fmt --all -- --checkandcargo clippy -p lance-index --tests --benches -- -D warningspass.Fixes #7867