diff --git a/crates/paimon/src/lumina/reader.rs b/crates/paimon/src/lumina/reader.rs index 9c365f4cd..25514a67f 100644 --- a/crates/paimon/src/lumina/reader.rs +++ b/crates/paimon/src/lumina/reader.rs @@ -45,6 +45,39 @@ fn convert_distance_to_score(distance: f32, metric: LuminaVectorMetric) -> f32 { } } +/// Order two search scores, best last, with NaN ranked below every real score. +/// +/// This is the score-domain mirror of `vindex::pkvector::metric`'s +/// `java_float_compare`, which ranks a NaN *distance* worst. `f32::total_cmp` +/// alone is unsuitable in either domain: it places a positive NaN above every +/// finite value, so a NaN score -- reachable here because a non-finite stored +/// vector yields a NaN distance and [`convert_distance_to_score`] passes NaN +/// through -- would outrank real neighbours. Both NaN signs lose, and two NaNs +/// compare equal so the caller's row-id tie-break decides between them. +fn compare_scores(a: f32, b: f32) -> std::cmp::Ordering { + match (a.is_nan(), b.is_nan()) { + (true, true) => std::cmp::Ordering::Equal, + (true, false) => std::cmp::Ordering::Less, + (false, true) => std::cmp::Ordering::Greater, + (false, false) => a.total_cmp(&b), + } +} + +/// Allocate the label buffer for one native search, filled with [`SENTINEL`]. +/// +/// [`SENTINEL`] is the "no result" marker (the C ABI's `-1`), and +/// [`collect_results`] drops any slot carrying it. Zero-filling the buffer +/// instead makes that marker unreliable: `0` is a *legal* row id, so a slot the +/// searcher leaves untouched is indistinguishable from a real hit, and it pairs +/// with distance `0.0` -- the best distance both L2 and cosine can report -- so +/// a search returning fewer neighbours than requested would surface row 0 as its +/// top match. The FFI reports only a status code, never how many slots it wrote, +/// so the Rust side cannot detect a short return; allocating the sentinel is +/// correct either way. +fn new_label_buffer(len: usize) -> Vec { + vec![SENTINEL; len] +} + /// Post-filter search results to top_k. fn collect_results( labels: &[u64], @@ -64,8 +97,21 @@ fn collect_results( } } impl Ord for ScoredRow { + // Reversed on score so the heap top is the weakest candidate; among + // equal scores the larger row id sorts first and is therefore evicted + // first, which keeps the retained set independent of the order the + // searcher returned the pairs in. Same shape as + // `vector_search::ScoredRow`, except that scores are compared through + // `compare_scores` so a NaN cannot claim the strongest slot. fn cmp(&self, other: &Self) -> std::cmp::Ordering { - other.score.total_cmp(&self.score) + compare_scores(other.score, self.score).then_with(|| self.row_id.cmp(&other.row_id)) + } + } + + impl ScoredRow { + fn is_stronger_than(&self, other: &Self) -> bool { + compare_scores(self.score, other.score).then_with(|| other.row_id.cmp(&self.row_id)) + == std::cmp::Ordering::Greater } } @@ -75,13 +121,15 @@ fn collect_results( continue; } let score = convert_distance_to_score(distance, metric); + let entry = ScoredRow { row_id, score }; if min_heap.len() < top_k { - min_heap.push(ScoredRow { row_id, score }); - } else if let Some(peek) = min_heap.peek() { - if score > peek.score { - min_heap.pop(); - min_heap.push(ScoredRow { row_id, score }); - } + min_heap.push(entry); + } else if min_heap + .peek() + .is_some_and(|weakest| entry.is_stronger_than(weakest)) + { + min_heap.pop(); + min_heap.push(entry); } } @@ -272,7 +320,7 @@ fn search_lumina( } let ek = std::cmp::min(effective_k, filter_id_list.len()); let mut distances = vec![0.0f32; ek]; - let mut labels = vec![0u64; ek]; + let mut labels = new_label_buffer(ek); let mut search_opts: HashMap = search_options_base.clone(); search_opts.insert("search.thread_safe_filter".to_string(), "true".to_string()); ensure_search_list_size(&mut search_opts, ek); @@ -288,7 +336,7 @@ fn search_lumina( (distances, labels) } else { let mut distances = vec![0.0f32; effective_k]; - let mut labels = vec![0u64; effective_k]; + let mut labels = new_label_buffer(effective_k); let mut search_opts: HashMap = search_options_base.clone(); ensure_search_list_size(&mut search_opts, effective_k); searcher.search( @@ -371,7 +419,7 @@ fn search_lumina_batch( } let mut distances = vec![0.0f32; vector_searches.len() * effective_k]; - let mut labels = vec![0u64; vector_searches.len() * effective_k]; + let mut labels = new_label_buffer(vector_searches.len() * effective_k); let mut search_opts: HashMap = search_options_base.clone(); ensure_search_list_size(&mut search_opts, effective_k); searcher.search( @@ -500,6 +548,128 @@ mod tests { assert!(!result.contains_key(&2)); // 0.1 is lowest } + /// Rows sharing a score must be kept by ascending row id, not by the order + /// the native searcher happened to return them in. Feeding the same tied + /// scores in two different label orders must select the same rows. + #[test] + fn test_collect_results_breaks_ties_by_row_id() { + // Rows 10, 20, 30 all score 0.5; row 40 scores higher and always wins. + let distances = vec![0.9, 0.5, 0.5, 0.5]; + + let forward = collect_results( + &[40, 10, 20, 30], + &distances, + 2, + LuminaVectorMetric::InnerProduct, + ); + let reversed = collect_results( + &[40, 30, 20, 10], + &distances, + 2, + LuminaVectorMetric::InnerProduct, + ); + + let mut forward_ids: Vec = forward.keys().copied().collect(); + forward_ids.sort_unstable(); + let mut reversed_ids: Vec = reversed.keys().copied().collect(); + reversed_ids.sort_unstable(); + + assert_eq!( + forward_ids, reversed_ids, + "tied rows must not depend on label order" + ); + assert_eq!( + forward_ids, + vec![10, 40], + "among equal scores the smallest row id wins" + ); + } + + /// The same invariant when every candidate ties: the retained set is the + /// `top_k` smallest row ids regardless of input order. + #[test] + fn test_collect_results_all_tied_keeps_smallest_row_ids() { + let distances = vec![0.25; 5]; + + let forward = collect_results( + &[1, 2, 3, 4, 5], + &distances, + 3, + LuminaVectorMetric::InnerProduct, + ); + let shuffled = collect_results( + &[4, 1, 5, 3, 2], + &distances, + 3, + LuminaVectorMetric::InnerProduct, + ); + + let mut forward_ids: Vec = forward.keys().copied().collect(); + forward_ids.sort_unstable(); + let mut shuffled_ids: Vec = shuffled.keys().copied().collect(); + shuffled_ids.sort_unstable(); + + assert_eq!(forward_ids, vec![1, 2, 3]); + assert_eq!(shuffled_ids, vec![1, 2, 3]); + } + + /// A NaN score must never outrank a finite one. `f32::total_cmp` alone ranks + /// a positive NaN above every finite value, so using it here would let a NaN + /// score -- which a non-finite stored vector can produce -- take the only + /// top-1 slot. Both NaN signs must lose, in either arrival order. + #[test] + fn test_collect_results_ranks_nan_below_finite_scores() { + for (labels, distances) in [ + (vec![7u64, 8], vec![f32::NAN, 0.5]), + (vec![8u64, 7], vec![0.5, f32::NAN]), + (vec![7u64, 8], vec![-f32::NAN, 0.5]), + (vec![8u64, 7], vec![0.5, -f32::NAN]), + ] { + let result = collect_results(&labels, &distances, 1, LuminaVectorMetric::InnerProduct); + assert_eq!(result.len(), 1); + assert!( + result.contains_key(&8), + "the finite score must win regardless of arrival order, got {result:?}" + ); + } + } + + /// The label buffer handed to the native searcher must start out as + /// [`SENTINEL`], never zero -- row id `0` is a legal result. + #[test] + fn test_new_label_buffer_is_sentinel_filled() { + assert_eq!(new_label_buffer(3), vec![SENTINEL; 3]); + assert!(new_label_buffer(0).is_empty()); + } + + /// A search that returns fewer neighbours than requested leaves the tail of + /// the buffer exactly as it was allocated. Those slots must not surface as + /// rows: with a zero-filled buffer the tail reads as row `0` at distance + /// `0.0`, and `1.0 / (1.0 + 0.0)` is the highest score L2 can produce, so + /// row 0 would be reported as the best match for every short result. + #[test] + fn test_unfilled_label_slots_are_not_reported_as_hits() { + let mut labels = new_label_buffer(4); + let mut distances = vec![0.0f32; 4]; + labels[0] = 11; + distances[0] = 3.0; + labels[1] = 22; + distances[1] = 7.0; + + let result = collect_results(&labels, &distances, 4, LuminaVectorMetric::L2); + + assert_eq!( + result.len(), + 2, + "only the slots the searcher filled may be reported, got {result:?}" + ); + assert!(result.contains_key(&11) && result.contains_key(&22)); + assert!( + !result.contains_key(&0), + "an untouched slot must not become row 0" + ); + } + #[test] fn test_reader_new() { let m = GlobalIndexIOMeta::new("a".into(), 100, vec![]);