From da828f81c18880c9ce44b038e437e22926f395eb Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 11 Sep 2026 14:09:43 +0800 Subject: [PATCH 01/10] fix(index): reject unattainable partition counts in weighted kmeans MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming coreset trainer padded a short cluster list by duplicating the heaviest cluster until it reached the requested partition count. argmin ties always resolve to the first copy, so every duplicate stayed empty for the life of the index — dead partitions inflating num_partitions and probe lists. Return the same descriptive error the flat hierarchical trainer already uses, telling the caller to reduce num_partitions. Assisted-by: GLM-5.3 --- rust/lance/src/index/vector/ivf.rs | 63 +++++++++++++++++++++++------- 1 file changed, 48 insertions(+), 15 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index da21797e09c..338c0848739 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4497,21 +4497,19 @@ fn train_weighted_hierarchical_f32_kmeans( let mut clusters = heap.into_vec(); clusters.sort_by_key(|cluster| cluster.id); - while clusters.len() < target_k { - let duplicate = clusters - .iter() - .max_by(|left, right| { - left.weight - .partial_cmp(&right.weight) - .unwrap_or(std::cmp::Ordering::Equal) - }) - .cloned() - .ok_or_else(|| Error::index("No weighted clusters were trained"))?; - clusters.push(WeightedCluster { - id: next_cluster_id, - ..duplicate - }); - next_cluster_id += 1; + if clusters.len() < target_k { + // Duplicating the heaviest cluster to pad the count used to leave + // permanently empty partitions: argmin ties always resolve to the + // first copy, so no vector is ever assigned to a duplicate. Match + // the flat hierarchical trainer, which rejects the request. + return Err(Error::invalid_input(format!( + "Cannot create {target_k} IVF partitions: weighted k-means could only form {} \ + non-empty clusters. The dataset is likely too small or has too many \ + (near-)duplicate vectors for this many partitions. Reduce num_partitions to \ + <= {} or provide more diverse data.", + clusters.len(), + clusters.len() + ))); } clusters.truncate(target_k); @@ -6413,6 +6411,41 @@ mod tests { ); } + /// Duplicating the heaviest cluster to pad the partition count used to + /// yield permanently empty partitions (argmin ties always resolve to + /// the first copy); reject like the flat hierarchical trainer. + #[test] + fn test_weighted_hierarchical_rejects_padding_with_duplicates() { + let dimension = 8; + let num_points = 6; + let values: Vec = (0..num_points * dimension) + .map(|i| (i % 251) as f32 * 0.01) + .collect(); + let data = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension as i32) + .unwrap(); + // Six distinct single-member clusters can never reach 16, so the + // trainer used to pad with duplicates of the heaviest cluster. + let params = WeightedHierarchicalKMeansParams { + dimension, + target_k: 16, + metric_type: MetricType::L2, + max_iters: 2, + on_progress: Arc::new(|_, _| {}), + }; + let err = train_weighted_hierarchical_f32_kmeans( + &data, + &vec![1.0; num_points], + &vec![0.1; num_points], + ¶ms, + ) + .unwrap_err(); + assert!( + err.to_string().contains("Cannot create 16 IVF partitions"), + "got: {err}" + ); + } + /// Regression test for a hang in the streaming *coreset* trainer /// (`train_streaming_coreset_ivf_model`, taken when `num_partitions > 256`). /// From 945c6f617d5f4cd84c2f383d24290fc7d75f3c80 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 17 Sep 2026 18:39:10 +0800 Subject: [PATCH 02/10] fix(index): reject unattainable partition counts in weighted kmeans --- rust/lance/src/index/vector/ivf.rs | 70 ++++++++++++++++++++++-------- 1 file changed, 51 insertions(+), 19 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 338c0848739..82f10dc27d8 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4499,19 +4499,29 @@ fn train_weighted_hierarchical_f32_kmeans( clusters.sort_by_key(|cluster| cluster.id); if clusters.len() < target_k { // Duplicating the heaviest cluster to pad the count used to leave - // permanently empty partitions: argmin ties always resolve to the - // first copy, so no vector is ever assigned to a duplicate. Match - // the flat hierarchical trainer, which rejects the request. + // permanently empty partitions: a padded centroid is a bit-identical + // clone, and assignment keeps the incumbent on a tie, so no vector is + // ever assigned to a duplicate. Reject instead, like the flat + // hierarchical trainer. The shortfall has two causes worth + // distinguishing in a bug report but not in the advice: the coreset + // rows are too alike to split further, or splitting stopped early. return Err(Error::invalid_input(format!( "Cannot create {target_k} IVF partitions: weighted k-means could only form {} \ - non-empty clusters. The dataset is likely too small or has too many \ - (near-)duplicate vectors for this many partitions. Reduce num_partitions to \ - <= {} or provide more diverse data.", + non-empty clusters from {} weighted coreset rows, and the rest could not be \ + split further. Reduce num_partitions to <= {} or provide more diverse data.", clusters.len(), + data.len(), clusters.len() ))); } - clusters.truncate(target_k); + // The split loop stops at `target_k`, and every `cluster_k` it takes is + // bounded by what is left, so the count lands exactly on `target_k`. + debug_assert_eq!( + clusters.len(), + target_k, + "weighted kmeans formed {} clusters for target_k {target_k}", + clusters.len() + ); let mut values = Vec::with_capacity(target_k * dimension); for cluster in clusters { @@ -6412,23 +6422,38 @@ mod tests { } /// Duplicating the heaviest cluster to pad the partition count used to - /// yield permanently empty partitions (argmin ties always resolve to - /// the first copy); reject like the flat hierarchical trainer. + /// yield permanently empty partitions (a padded centroid is a clone, and + /// assignment keeps the incumbent on a tie); reject like the flat + /// hierarchical trainer. + /// + /// Two shapes reach the shortfall. Fewer rows than partitions is the + /// obvious one; the one the only production caller can actually produce is + /// a coreset with enough rows whose clusters cannot be split further, + /// because the coreset always carries at least `num_partitions` rows. + #[rstest::rstest] + #[case::fewer_rows_than_partitions(6, 16)] + #[case::enough_rows_all_identical(32, 16)] #[test] - fn test_weighted_hierarchical_rejects_padding_with_duplicates() { + fn test_weighted_hierarchical_rejects_padding_with_duplicates( + #[case] num_points: usize, + #[case] target_k: usize, + ) { let dimension = 8; - let num_points = 6; - let values: Vec = (0..num_points * dimension) - .map(|i| (i % 251) as f32 * 0.01) - .collect(); + // Every row identical in the second case, a ramp in the first: either + // way the trainer cannot reach `target_k` non-empty clusters. + let values: Vec = if num_points < target_k { + (0..num_points * dimension) + .map(|i| i as f32 * 0.01) + .collect() + } else { + vec![0.5; num_points * dimension] + }; let data = FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension as i32) .unwrap(); - // Six distinct single-member clusters can never reach 16, so the - // trainer used to pad with duplicates of the heaviest cluster. let params = WeightedHierarchicalKMeansParams { dimension, - target_k: 16, + target_k, metric_type: MetricType::L2, max_iters: 2, on_progress: Arc::new(|_, _| {}), @@ -6440,9 +6465,16 @@ mod tests { ¶ms, ) .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains(&format!("Cannot create {target_k} IVF partitions")), + "got: {msg}" + ); + // The row count belongs in the message: it is the coreset size, not + // the dataset's, so a reader can tell the two shortfalls apart. assert!( - err.to_string().contains("Cannot create 16 IVF partitions"), - "got: {err}" + msg.contains(&format!("{num_points} weighted coreset rows")), + "got: {msg}" ); } From 6488b60a9eeef09ffc5c2bea7c7ec21bae7b42e9 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 17 Sep 2026 19:21:27 +0800 Subject: [PATCH 03/10] fix(index): exhaust splittable clusters before rejecting --- rust/lance/src/index/vector/ivf.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 82f10dc27d8..8a91eb9a143 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4434,11 +4434,21 @@ fn train_weighted_hierarchical_f32_kmeans( let mut cluster = heap .pop() .ok_or_else(|| Error::index("No weighted cluster can be further split"))?; - if cluster.finalized || cluster.indices.len() <= 1 { - cluster.finalized = true; + if cluster.finalized { + // Ordering pops non-finalized clusters first, so a finalized + // cluster at the top means nothing splittable is left. heap.push(cluster); break; } + if cluster.indices.len() <= 1 { + // Priority is loss, not size, so a heavy singleton can outrank + // clusters that still have rows to give. Finalize it so it sinks + // below them and keep splitting the rest; giving up here would + // reject a build that can still reach target_k. + cluster.finalized = true; + heap.push(cluster); + continue; + } let remaining_k = target_k - heap.len(); let cluster_k = if cluster.indices.len() <= 16 { From 357e83d4d8a803db4d0c53243a6be02fce9d3db8 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 17 Sep 2026 19:57:58 +0800 Subject: [PATCH 04/10] test(index): guard against over-rejecting an attainable partition count --- rust/lance/src/index/vector/ivf.rs | 40 +++++++++++++++++++++++++++--- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 8a91eb9a143..0ab13382f4a 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4441,10 +4441,10 @@ fn train_weighted_hierarchical_f32_kmeans( break; } if cluster.indices.len() <= 1 { - // Priority is loss, not size, so a heavy singleton can outrank - // clusters that still have rows to give. Finalize it so it sinks - // below them and keep splitting the rest; giving up here would - // reject a build that can still reach target_k. + // Priority is loss, not size, so a singleton can outrank clusters + // that still have rows to give. Finalize it so it sinks below them + // and keep splitting the rest; giving up here would reject a build + // that can still reach target_k. cluster.finalized = true; heap.push(cluster); continue; @@ -6488,6 +6488,38 @@ mod tests { ); } + /// Splitting must not stop at the first unsplittable cluster while others + /// still have rows to give. With one more row than partitions, the loop + /// has to walk past singletons to find the last splittable cluster, so + /// abandoning the loop there loses a partition that was attainable. + #[test] + fn test_weighted_hierarchical_reaches_target_past_unsplittable_clusters() { + let dimension = 4; + let num_points = 258; + let target_k = 257; + let values: Vec = (0..num_points) + .flat_map(|row| std::iter::repeat_n(row as f32, dimension)) + .collect(); + let data = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension as i32) + .unwrap(); + let params = WeightedHierarchicalKMeansParams { + dimension, + target_k, + metric_type: MetricType::L2, + max_iters: 2, + on_progress: Arc::new(|_, _| {}), + }; + let centroids = train_weighted_hierarchical_f32_kmeans( + &data, + &vec![1.0; num_points], + &vec![0.1; num_points], + ¶ms, + ) + .expect("258 distinct rows can reach 257 partitions"); + assert_eq!(centroids.len(), target_k); + } + /// Regression test for a hang in the streaming *coreset* trainer /// (`train_streaming_coreset_ivf_model`, taken when `num_partitions > 256`). /// From 78b0d6c3d6873169bfd36dfbaf8b407bf67ec08b Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 17 Sep 2026 22:48:55 +0800 Subject: [PATCH 05/10] test(index): pin that a high-loss singleton does not stop splitting --- rust/lance/src/index/vector/ivf.rs | 49 +++++++++++++++--------------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 0ab13382f4a..2795bac8f87 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4441,10 +4441,11 @@ fn train_weighted_hierarchical_f32_kmeans( break; } if cluster.indices.len() <= 1 { - // Priority is loss, not size, so a singleton can outrank clusters - // that still have rows to give. Finalize it so it sinks below them - // and keep splitting the rest; giving up here would reject a build - // that can still reach target_k. + // Priority is loss, not size, and `assign_weighted_f32_points` + // folds each supplied base loss into the cluster loss that orders + // the heap, so one huge base loss puts a singleton on top. Finalize + // it so it sinks and keep splitting the rest; giving up here + // rejected counts that were attainable. cluster.finalized = true; heap.push(cluster); continue; @@ -6488,35 +6489,35 @@ mod tests { ); } - /// Splitting must not stop at the first unsplittable cluster while others - /// still have rows to give. With one more row than partitions, the loop - /// has to walk past singletons to find the last splittable cluster, so - /// abandoning the loop there loses a partition that was attainable. + /// A high-loss singleton must not stop training while other clusters still + /// have distinct rows to split. `assign_weighted_f32_points` folds each + /// supplied `base_losses[row]` into the cluster loss that orders the heap, + /// so one huge base loss puts a single-member cluster at the top; breaking + /// there rejected a 257-partition request that 258 distinct rows can serve, + /// advising a maximum of 16. #[test] - fn test_weighted_hierarchical_reaches_target_past_unsplittable_clusters() { - let dimension = 4; + fn test_weighted_hierarchical_exhausts_splittable_clusters() { + let dimension = 1; let num_points = 258; let target_k = 257; - let values: Vec = (0..num_points) - .flat_map(|row| std::iter::repeat_n(row as f32, dimension)) - .collect(); + let mut values = Vec::with_capacity(num_points); + values.push(1.0e9); + values.extend((1..num_points).map(|i| i as f32 * 100.0)); let data = - FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension as i32) - .unwrap(); + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap(); let params = WeightedHierarchicalKMeansParams { - dimension, + dimension: dimension as usize, target_k, metric_type: MetricType::L2, - max_iters: 2, + max_iters: 20, on_progress: Arc::new(|_, _| {}), }; - let centroids = train_weighted_hierarchical_f32_kmeans( - &data, - &vec![1.0; num_points], - &vec![0.1; num_points], - ¶ms, - ) - .expect("258 distinct rows can reach 257 partitions"); + let mut losses = vec![0.0; num_points]; + losses[0] = 1.0e20; + + let centroids = + train_weighted_hierarchical_f32_kmeans(&data, &vec![1.0; num_points], &losses, ¶ms) + .unwrap(); assert_eq!(centroids.len(), target_k); } From 52e960e871facdf9fe44a26314f64f3db36d49cf Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 16 Oct 2026 17:38:33 +0800 Subject: [PATCH 06/10] fix(index): bisect clusters k-means cannot split before rejecting a shortfall --- rust/lance/src/index/vector/ivf.rs | 140 ++++++++++++++++++++++++++++- 1 file changed, 137 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 333a0c65367..e69666fe76d 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4830,6 +4830,79 @@ fn weighted_subset( )) } +/// Splits `cluster` in two by L2 around its two farthest-apart rows, or returns +/// it unchanged when all its rows are identical. Each seed row is nearest to +/// itself, so both halves are non-empty whenever the rows differ. +fn bisect_weighted_cluster( + cluster: WeightedCluster, + data_values: &[f32], + weights: &[f64], + losses: &[f64], + dimension: usize, + metric_policy: &dyn StreamingFloatKMeansMetricPolicy, + next_cluster_id: &mut usize, +) -> std::result::Result<[WeightedCluster; 2], WeightedCluster> { + let row = |idx: usize| &data_values[idx * dimension..(idx + 1) * dimension]; + let farthest_from = |anchor: &[f32]| { + cluster + .indices + .iter() + .copied() + .max_by(|&a, &b| squared_l2(row(a), anchor).total_cmp(&squared_l2(row(b), anchor))) + }; + let Some(first) = farthest_from(&cluster.centroid) else { + return Err(cluster); + }; + let Some(second) = farthest_from(row(first)) else { + return Err(cluster); + }; + if squared_l2(row(first), row(second)) == 0.0 { + return Err(cluster); + } + let (near_first, near_second): (Vec, Vec) = cluster + .indices + .iter() + .partition(|&&idx| squared_l2(row(idx), row(first)) <= squared_l2(row(idx), row(second))); + Ok([near_first, near_second].map(|indices| { + let weight = indices.iter().map(|&idx| weights[idx]).sum::(); + let mut centroid = vec![0.0_f32; dimension]; + for &idx in &indices { + for (sum, value) in centroid.iter_mut().zip(row(idx)) { + *sum += *value * weights[idx] as f32; + } + } + if weight > 0.0 { + centroid + .iter_mut() + .for_each(|value| *value /= weight as f32); + } else { + centroid.copy_from_slice(row(indices[0])); + } + let loss = indices + .iter() + .map(|&idx| losses[idx] + weights[idx] * squared_l2(row(idx), ¢roid)) + .sum(); + let split_priority = metric_policy.split_priority( + data_values, + weights, + &indices, + ¢roid, + dimension, + loss, + ); + let id = *next_cluster_id; + *next_cluster_id += 1; + WeightedCluster { + id, + indices, + centroid, + weight, + split_priority, + finalized: true, + } + })) +} + fn train_weighted_hierarchical_f32_kmeans( data: &FixedSizeListArray, weights: &[f64], @@ -4981,16 +5054,42 @@ fn train_weighted_hierarchical_f32_kmeans( } } + // A k-means split can leave every row on one side even when the rows + // differ, which happens under Dot, and the cluster is then finalized whole. + // Bisect such clusters by L2 until `target_k` is reached, so the count + // falls short only when the coreset holds fewer distinct rows. let mut clusters = heap.into_vec(); + let mut unsplittable = Vec::new(); + while clusters.len() + unsplittable.len() < target_k { + let Some(position) = (0..clusters.len()) + .filter(|&position| clusters[position].indices.len() > 1) + .max_by(|&a, &b| clusters[a].cmp(&clusters[b])) + else { + break; + }; + let cluster = clusters.swap_remove(position); + match bisect_weighted_cluster( + cluster, + data_values, + weights, + losses, + dimension, + metric_policy, + &mut next_cluster_id, + ) { + Ok(halves) => clusters.extend(halves), + Err(cluster) => unsplittable.push(cluster), + } + } + clusters.append(&mut unsplittable); clusters.sort_by_key(|cluster| cluster.id); if clusters.len() < target_k { // Duplicating the heaviest cluster to pad the count used to leave // permanently empty partitions: a padded centroid is a bit-identical // clone, and assignment keeps the incumbent on a tie, so no vector is // ever assigned to a duplicate. Reject instead, like the flat - // hierarchical trainer. The shortfall has two causes worth - // distinguishing in a bug report but not in the advice: the coreset - // rows are too alike to split further, or splitting stopped early. + // hierarchical trainer: after the bisection above, the coreset has + // fewer distinct rows than `target_k`. return Err(Error::invalid_input(format!( "Cannot create {target_k} IVF partitions: weighted k-means could only form {} \ non-empty clusters from {} weighted coreset rows, and the rest could not be \ @@ -7001,6 +7100,41 @@ mod tests { assert_eq!(centroids.len(), target_k); } + /// Under Dot, positive 1-D rows all go to the largest centroid, so every + /// k-means split leaves its rows on one side. The rows are distinct, so the + /// trainer must still reach `target_k` distinct centroids rather than + /// reject the request. + #[test] + fn test_weighted_hierarchical_bisects_clusters_kmeans_cannot_split() { + let num_points = 8; + let target_k = 4; + let values = (1..=num_points).map(|i| i as f32).collect::>(); + let data = FixedSizeListArray::try_new_from_values(Float32Array::from(values), 1).unwrap(); + let params = WeightedHierarchicalKMeansParams { + dimension: 1, + target_k, + metric_policy: &DOT_METRIC_POLICY, + max_iters: 5, + on_progress: Arc::new(|_, _| {}), + }; + + let centroids = train_weighted_hierarchical_f32_kmeans( + &data, + &vec![1.0; num_points], + &vec![0.0; num_points], + ¶ms, + ) + .unwrap(); + let mut values = centroids + .values() + .as_primitive::() + .values() + .to_vec(); + values.sort_by(f32::total_cmp); + values.dedup(); + assert_eq!(values.len(), target_k, "centroids: {values:?}"); + } + #[tokio::test] async fn test_build_ivf_model_streaming_dot_coreset_training() { const DIMENSION: usize = 4; From 2a2220ebf6c4097a35cd8f824e5a86ca00a16eed Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 16 Oct 2026 21:42:57 +0800 Subject: [PATCH 07/10] fix(index): leave the singleton exit to the bisection and name a short coreset --- rust/lance/src/index/vector/ivf.rs | 96 ++++++++++-------------------- 1 file changed, 31 insertions(+), 65 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index e69666fe76d..f073ee5742c 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4830,9 +4830,10 @@ fn weighted_subset( )) } -/// Splits `cluster` in two by L2 around its two farthest-apart rows, or returns -/// it unchanged when all its rows are identical. Each seed row is nearest to -/// itself, so both halves are non-empty whenever the rows differ. +/// Splits `cluster` in two by L2 around the row farthest from its centroid and +/// the row farthest from that one, or returns it unchanged when all its rows are +/// identical. Each seed row is nearest to itself, so both halves are non-empty +/// whenever the rows differ. fn bisect_weighted_cluster( cluster: WeightedCluster, data_values: &[f32], @@ -4973,21 +4974,10 @@ fn train_weighted_hierarchical_f32_kmeans( let mut cluster = heap .pop() .ok_or_else(|| Error::index("No weighted cluster can be further split"))?; - if cluster.finalized { - // Ordering pops non-finalized clusters first, so a finalized - // cluster at the top means nothing splittable is left. - heap.push(cluster); - break; - } - if cluster.indices.len() <= 1 { - // Priority is not size. Under L2 it is the cluster loss, into which - // `assign_weighted_f32_points` folds each supplied base loss, so one - // huge base loss puts a singleton on top. Finalize it so it sinks - // and keep splitting the rest; giving up here rejected counts that - // were attainable. + if cluster.finalized || cluster.indices.len() <= 1 { cluster.finalized = true; heap.push(cluster); - continue; + break; } let remaining_k = target_k - heap.len(); @@ -5084,12 +5074,22 @@ fn train_weighted_hierarchical_f32_kmeans( clusters.append(&mut unsplittable); clusters.sort_by_key(|cluster| cluster.id); if clusters.len() < target_k { - // Duplicating the heaviest cluster to pad the count used to leave - // permanently empty partitions: a padded centroid is a bit-identical - // clone, and assignment keeps the incumbent on a tie, so no vector is - // ever assigned to a duplicate. Reject instead, like the flat - // hierarchical trainer: after the bisection above, the coreset has - // fewer distinct rows than `target_k`. + // Padding the count with clones of the heaviest cluster left partitions + // whose centroid ties exactly with the original, which assignment keeps, + // so they started empty. Reject instead, like the flat hierarchical + // trainer: after the bisection above, the coreset has fewer distinct + // rows than `target_k`. + if data.len() < target_k { + // Not the data's fault: the coreset summarizing the training sample + // is smaller than the partition count, as Dot's can be when each + // chunk samples exactly `num_partitions` rows. + return Err(Error::invalid_input(format!( + "Cannot create {target_k} IVF partitions from {} weighted coreset rows: the \ + coreset summarizing the training sample is smaller than the partition \ + count. Increase sample_rate or reduce num_partitions.", + data.len() + ))); + } return Err(Error::invalid_input(format!( "Cannot create {target_k} IVF partitions: weighted k-means could only form {} \ non-empty clusters from {} weighted coreset rows, and the rest could not be \ @@ -7011,22 +7011,19 @@ mod tests { ); } - /// Duplicating the heaviest cluster to pad the partition count used to - /// yield permanently empty partitions (a padded centroid is a clone, and - /// assignment keeps the incumbent on a tie); reject like the flat - /// hierarchical trainer. - /// - /// Two shapes reach the shortfall. Fewer rows than partitions is the - /// obvious one; the one the only production caller can actually produce is - /// a coreset with enough rows whose clusters cannot be split further, - /// because the coreset always carries at least `num_partitions` rows. + /// Padding the partition count with clones of the heaviest cluster left + /// partitions that started empty; reject like the flat hierarchical + /// trainer. A coreset smaller than the partition count and a coreset whose + /// rows are identical both reach the shortfall, and each gets its own + /// advice. #[rstest::rstest] - #[case::fewer_rows_than_partitions(6, 16)] - #[case::enough_rows_all_identical(32, 16)] + #[case::fewer_rows_than_partitions(6, 16, "Increase sample_rate")] + #[case::enough_rows_all_identical(32, 16, "provide more diverse data")] #[test] fn test_weighted_hierarchical_rejects_padding_with_duplicates( #[case] num_points: usize, #[case] target_k: usize, + #[case] advice: &str, ) { let dimension = 8; // Every row identical in the second case, a ramp in the first: either @@ -7066,38 +7063,7 @@ mod tests { msg.contains(&format!("{num_points} weighted coreset rows")), "got: {msg}" ); - } - - /// A high-loss singleton must not stop training while other clusters still - /// have distinct rows to split. `assign_weighted_f32_points` folds each - /// supplied `base_losses[row]` into the cluster loss that orders the heap, - /// so one huge base loss puts a single-member cluster at the top; breaking - /// there rejected a 257-partition request that 258 distinct rows can serve, - /// advising a maximum of 16. - #[test] - fn test_weighted_hierarchical_exhausts_splittable_clusters() { - let dimension = 1; - let num_points = 258; - let target_k = 257; - let mut values = Vec::with_capacity(num_points); - values.push(1.0e9); - values.extend((1..num_points).map(|i| i as f32 * 100.0)); - let data = - FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap(); - let params = WeightedHierarchicalKMeansParams { - dimension: dimension as usize, - target_k, - metric_policy: &L2_METRIC_POLICY, - max_iters: 20, - on_progress: Arc::new(|_, _| {}), - }; - let mut losses = vec![0.0; num_points]; - losses[0] = 1.0e20; - - let centroids = - train_weighted_hierarchical_f32_kmeans(&data, &vec![1.0; num_points], &losses, ¶ms) - .unwrap(); - assert_eq!(centroids.len(), target_k); + assert!(msg.contains(advice), "got: {msg}"); } /// Under Dot, positive 1-D rows all go to the largest centroid, so every From f14a8147cdc3ef2dc299d613accc10d7e3aa6215 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 16 Oct 2026 22:20:13 +0800 Subject: [PATCH 08/10] docs(index): say which clusters the bisection reaches --- rust/lance/src/index/vector/ivf.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index f073ee5742c..670a7653954 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -5044,10 +5044,11 @@ fn train_weighted_hierarchical_f32_kmeans( } } - // A k-means split can leave every row on one side even when the rows - // differ, which happens under Dot, and the cluster is then finalized whole. - // Bisect such clusters by L2 until `target_k` is reached, so the count - // falls short only when the coreset holds fewer distinct rows. + // The split loop can stop with multi-row clusters left: a k-means split can + // leave every row on one side even when the rows differ, which happens + // under Dot, and a singleton at the top of the heap ends the loop. Bisect + // the remaining multi-row clusters by L2 until `target_k` is reached, so + // the count falls short only when the coreset holds fewer distinct rows. let mut clusters = heap.into_vec(); let mut unsplittable = Vec::new(); while clusters.len() + unsplittable.len() < target_k { @@ -5081,7 +5082,7 @@ fn train_weighted_hierarchical_f32_kmeans( // rows than `target_k`. if data.len() < target_k { // Not the data's fault: the coreset summarizing the training sample - // is smaller than the partition count, as Dot's can be when each + // is smaller than the partition count, as Dot's can be when its only // chunk samples exactly `num_partitions` rows. return Err(Error::invalid_input(format!( "Cannot create {target_k} IVF partitions from {} weighted coreset rows: the \ From 0c6cedfc41438d5b73a9d43d5c92506fc5e913e0 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 22 Oct 2026 14:38:21 +0800 Subject: [PATCH 09/10] refactor(index): build bisected halves without expect; split the f32/f64 bisect test --- rust/lance/src/index/vector/ivf.rs | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index ea739ed8ab7..15965be404b 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4895,8 +4895,7 @@ fn bisect_weighted_cluster( .indices .iter() .partition(|&&idx| squared_l2(row(idx), row(first)) <= squared_l2(row(idx), row(second))); - let mut halves = Vec::with_capacity(2); - for indices in [near_first, near_second] { + let mut make_half = |indices: Vec| -> Result> { let weight = indices.iter().map(|&idx| weights[idx]).sum::(); let mut centroid = vec![T::Native::default(); dimension]; for &idx in &indices { @@ -4927,17 +4926,17 @@ fn bisect_weighted_cluster( ); let id = *next_cluster_id; *next_cluster_id += 1; - halves.push(WeightedCluster { + Ok(WeightedCluster { id, indices, centroid, weight, split_priority, finalized: true, - }); - } - let second = halves.pop().expect("two halves"); - let first = halves.pop().expect("two halves"); + }) + }; + let first = make_half(near_first)?; + let second = make_half(near_second)?; Ok(Some([first, second])) } @@ -7172,10 +7171,11 @@ mod tests { assert_eq!(values.len(), target_k, "centroids: {values:?}"); } - #[test] - fn test_weighted_hierarchical_bisects_clusters_kmeans_cannot_split() { - bisects_clusters_kmeans_cannot_split::(); - bisects_clusters_kmeans_cannot_split::(); + #[rstest::rstest] + #[case::f32(bisects_clusters_kmeans_cannot_split:: as fn())] + #[case::f64(bisects_clusters_kmeans_cannot_split:: as fn())] + fn test_weighted_hierarchical_bisects_clusters_kmeans_cannot_split(#[case] check: fn()) { + check(); } #[tokio::test] From 3c70e999983905b813ef60b54d37b6ef61e59aec Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 08:03:34 +0800 Subject: [PATCH 10/10] test(index): expect rejection when coreset cannot serve target_k --- rust/lance/src/index/vector/ivf.rs | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 493fb3cb519..0c471697297 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -8532,8 +8532,23 @@ mod tests { &vec![1.0; 32], &losses, ¶ms, - ) - .unwrap(); + ); + if target_k > data.len() { + let Err(error) = result else { + panic!( + "a coreset of {} rows cannot serve {target_k} partitions", + data.len() + ); + }; + assert!( + error + .to_string() + .contains("smaller than the partition count"), + "unexpected error: {error}" + ); + return; + } + let result = result.unwrap(); assert_eq!(result.len(), target_k); let values = result.values().as_primitive::().values(); let mut sorted = values.to_vec();