From 14a048a06f9917ad766c91b24e8132a509401d7a Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 14 Sep 2026 02:39:43 +0800 Subject: [PATCH 1/2] fix(index): return zero balance factor when no vector is assigned 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 --- rust/lance-index/src/vector/kmeans.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/rust/lance-index/src/vector/kmeans.rs b/rust/lance-index/src/vector/kmeans.rs index 9a0784d4845..c09b5203aed 100644 --- a/rust/lance-index/src/vector/kmeans.rs +++ b/rust/lance-index/src/vector/kmeans.rs @@ -296,6 +296,12 @@ fn compute_cluster_sizes( } }); + if max_cluster_size == 0 { + // No vector was assigned (all memberships are None): the division + // below yields NaN. Return 0 so the balance schedule degrades + // gracefully instead of depending on f32::min swallowing NaN. + return 0.0; + } (radius[max_cluster_id] - losses[max_cluster_id] as f32 / cluster_sizes[max_cluster_id] as f32) / membership.len() as f32 } From 556509dd27cd7ce726bdc0c99e2b06ffeba4f0b3 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 19 Sep 2026 01:04:22 +0800 Subject: [PATCH 2/2] fix(index): return zero balance factor when no vector is assigned --- rust/lance-index/src/vector/kmeans.rs | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/rust/lance-index/src/vector/kmeans.rs b/rust/lance-index/src/vector/kmeans.rs index c09b5203aed..94b801d1553 100644 --- a/rust/lance-index/src/vector/kmeans.rs +++ b/rust/lance-index/src/vector/kmeans.rs @@ -297,11 +297,13 @@ fn compute_cluster_sizes( }); if max_cluster_size == 0 { - // No vector was assigned (all memberships are None): the division - // below yields NaN. Return 0 so the balance schedule degrades - // gracefully instead of depending on f32::min swallowing NaN. + // Nothing was assigned (every membership is `None`, which is what an + // all-non-finite batch produces), so the division below would be by + // zero and the caller's `min` would silently drop the NaN. Say "no + // balancing" instead. return 0.0; } + (radius[max_cluster_id] - losses[max_cluster_id] as f32 / cluster_sizes[max_cluster_id] as f32) / membership.len() as f32 } @@ -2661,6 +2663,25 @@ mod tests { assert_eq!(allocate_quotas(&[2, 2], 10), vec![2, 2]); } + /// With every row unassigned (all-NaN input leaves membership `None`), the + /// largest cluster holds zero rows, so the balance factor divided by that + /// size is NaN. The NaN used to survive until `adjusted_balance_factor.min( + /// params.balance_factor)` quietly dropped it, which means the schedule + /// depended on `f32::min`'s NaN handling rather than on a decision. + #[test] + fn test_compute_cluster_sizes_without_any_assignment() { + let membership = vec![None, None, None]; + let radius = vec![0.0f32; 2]; + let losses = vec![0.0f64; 2]; + let mut cluster_sizes = vec![7usize, 9]; + + let factor = compute_cluster_sizes(&membership, &radius, &losses, &mut cluster_sizes); + + assert_eq!(factor, 0.0, "expected no balance factor, got {factor}"); + assert!(factor.is_finite()); + assert_eq!(cluster_sizes, vec![0, 0]); + } + #[test] fn test_split_clusters_takes_half_of_the_largest() { let mut cnts = vec![0, 10, 4];