Skip to content

fix(index): reject cosine and hamming in float kmeans instead of panicking - #9027

Open
LuciferYang wants to merge 1 commit into
lance-format:mainfrom
LuciferYang:fix/kmeans-reject-cosine
Open

fix(index): reject cosine and hamming in float kmeans instead of panicking#9027
LuciferYang wants to merge 1 commit into
lance-format:mainfrom
LuciferYang:fix/kmeans-reject-cosine

Conversation

@LuciferYang

@LuciferYang LuciferYang commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

Problem

KMeans::new_with_params dispatches a float column against any distance type, so cosine and hamming reach the float algorithm, whose membership pass panics with KMeans::find_partitions: cosine is not supported. compute_membership_and_distances has the same shape, so a cosine model built with the public with_centroids panics on assignment too. Both matches already end in an arm that returns a descriptive error.

There is one exception, and it is why lance.util.KMeans(k, "cosine") looked like it worked: when the centroid set is large enough for may_train_index to build an HNSW over it, the assignment goes through the index and never reads the distance type, so training completes. Below that threshold the same call panics.

Fixes #9026.

What this changes

Narrow the float arms in both dispatches to L2 | Dot, so cosine and hamming fall through to the error arm that was already there. Dot stays because the float path implements it.

lance.util.KMeans documents cosine, so the binding now normalizes its input and clusters with l2, which is what the index build path does before every training call. Cosine then behaves the same at every k rather than panicking below the HNSW threshold, and predict normalizes its input so a row and a scaled copy of it land in the same cluster.

The module and new_with_params doc comments claimed cosine inputs are normalized each iteration. That was implemented in #1723 and removed in #2015 while the comment was carried forward, so both now say what callers actually do.

One thing this does not cover: calling the lower-level compute_partitions or KMeansAlgoFloat::compute_membership_and_dist directly still panics for cosine, since the check lives in the dispatch rather than in the algorithm. Moving it would mean making the hot membership loop fallible.

Test plan

test_kmeans_rejects_cosine_and_hamming_for_floats covers training through both dispatch blocks: cosine and hamming over an f32 column at k=2, then cosine at k=257 for the hierarchical route. test_compute_membership_rejects_cosine_and_hamming_for_floats covers assignment through with_centroids, and asserts l2 still works so the narrowing is not over-wide. Restoring the wildcard arms makes each of them panic at kmeans.rs:511 instead.

test_kmeans_cosine covers the Python surface at k=8, which is the size that used to panic: it trains, and it asserts that scaling the input by 7 does not move any row's cluster, which is what fails if either fit or predict skips the normalization.

  • cargo test -p lance-index --lib vector::kmeans 17 passed, 1 ignored
  • cargo test -p lance --lib index::vector 285 passed, 1 ignored
  • uv run pytest python/tests/test_kmeans.py 4 passed
  • cargo clippy --all --tests --benches -- -D warnings and the same for python/Cargo.toml clean
  • cargo fmt --all --check clean, uv run make lint clean

@github-actions github-actions Bot added A-python Python bindings A-index Vector index, linalg, tokenizer bug Something isn't working labels Sep 7, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: request changes.

The training guards are sound, but assignment still validates unsupported metrics after optional index construction. Reject the data/centroid/metric combination before building SimpleIndex; that keeps invalid requests out of HNSW and makes the existing error arm authoritative.

))
}
(DataType::Float32, DataType::Float32, _) => {
(DataType::Float32, DataType::Float32, DistanceType::L2 | DistanceType::Dot) => {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This guard runs too late to guarantee the promised rejection: SimpleIndex::may_train_index has already consumed self.distance_type. For float/Hamming, the enabled (and sufficiently large automatic) index path builds FlatFloatStorage, then HNSW calls DistanceType::func::<f32> and panics at its Hamming => todo!() before this match can return InvalidArgumentError. Validate the type/metric combination before may_train_index, and only build the optional index after that validation.

Reproducer
LANCE_USE_HNSW_SPEEDUP_INDEXING=enabled cargo test -p lance-index --lib vector::kmeans::tests::test_compute_membership_rejects_cosine_and_hamming_for_floats -- --exact --nocapture

On the current head, the test fails with rust/lance-linalg/src/distance.rs:313: not yet implemented; the expected result is for its existing unwrap_err assertions to pass.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Sep 7, 2026
…cking

The training and assignment dispatches accepted any distance type for a float
column, and the membership pass then panicked with "not supported" unless the
centroid HNSW happened to be built, which bypasses that match. Narrow both
dispatches to L2 and Dot so the existing descriptive error arms handle the rest.

lance.util.KMeans documents cosine, so the binding now normalizes its input and
clusters with l2, the way the index build path does. That makes cosine behave
the same at every k instead of panicking below the HNSW threshold.

The module and new_with_params doc comments claimed cosine inputs are
normalized each iteration. That was implemented in lance-format#1723 and removed in lance-format#2015
while the comment was carried forward.
@LuciferYang
LuciferYang force-pushed the fix/kmeans-reject-cosine branch from 5450d01 to 8438ee9 Compare September 7, 2026 05:20
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Sep 7, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: request changes.

The Python cosine path now consistently normalizes and uses L2, but the indexed-assignment finding remains. Validate metric/type compatibility before SimpleIndex::may_train_index so float/Hamming cannot panic during HNSW construction.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Sep 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-index Vector index, linalg, tokenizer A-python Python bindings bug Something isn't working K-changes Latest Gatekeeper recommendation requests changes.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug: KMeans panics on cosine over a float column instead of rejecting it

1 participant