diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index d3b1f8bda97..0c471697297 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4993,6 +4993,85 @@ fn weighted_subset( )) } +/// Splits `cluster` in two by L2 around the row farthest from its centroid and +/// the row farthest from that one, or returns `None` 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: &[T::Native], + weights: &[f64], + losses: &[f64], + dimension: usize, + metric_policy: &dyn StreamingFloatKMeansMetricPolicy, + next_cluster_id: &mut usize, +) -> Result; 2]>> { + let row = |idx: usize| &data_values[idx * dimension..(idx + 1) * dimension]; + let farthest_from = |anchor: &[T::Native]| { + 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 Ok(None); + }; + let Some(second) = farthest_from(row(first)) else { + return Ok(None); + }; + if squared_l2(row(first), row(second)) == 0.0 { + return Ok(None); + } + let (near_first, near_second): (Vec, Vec) = cluster + .indices + .iter() + .partition(|&&idx| squared_l2(row(idx), row(first)) <= squared_l2(row(idx), row(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 { + let row_weight = cast_weight::(weights[idx], format_args!("bisected row {idx}"))?; + for (sum, value) in centroid.iter_mut().zip(row(idx)) { + *sum += *value * row_weight; + } + } + if weight > 0.0 { + let weight = cast_weight::(weight, "bisected cluster")?; + centroid + .iter_mut() + .for_each(|value| *value = *value / weight); + } 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; + Ok(WeightedCluster { + id, + indices, + centroid, + weight, + split_priority, + finalized: true, + }) + }; + let first = make_half(near_first)?; + let second = make_half(near_second)?; + Ok(Some([first, second])) +} + fn train_weighted_hierarchical_kmeans( data: &FixedSizeListArray, weights: &[f64], @@ -5139,25 +5218,70 @@ fn train_weighted_hierarchical_kmeans( } } + // 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 { + 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, + )? { + Some(halves) => clusters.extend(halves), + None => unsplittable.push(cluster), + } + } + clusters.append(&mut unsplittable); 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 { + // 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 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 \ + 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 \ + 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 { @@ -7377,6 +7501,105 @@ mod tests { ); } + /// 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, "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 + // 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 = streaming_kmeans_fsl_from_values::(values, dimension).unwrap(); + let params = WeightedHierarchicalKMeansParams { + dimension, + target_k, + metric_policy: &L2_METRIC_POLICY, + max_iters: 2, + on_progress: Arc::new(|_, _| {}), + }; + let err = train_weighted_hierarchical_kmeans::( + &data, + &vec![1.0; num_points], + &vec![0.1; num_points], + ¶ms, + ) + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got: {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!( + msg.contains(&format!("{num_points} weighted coreset rows")), + "got: {msg}" + ); + assert!(msg.contains(advice), "got: {msg}"); + } + + /// 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, for f32 and f64 centroids alike. + fn bisects_clusters_kmeans_cannot_split() { + let num_points = 8; + let target_k = 4; + let values = (1..=num_points) + .map(|i| T::Native::from_f64(i as f64).unwrap()) + .collect::>(); + let data = streaming_kmeans_fsl_from_values::(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_kmeans::( + &data, + &vec![1.0; num_points], + &vec![0.0; num_points], + ¶ms, + ) + .unwrap(); + let mut values = centroids + .values() + .as_primitive::() + .values() + .iter() + .map(|value| value.as_()) + .collect::>(); + values.sort_by(f64::total_cmp); + values.dedup(); + assert_eq!(values.len(), target_k, "centroids: {values:?}"); + } + + #[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] async fn test_build_ivf_model_streaming_dot_coreset_training() { const DIMENSION: usize = 4; @@ -8309,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();