Skip to content

fix(index): return zero balance factor when no vector is assigned - #9407

Open
LuciferYang wants to merge 3 commits into
lance-format:mainfrom
LuciferYang:fix/ivf-rq-kmeans-nan-balance-factor
Open

LuciferYang wants to merge 3 commits into
lance-format:mainfrom
LuciferYang:fix/ivf-rq-kmeans-nan-balance-factor

Conversation

@LuciferYang

Copy link
Copy Markdown
Contributor

Closes #9406

compute_cluster_sizes divided the largest cluster's loss by that cluster's size, which is zero when every membership is None, so it returned NaN. The caller's adjusted_balance_factor.min(params.balance_factor) then dropped the NaN because that is what f32::min does, leaving the next iteration's factor decided by an accident rather than by the code. It now returns 0.0 for that case, which says what the situation is: nothing was assigned, so there is nothing to balance.

The only inputs that reach it are batches where every vector is non-finite, which the index build path filters out before training; direct KMeans users can still hit it.

How was this patch tested?

test_compute_cluster_sizes_without_any_assignment calls the function with an all-None membership and asserts the factor is 0.0 and finite, and that the cluster sizes were reset. Removing the guard returns NaN and fails the assertion.

compute_cluster_sizes returned 0/0 = NaN when every membership was
None, and correctness of the next iteration silently depended on
Rust's f32::min returning the non-NaN operand. Guard explicitly.

Assisted-by: GLM-5.3
@github-actions github-actions Bot added bug Something isn't working A-index Vector index, linalg, tokenizer labels Sep 18, 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: approve.

The zero-assignment guard uses the documented 0.0 “no balancing” value, keeps NaN out of the next iteration, and preserves the cluster-size reset. The focused regression test covers the all-None boundary.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Sep 18, 2026
@LuciferYang

Copy link
Copy Markdown
Contributor Author

The rustdoc job is failing on main as well, not because of this change: run 35372502724 on main (83d0084) fails with 'error: could not document lance' from rust/lance/src/io/exec/knn.rs:563, a private intra-doc link to Self::try_new_batch under -D rustdoc::private-intra-doc-links. This PR touches rust/lance-index/src/vector/kmeans.rs only.

@lance-gatekeeper lance-gatekeeper Bot added K-approved Latest Gatekeeper recommendation permits acceptance. and removed K-approved Latest Gatekeeper recommendation permits acceptance. labels Sep 18, 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 bug Something isn't working K-approved Latest Gatekeeper recommendation permits acceptance.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug: the kmeans balance factor is NaN when nothing is assigned, and f32::min hides it

1 participant