From 037deecb147c469dbcbb34fd5b07d270b090ff81 Mon Sep 17 00:00:00 2001 From: Daoyuan Wang Date: Mon, 28 Sep 2026 16:31:51 +0800 Subject: [PATCH] fix: skip NaN scores in IVF top-k heaps Reject both signs of NaN in FlatTopKHeap, ReaderTopKHeap, IVF-PQ TopKHeap, and the unified TopKHeap before insertion or duplicate updates. Use total_cmp for consistent heap comparisons and sorting, and reject NaN candidates in should_consider. Add regression coverage for both NaN signs, heap capacity, duplicate IDs, and IVF-FLAT in-memory and reader inner-product searches, including padding when valid neighbors are fewer than k. Validation: Rust workspace tests, formatting, and Clippy all pass. --- core/src/ivfflat.rs | 71 +++++++++++++++++++++++++++++++++++++++--- core/src/ivfflat_io.rs | 51 +++++++++++++++++++++++++++--- core/src/ivfpq.rs | 37 +++++++++++++++++++--- core/src/topk.rs | 53 ++++++++++++++++++++++++++++--- 4 files changed, 195 insertions(+), 17 deletions(-) diff --git a/core/src/ivfflat.rs b/core/src/ivfflat.rs index 67251d2..46d0d8e 100644 --- a/core/src/ivfflat.rs +++ b/core/src/ivfflat.rs @@ -194,7 +194,7 @@ impl FlatTopKHeap { } fn push(&mut self, dist: f32, id: i64) { - if self.k == 0 { + if self.k == 0 || dist.is_nan() { return; } if self.data.len() < self.k { @@ -205,16 +205,16 @@ impl FlatTopKHeap { .data .iter() .enumerate() - .max_by(|(_, a), (_, b)| a.0.partial_cmp(&b.0).unwrap()) + .max_by(|(_, a), (_, b)| a.0.total_cmp(&b.0)) { - if dist < self.data[worst_idx].0 { + if dist.total_cmp(&self.data[worst_idx].0).is_lt() { self.data[worst_idx] = (dist, id); } } } fn into_sorted(mut self) -> Vec<(f32, i64)> { - self.data.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + self.data.sort_by(|a, b| a.0.total_cmp(&b.0)); self.data } } @@ -310,4 +310,67 @@ mod tests { assert!(distances[0] > 0.0); assert_eq!(distances[1], f32::MAX); } + + #[test] + fn test_flat_topk_heap_skips_both_nan_signs() { + let expected = [(0.25, 5), (0.5, 3), (1.0, 1), (2.0, 4)]; + for k in [0, 1, 3, 5] { + let mut heap = FlatTopKHeap::new(k); + for (distance, id) in [ + (f32::NAN, 10), + (-f32::NAN, 11), + (1.0, 1), + (f32::NAN, 12), + (0.5, 3), + (-f32::NAN, 13), + (2.0, 4), + (f32::NAN, 14), + (-f32::NAN, 15), + (0.25, 5), + ] { + heap.push(distance, id); + } + assert_eq!( + heap.into_sorted(), + expected[..k.min(expected.len())], + "k={k}" + ); + } + } + + #[test] + fn test_ivfflat_inner_product_search_skips_both_nan_signs() { + for nan in [f32::NAN, -f32::NAN] { + for (data, ids) in [ + ([nan, 1.0, 2.0], [99, 1, 2]), + ([1.0, nan, 2.0], [1, 99, 2]), + ([1.0, 2.0, nan], [1, 2, 99]), + ] { + let mut index = IVFFlatIndex::new(1, 1, MetricType::InnerProduct); + index.set_quantizer_centroids(vec![0.0]); + index.add(&data, &ids, 3); + for k in 1..=3 { + let mut distances = vec![0.0; k]; + let mut labels = vec![0; k]; + index.search(&[1.0], 1, k, 1, &mut distances, &mut labels); + assert_eq!(labels, [2, 1, -1][..k]); + assert_eq!(distances, [-2.0, -1.0, f32::MAX][..k]); + } + } + } + } + + #[test] + fn test_ivfflat_search_with_only_nan_returns_missing_neighbors() { + for metric in [MetricType::L2, MetricType::InnerProduct] { + let mut index = IVFFlatIndex::new(1, 1, metric); + index.set_quantizer_centroids(vec![0.0]); + index.add(&[f32::NAN, -f32::NAN], &[1, 2], 2); + let mut distances = [0.0; 3]; + let mut labels = [0; 3]; + index.search(&[1.0], 1, 3, 1, &mut distances, &mut labels); + assert_eq!(labels, [-1; 3]); + assert_eq!(distances, [f32::MAX; 3]); + } + } } diff --git a/core/src/ivfflat_io.rs b/core/src/ivfflat_io.rs index cadf1d8..0f11ceb 100644 --- a/core/src/ivfflat_io.rs +++ b/core/src/ivfflat_io.rs @@ -1410,7 +1410,7 @@ impl ReaderTopKHeap { #[inline] fn push(&mut self, dist: f32, id: i64) { - if self.k == 0 { + if self.k == 0 || dist.is_nan() { return; } if self.data.len() < self.k { @@ -1423,7 +1423,7 @@ impl ReaderTopKHeap { let worst_index = self .worst_index .expect("a full IVF-FLAT top-k heap has a worst entry"); - if dist < self.data[worst_index].0 { + if dist.total_cmp(&self.data[worst_index].0).is_lt() { self.data[worst_index] = (dist, id); self.refresh_worst(); } @@ -1434,12 +1434,12 @@ impl ReaderTopKHeap { .data .iter() .enumerate() - .max_by(|(_, a), (_, b)| a.0.partial_cmp(&b.0).unwrap()) + .max_by(|(_, a), (_, b)| a.0.total_cmp(&b.0)) .map(|(index, _)| index); } fn into_sorted(mut self) -> Vec<(f32, i64)> { - self.data.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + self.data.sort_by(|a, b| a.0.total_cmp(&b.0)); self.data } } @@ -2396,4 +2396,47 @@ mod tests { }; assert!(err.to_string().contains("reserved bytes must be zero")); } + + #[test] + fn test_reader_topk_heap_skips_both_nan_signs() { + let expected = [(0.25, 5), (0.5, 3), (1.0, 1), (2.0, 4)]; + for k in [0, 1, 3, 5] { + let mut heap = ReaderTopKHeap::new(k); + for (distance, id) in [ + (f32::NAN, 10), + (-f32::NAN, 11), + (1.0, 1), + (f32::NAN, 12), + (0.5, 3), + (-f32::NAN, 13), + (2.0, 4), + (f32::NAN, 14), + (-f32::NAN, 15), + (0.25, 5), + ] { + heap.push(distance, id); + } + assert_eq!( + heap.into_sorted(), + expected[..k.min(expected.len())], + "k={k}" + ); + } + } + + #[test] + fn test_ivfflat_reader_inner_product_search_skips_both_nan_signs() { + for nan in [f32::NAN, -f32::NAN] { + let mut index = IVFFlatIndex::new(1, 1, MetricType::InnerProduct); + index.set_quantizer_centroids(vec![0.0]); + index.add(&[1.0, nan, 2.0], &[1, 2, 3], 3); + let mut reader = + IVFFlatIndexReader::open(Cursor::new(serialized_flat_index(&index))).unwrap(); + for k in 1..=3 { + let (labels, distances) = reader.search(&[1.0], k, 1).unwrap(); + assert_eq!(labels, [3, 1, -1][..k]); + assert_eq!(distances, [-2.0, -1.0, f32::MAX][..k]); + } + } + } } diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 000bfa4..4bdb3b1 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -3322,7 +3322,7 @@ impl TopKHeap { #[inline] fn push(&mut self, dist: f32, id: i64) { - if self.k == 0 { + if self.k == 0 || dist.is_nan() { return; } if self.data.len() < self.k { @@ -3331,14 +3331,14 @@ impl TopKHeap { build_max_heap(&mut self.data); self.built = true; } - } else if dist < self.data[0].0 { + } else if dist.total_cmp(&self.data[0].0).is_lt() { self.data[0] = (dist, id); sift_down(&mut self.data, 0); } } fn into_sorted(mut self) -> Vec<(f32, i64)> { - self.data.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + self.data.sort_by(|a, b| a.0.total_cmp(&b.0)); self.data } } @@ -3419,10 +3419,10 @@ fn sift_down(heap: &mut [(f32, i64)], mut i: usize) { let left = 2 * i + 1; let right = 2 * i + 2; - if left < n && heap[left].0 > heap[largest].0 { + if left < n && heap[left].0.total_cmp(&heap[largest].0).is_gt() { largest = left; } - if right < n && heap[right].0 > heap[largest].0 { + if right < n && heap[right].0.total_cmp(&heap[largest].0).is_gt() { largest = right; } if largest == i { @@ -6246,4 +6246,31 @@ mod tests { fn ivfpq_rejects_odd_4bit_subquantizer_count_at_construction() { let _ = IVFPQIndex::with_nbits(12, 4, 3, 4, MetricType::L2, false); } + + #[test] + fn test_topk_heap_skips_both_nan_signs() { + let expected = [(0.25, 5), (0.5, 3), (1.0, 1), (2.0, 4)]; + for k in [0, 1, 3, 5] { + let mut heap = TopKHeap::new(k); + for (distance, id) in [ + (f32::NAN, 10), + (-f32::NAN, 11), + (1.0, 1), + (f32::NAN, 12), + (0.5, 3), + (-f32::NAN, 13), + (2.0, 4), + (f32::NAN, 14), + (-f32::NAN, 15), + (0.25, 5), + ] { + heap.push(distance, id); + } + assert_eq!( + heap.into_sorted(), + expected[..k.min(expected.len())], + "k={k}" + ); + } + } } diff --git a/core/src/topk.rs b/core/src/topk.rs index 499f210..0df90d1 100644 --- a/core/src/topk.rs +++ b/core/src/topk.rs @@ -46,11 +46,11 @@ impl TopKHeap { } pub(crate) fn push(&mut self, dist: f32, id: i64) { - if self.k == 0 || dist >= self.max_distance { + if self.k == 0 || dist.is_nan() || dist >= self.max_distance { return; } if let Some(&idx) = self.positions.get(&id) { - if dist < self.data[idx].0 { + if dist.total_cmp(&self.data[idx].0).is_lt() { self.data[idx].0 = dist; self.sift_down(idx); } @@ -62,7 +62,7 @@ impl TopKHeap { self.sift_up(self.data.len() - 1); return; } - if dist < self.data[0].0 { + if dist.total_cmp(&self.data[0].0).is_lt() { let old_id = self.data[0].1; self.positions.remove(&old_id); self.data[0] = (dist, id); @@ -73,10 +73,11 @@ impl TopKHeap { pub(crate) fn should_consider(&self, lower_bound: f32) -> bool { self.k > 0 + && !lower_bound.is_nan() && if self.data.len() < self.k { lower_bound < self.max_distance } else { - lower_bound < self.data[0].0 + lower_bound.total_cmp(&self.data[0].0).is_lt() } } @@ -190,4 +191,48 @@ mod tests { assert_eq!(heap.into_sorted(), vec![(3.0, 3), (4.0, 2)]); } + + #[test] + fn test_topk_heap_skips_both_nan_signs() { + let expected = [(0.25, 5), (0.5, 3), (1.0, 1), (2.0, 4)]; + for k in [0, 1, 3, 5] { + let mut heap = TopKHeap::new(k); + for (distance, id) in [ + (f32::NAN, 10), + (-f32::NAN, 11), + (1.0, 1), + (f32::NAN, 12), + (0.5, 3), + (-f32::NAN, 13), + (2.0, 4), + (f32::NAN, 14), + (-f32::NAN, 15), + (0.25, 5), + ] { + heap.push(distance, id); + } + assert_eq!( + heap.into_sorted(), + expected[..k.min(expected.len())], + "k={k}" + ); + } + } + + #[test] + fn test_topk_heap_rejects_nan_for_duplicates_and_collector() { + for nan in [f32::NAN, -f32::NAN] { + let mut heap = TopKHeap::new(2); + assert!(!heap.should_consider(nan)); + heap.push(6.0, 2); + assert!(!heap.should_consider(nan)); + heap.push(8.0, 1); + assert!(!heap.should_consider(nan)); + heap.push(nan, 2); + Collector::push(&mut heap, 3, nan).unwrap(); + heap.push(4.0, 1); + heap.push(5.0, 3); + assert_eq!(heap.into_sorted(), vec![(4.0, 1), (5.0, 3)]); + } + } }