diff --git a/crates/paimon/src/deletion_vector/core.rs b/crates/paimon/src/deletion_vector/core.rs index 11c2d06c..4164da64 100644 --- a/crates/paimon/src/deletion_vector/core.rs +++ b/crates/paimon/src/deletion_vector/core.rs @@ -59,6 +59,15 @@ impl DeletionVector { self.bitmap.len() } + /// Returns true if `position` is deleted. Positions above `u32::MAX` cannot be + /// present in a roaring32 bitmap and are therefore reported as not deleted. + /// Mirrors Java `BitmapDeletionVector#isDeleted` / the searchers' `LongPredicate`. + pub fn is_deleted(&self, position: u64) -> bool { + u32::try_from(position) + .ok() + .is_some_and(|p| self.bitmap.contains(p)) + } + /// Returns an iterator over deleted positions that supports [DeletionVectorIterator::advance_to]. /// Required for efficient row selection building when skipping row groups (avoid re-scanning /// deletes in skipped ranges). @@ -249,4 +258,15 @@ mod tests { let expected_bitmap = RoaringBitmap::from_iter([1u32, 2u32]); assert_eq!(dv.bitmap(), &expected_bitmap, "bitmap should be [1, 2]"); } + + #[test] + fn test_is_deleted_reports_membership_and_guards_u32_overflow() { + let mut bitmap = RoaringBitmap::new(); + bitmap.insert(2); + let dv = DeletionVector::from_bitmap(bitmap); + assert!(dv.is_deleted(2), "position 2 was deleted"); + assert!(!dv.is_deleted(0), "position 0 was not deleted"); + // Positions above u32::MAX cannot exist in a roaring32 bitmap -> not deleted. + assert!(!dv.is_deleted(u64::from(u32::MAX) + 1)); + } } diff --git a/crates/paimon/src/spec/pk_vector_source.rs b/crates/paimon/src/spec/pk_vector_source.rs index 2685f45e..2c859e0e 100644 --- a/crates/paimon/src/spec/pk_vector_source.rs +++ b/crates/paimon/src/spec/pk_vector_source.rs @@ -157,8 +157,8 @@ impl PkVectorSourceMeta { } /// Minimal big-endian reader mirroring the Java `DataInput` primitives the -/// `_SOURCE_META` frame uses. Module-private by design (see the PR1 spec: no -/// shared `common/` abstraction until a second consumer exists). +/// `_SOURCE_META` frame uses. Module-private by design: no shared `common/` +/// abstraction until a second consumer exists. struct DataInputCursor<'a> { bytes: &'a [u8], position: usize, diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs index c0d566ac..77457efd 100644 --- a/crates/paimon/src/vindex/mod.rs +++ b/crates/paimon/src/vindex/mod.rs @@ -17,6 +17,8 @@ pub mod reader; +pub mod pkvector; + use crate::spec::{DataField, DataType}; use paimon_vindex_core::index::VectorIndexConfig; use std::collections::HashMap; diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs new file mode 100644 index 00000000..f87e7f98 --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -0,0 +1,459 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use super::bucket::BucketAnnSegment; +use super::data_invalid; +use super::metric::VectorSearchMetric; +use super::result::PkVectorSearchResult; +use crate::deletion_vector::DeletionVector; +use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; +use crate::vector_search::VectorSearch; + +/// Build the live-row-id mask for the ANN reader's `include_row_ids` filter, in +/// segment-ordinal space (source files concatenated in order). Mirrors Java +/// `PkVectorAnnSegmentSearcher.liveRowPositions`. +/// +/// Only source files present in `active_source_files` contribute live ordinals; +/// inactive sources' ordinal ranges are masked out entirely (their rows are no +/// longer readable in this snapshot). Deletion vectors are applied only to active +/// sources. +/// +/// Returns `None` only when every source file is active AND no deletion vector is +/// relevant — nothing to mask. Otherwise returns the masked live ids. +pub(crate) fn build_live_row_ids( + source_files: &[PkVectorSourceFile], + active_source_files: &HashSet, + deletion_vectors: &HashMap>, +) -> crate::Result> { + let all_active = source_files + .iter() + .all(|f| active_source_files.contains(f.file_name())); + let has_relevant_dv = source_files + .iter() + .any(|f| deletion_vectors.contains_key(f.file_name())); + if all_active && !has_relevant_dv { + return Ok(None); + } + + let mut live = roaring::RoaringTreemap::new(); + let mut deleted = roaring::RoaringTreemap::new(); + let mut file_offset: u64 = 0; + for source_file in source_files { + let row_count = u64::try_from(source_file.row_count()) + .map_err(|_| data_invalid("vector source row count must not be negative"))?; + let end = file_offset + .checked_add(row_count) + .ok_or_else(|| data_invalid("vector source row counts overflow u64"))?; + let active = active_source_files.contains(source_file.file_name()); + if active && row_count > 0 { + live.insert_range(file_offset..end); + } + if active { + if let Some(dv) = deletion_vectors.get(source_file.file_name()) { + for position in dv.iter() { + let global = file_offset.checked_add(position).ok_or_else(|| { + data_invalid("vector source deleted position overflows u64") + })?; + deleted.insert(global); + } + } + } + file_offset = end; + } + live -= deleted; + Ok(Some(live)) +} + +/// Map ANN `(ordinal, score)` pairs to physical `(data file, position)` results, +/// validating ordinals against source metadata, rejecting hits that resolve to an +/// inactive source file, and rejecting hits on snapshot-deleted rows. Mirrors the +/// post-processing loop of Java `PkVectorAnnSegmentSearcher.search`. Results are +/// sorted BEST_FIRST. +pub(crate) fn map_ann_results( + scored: &[(u64, f32)], + source_meta: &PkVectorSourceMeta, + active_source_files: &HashSet, + deletion_vectors: &HashMap>, + metric: VectorSearchMetric, +) -> crate::Result> { + let mut results = Vec::with_capacity(scored.len()); + for &(ordinal, score) in scored { + let ordinal_i64 = i64::try_from(ordinal) + .map_err(|_| data_invalid(format!("ANN ordinal {ordinal} exceeds i64::MAX")))?; + let (data_file_name, row_position) = source_meta.resolve(ordinal_i64)?; + if !active_source_files.contains(&data_file_name) { + return Err(data_invalid(format!( + "ANN segment returned inactive source {data_file_name}" + ))); + } + if let Some(dv) = deletion_vectors.get(&data_file_name) { + let pos = u64::try_from(row_position) + .map_err(|_| data_invalid("resolved row position must not be negative"))?; + if dv.is_deleted(pos) { + return Err(data_invalid(format!( + "ANN segment returned snapshot-deleted row position {row_position} in {data_file_name}" + ))); + } + } + results.push(PkVectorSearchResult { + data_file_name, + row_position, + distance: metric.score_to_distance(score), + }); + } + results.sort_by(|a, b| { + a.distance + .total_cmp(&b.distance) + .then_with(|| a.data_file_name.cmp(&b.data_file_name)) + .then_with(|| a.row_position.cmp(&b.row_position)) + }); + Ok(results) +} + +/// One ANN segment's search dependency for the bucket kernel. Bucket tests fake +/// this (mirroring Java's mock of `PkVectorAnnSegmentSearcher`). +pub(crate) trait PkVectorAnnSearcher { + #[allow(clippy::too_many_arguments)] + fn search( + &self, + segment: &BucketAnnSegment, + query: &[f32], + metric: VectorSearchMetric, + limit: usize, + active_source_files: &HashSet, + deletion_vectors: &HashMap>, + search_options: &HashMap, + ) -> crate::Result>; +} + +/// Scorer seam: drives the underlying vindex ANN reader. Returns `ordinal -> +/// score` (higher-is-better). Any negative labels are skipped by the existing +/// `vindex` reader (`collect_results` drops `row_id < 0`), so this seam only +/// ever yields non-negative `u64` ordinals — no signed-label handling is needed +/// downstream. +/// +/// The production scorer drives `VindexVectorGlobalIndexReader::visit_vector_search` +/// with a segment's index bytes; tests inject a synthetic scorer. The adapter's +/// own logic (live-row masking, ordinal mapping, deletion checks, ordering) is +/// exercised independently of the scorer. +type Scorer = Box crate::Result>>>; + +/// Structural vindex-backed `PkVectorAnnSearcher`. Composes the pure helpers +/// (`build_live_row_ids`, `map_ann_results`) around the scorer seam. +pub(crate) struct VindexAnnSearcher { + field_name: String, + scorer: Scorer, +} + +impl VindexAnnSearcher { + pub(crate) fn new(field_name: String, scorer: Scorer) -> Self { + Self { field_name, scorer } + } +} + +impl PkVectorAnnSearcher for VindexAnnSearcher { + fn search( + &self, + segment: &BucketAnnSegment, + query: &[f32], + metric: VectorSearchMetric, + limit: usize, + active_source_files: &HashSet, + deletion_vectors: &HashMap>, + search_options: &HashMap, + ) -> crate::Result> { + if limit == 0 { + return Err(data_invalid("vector search limit must be positive")); + } + let source_files = segment.source_meta.source_files(); + let mut search = VectorSearch::new(query.to_vec(), limit, self.field_name.clone())? + .with_options(search_options.clone()); + if let Some(live) = build_live_row_ids(source_files, active_source_files, deletion_vectors)? + { + search = search.with_include_row_ids(live); + } + let scored = match (self.scorer)(&search)? { + Some(map) => map, + None => return Ok(Vec::new()), + }; + let scored: Vec<(u64, f32)> = scored.into_iter().collect(); + map_ann_results( + &scored, + &segment.source_meta, + active_source_files, + deletion_vectors, + metric, + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use roaring::RoaringBitmap; + + fn source_meta(files: &[(&str, i64)]) -> PkVectorSourceMeta { + let files = files + .iter() + .map(|(name, rows)| PkVectorSourceFile::new((*name).to_string(), *rows).unwrap()) + .collect(); + PkVectorSourceMeta::new(files).unwrap() + } + + fn dv(deleted: &[u32]) -> Arc { + let mut bitmap = RoaringBitmap::new(); + for &p in deleted { + bitmap.insert(p); + } + Arc::new(DeletionVector::from_bitmap(bitmap)) + } + + fn active_set(names: &[&str]) -> HashSet { + names.iter().map(|n| (*n).to_string()).collect() + } + + #[test] + fn test_build_live_row_ids_none_when_all_active_and_no_relevant_dv() { + let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; + let active = active_set(&["f0"]); + // All active + empty map -> None. + assert!(build_live_row_ids(&files, &active, &HashMap::new()) + .unwrap() + .is_none()); + // All active + non-empty map but no matching file name -> None. + let mut dvs = HashMap::new(); + dvs.insert("other".to_string(), dv(&[0])); + assert!(build_live_row_ids(&files, &active, &dvs).unwrap().is_none()); + } + + #[test] + fn test_build_live_row_ids_masks_inactive_source_ordinal_range() { + // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). f1 is inactive, + // so its whole ordinal range is masked out; f0 stays fully live. No DV. + let files = vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 2).unwrap(), + ]; + let live = build_live_row_ids(&files, &active_set(&["f0"]), &HashMap::new()) + .unwrap() + .unwrap(); + assert_eq!(live.iter().collect::>(), vec![0, 1, 2]); + } + + #[test] + fn test_build_live_row_ids_masks_deleted_positions_with_file_offsets() { + // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). + let files = vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 2).unwrap(), + ]; + let mut dvs = HashMap::new(); + dvs.insert("f0".to_string(), dv(&[1])); // deletes global 1 + dvs.insert("f1".to_string(), dv(&[0])); // deletes global 3 + let live = build_live_row_ids(&files, &active_set(&["f0", "f1"]), &dvs) + .unwrap() + .unwrap(); + assert_eq!(live.iter().collect::>(), vec![0, 2, 4]); + } + + #[test] + fn test_map_ann_results_maps_ordinals_to_positions_and_scores() { + let meta = source_meta(&[("f0", 3), ("f1", 5)]); + // ordinal 3 -> (f1, 0); ordinal 0 -> (f0, 0). l2 score_to_distance(0.5)=1.0. + let scored = [(3u64, 0.5f32), (0u64, 0.5f32)]; + let results = map_ann_results( + &scored, + &meta, + &active_set(&["f0", "f1"]), + &HashMap::new(), + VectorSearchMetric::L2, + ) + .unwrap(); + assert_eq!( + results, + vec![ + PkVectorSearchResult { + data_file_name: "f0".into(), + row_position: 0, + distance: 1.0 + }, + PkVectorSearchResult { + data_file_name: "f1".into(), + row_position: 0, + distance: 1.0 + }, + ] + ); + } + + #[test] + fn test_map_ann_results_rejects_out_of_range_ordinal() { + let meta = source_meta(&[("f0", 3)]); + let err = map_ann_results( + &[(3u64, 0.5)], + &meta, + &active_set(&["f0"]), + &HashMap::new(), + VectorSearchMetric::L2, + ) + .unwrap_err(); + assert!(err.to_string().contains("out of range") || err.to_string().contains("ordinal")); + } + + #[test] + fn test_map_ann_results_rejects_hit_resolving_to_inactive_source() { + // ordinal 3 resolves to f1, which is not in the active set -> error. + let meta = source_meta(&[("f0", 3), ("f1", 5)]); + let err = map_ann_results( + &[(3u64, 0.5)], + &meta, + &active_set(&["f0"]), + &HashMap::new(), + VectorSearchMetric::L2, + ) + .unwrap_err(); + assert!(err.to_string().contains("inactive")); + } + + #[test] + fn test_map_ann_results_rejects_hit_on_deleted_position() { + let meta = source_meta(&[("f0", 3)]); + let mut dvs = HashMap::new(); + dvs.insert("f0".to_string(), dv(&[1])); // position 1 deleted + let err = map_ann_results( + &[(1u64, 0.5)], + &meta, + &active_set(&["f0"]), + &dvs, + VectorSearchMetric::L2, + ) + .unwrap_err(); + assert!(err.to_string().contains("deleted")); + } + + #[test] + fn test_vindex_adapter_composes_live_rows_and_maps_results() { + // Scorer records the VectorSearch it received and returns synthetic ordinals. + // The scorer must be `'static`, so share the recording cells via `Rc` moved + // into the closure rather than borrowing locals. + use std::cell::RefCell; + use std::rc::Rc; + let seen_limit = Rc::new(RefCell::new(0usize)); + let seen_has_filter = Rc::new(RefCell::new(false)); + let scorer_limit = Rc::clone(&seen_limit); + let scorer_has_filter = Rc::clone(&seen_has_filter); + let searcher = VindexAnnSearcher::new( + "embedding".to_string(), + Box::new(move |search: &VectorSearch| { + *scorer_limit.borrow_mut() = search.limit; + *scorer_has_filter.borrow_mut() = search.include_row_ids.is_some(); + let mut scores = HashMap::new(); + scores.insert(3u64, 0.5f32); // -> (f1, 0) + scores.insert(0u64, 0.25f32); // -> (f0, 0), l2 dist 3.0 + Ok(Some(scores)) + }), + ); + let segment = BucketAnnSegment { + source_meta: { + use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + PkVectorSourceMeta::new(vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 5).unwrap(), + ]) + .unwrap() + }, + }; + let mut dvs = HashMap::new(); + dvs.insert("f0".to_string(), dv(&[1])); + let results = searcher + .search( + &segment, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &active_set(&["f0", "f1"]), + &dvs, + &HashMap::new(), + ) + .unwrap(); + // Sorted BEST_FIRST by distance: (f1,0) dist 1.0 then (f0,0) dist 3.0. + assert_eq!(results[0].data_file_name, "f1"); + assert_eq!(results[1].data_file_name, "f0"); + assert_eq!(*seen_limit.borrow(), 2); + assert!( + *seen_has_filter.borrow(), + "DV present -> include_row_ids set" + ); + } + + #[test] + fn test_vindex_adapter_rejects_non_positive_limit() { + let searcher = VindexAnnSearcher::new( + "embedding".to_string(), + Box::new(|_: &VectorSearch| Ok(None)), + ); + let segment = BucketAnnSegment { + source_meta: { + use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()]) + .unwrap() + }, + }; + let err = searcher + .search( + &segment, + &[0.0, 0.0], + VectorSearchMetric::L2, + 0, + &active_set(&["f0"]), + &HashMap::new(), + &HashMap::new(), + ) + .unwrap_err(); + assert!(err.to_string().contains("positive")); + } + + #[test] + fn test_vindex_adapter_empty_scorer_result_is_empty() { + let searcher = VindexAnnSearcher::new( + "embedding".to_string(), + Box::new(|_: &VectorSearch| Ok(None)), + ); + let segment = BucketAnnSegment { + source_meta: { + use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()]) + .unwrap() + }, + }; + let results = searcher + .search( + &segment, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &active_set(&["f0"]), + &HashMap::new(), + &HashMap::new(), + ) + .unwrap(); + assert!(results.is_empty()); + } +} diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs new file mode 100644 index 00000000..23a2b1cb --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -0,0 +1,554 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::sync::Arc; + +use super::ann::PkVectorAnnSearcher; +use super::data_invalid; +use super::exact::exact_search; +use super::metric::VectorSearchMetric; +use super::reader::PkVectorReader; +use super::result::PkVectorSearchResult; +use crate::deletion_vector::DeletionVector; +use crate::spec::PkVectorSourceMeta; + +/// One ANN segment to be searched by the bucket kernel: the source metadata +/// resolving segment ordinals back to physical `(data file, position)`. Only +/// `source_meta` is needed for ordinal mapping and live-row masking. +pub(crate) struct BucketAnnSegment { + pub source_meta: PkVectorSourceMeta, +} + +/// A data file participating in the bucket search, with its row count. Used by +/// the bucket kernel to plan exact vs. ANN search over active files. +pub(crate) struct BucketActiveFile { + pub file_name: String, + pub row_count: i64, +} + +/// Total BEST_FIRST order over results: distance ASC, then data_file_name ASC, +/// then row_position ASC. `total_cmp` keeps it NaN-safe and panic-free. +fn best_first(a: &PkVectorSearchResult, b: &PkVectorSearchResult) -> Ordering { + a.distance + .total_cmp(&b.distance) + .then_with(|| a.data_file_name.cmp(&b.data_file_name)) + .then_with(|| a.row_position.cmp(&b.row_position)) +} + +/// A candidate wrapped so a max-heap keeps the WORST (BEST_FIRST-largest) +/// candidate on top; popping evicts the least-wanted one. Mirrors the +/// `PriorityQueue<>(limit, BEST_FIRST.reversed())` in Java +/// `PrimaryKeyVectorBucketSearch`. +struct WorstFirst(PkVectorSearchResult); + +impl PartialEq for WorstFirst { + fn eq(&self, other: &Self) -> bool { + best_first(&self.0, &other.0) == Ordering::Equal + } +} +impl Eq for WorstFirst {} +impl PartialOrd for WorstFirst { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for WorstFirst { + fn cmp(&self, other: &Self) -> Ordering { + best_first(&self.0, &other.0) + } +} + +/// Add `candidate` to a bounded (size `limit`) BEST_FIRST Top-K max-heap: push if +/// under capacity, else replace the current worst iff the candidate beats it. +/// `O(log limit)` per call. Mirrors Java `PrimaryKeyVectorBucketSearch.add`. +fn add_candidate(heap: &mut BinaryHeap, candidate: PkVectorSearchResult, limit: usize) { + if heap.len() < limit { + heap.push(WorstFirst(candidate)); + } else if heap + .peek() + .is_some_and(|worst| best_first(&candidate, &worst.0) == Ordering::Less) + { + heap.pop(); + heap.push(WorstFirst(candidate)); + } +} + +/// ANN + exact data-file fallback search for one snapshot bucket. Mirrors Java +/// `org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch.search`. +/// +/// `ann_searcher` may be `None` only when there are no ANN segments; segments +/// present with `None` is an error. +#[allow(clippy::too_many_arguments)] +pub(crate) fn bucket_search( + ann_searcher: Option<&dyn PkVectorAnnSearcher>, + ann_segments: &[BucketAnnSegment], + active_files: &[BucketActiveFile], + deletion_vectors: &HashMap>, + exact_reader_factory: &mut dyn FnMut( + &BucketActiveFile, + ) -> crate::Result>, + query: &[f32], + metric: VectorSearchMetric, + limit: usize, + search_options: &HashMap, +) -> crate::Result> { + if limit == 0 { + return Err(data_invalid("vector search limit must be positive")); + } + + let mut files_by_name: HashMap<&str, &BucketActiveFile> = HashMap::new(); + for file in active_files { + if file.row_count < 0 { + return Err(data_invalid(format!( + "active data file {} row count must not be negative: {}", + file.file_name, file.row_count + ))); + } + if files_by_name + .insert(file.file_name.as_str(), file) + .is_some() + { + return Err(data_invalid(format!( + "duplicate data file: {}", + file.file_name + ))); + } + } + + let mut heap: BinaryHeap = BinaryHeap::with_capacity(limit + 1); + let active_source_files: HashSet = + files_by_name.keys().map(|name| name.to_string()).collect(); + let mut covered: HashSet = HashSet::new(); + + for segment in ann_segments { + for source in segment.source_meta.source_files() { + // An ANN source that is no longer an active file (e.g. compacted away) + // is skipped, not rejected: its ordinal range is masked out of the ANN + // live-row bitmap and the remaining active sources are still searched. + // Mirrors Java master `PrimaryKeyVectorBucketSearch` (`file == null` + // -> continue). Active sources still require a row-count match. + match files_by_name.get(source.file_name()) { + Some(active) if active.row_count == source.row_count() => { + covered.insert(source.file_name().to_string()); + } + Some(_) => { + return Err(data_invalid(format!( + "ANN source {} does not match the active data file", + source.file_name() + ))); + } + None => continue, + } + } + let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is not configured"))?; + for result in searcher.search( + segment, + query, + metric, + limit, + &active_source_files, + deletion_vectors, + search_options, + )? { + add_candidate(&mut heap, result, limit); + } + } + + for file in active_files { + if covered.contains(&file.file_name) { + continue; + } + let dv = deletion_vectors.get(&file.file_name).cloned(); + let is_excluded = move |position: i64| -> bool { + match &dv { + Some(dv) => u64::try_from(position) + .map(|p| dv.is_deleted(p)) + .unwrap_or(false), + None => false, + } + }; + let mut reader = exact_reader_factory(file)?; + for result in exact_search( + &file.file_name, + reader.as_mut(), + query, + metric, + limit, + &is_excluded, + )? { + add_candidate(&mut heap, result, limit); + } + } + + let mut results: Vec = heap.into_iter().map(|w| w.0).collect(); + results.sort_by(best_first); + Ok(results) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::PkVectorSourceFile; + use crate::vindex::pkvector::ann::PkVectorAnnSearcher; + use crate::vindex::pkvector::reader::test_support::ArrayReader; + use roaring::RoaringBitmap; + use std::cell::RefCell; + + fn meta(files: &[(&str, i64)]) -> PkVectorSourceMeta { + PkVectorSourceMeta::new( + files + .iter() + .map(|(n, r)| PkVectorSourceFile::new((*n).into(), *r).unwrap()) + .collect(), + ) + .unwrap() + } + + fn active(name: &str, rows: i64) -> BucketActiveFile { + BucketActiveFile { + file_name: name.into(), + row_count: rows, + } + } + + /// Fake ANN searcher returning preset results and recording calls. + struct FakeAnnSearcher { + result: Vec, + } + impl PkVectorAnnSearcher for FakeAnnSearcher { + fn search( + &self, + _segment: &BucketAnnSegment, + _query: &[f32], + _metric: VectorSearchMetric, + _limit: usize, + _active_source_files: &HashSet, + _dvs: &HashMap>, + _opts: &HashMap, + ) -> crate::Result> { + Ok(self.result.clone()) + } + } + + #[test] + fn test_rejects_non_positive_limit() { + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + None, + &[], + &[], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 0, + &HashMap::new(), + ) + .unwrap_err(); + assert!(err.to_string().contains("positive")); + } + + #[test] + fn test_bounded_heap_evicts_by_best_first_tiebreak_over_limit() { + // All candidates share distance 1.0, so eviction is decided purely by the + // BEST_FIRST tie-break (data_file_name ASC, then row_position ASC). Feed + // more than `limit` ANN hits and assert the kept set is the smallest + // (file, position) pairs in that order. Locks the bounded-heap merge. + let segment = BucketAnnSegment { + source_meta: meta(&[("data-1", 3)]), + }; + let hit = |file: &str, pos: i64| PkVectorSearchResult { + data_file_name: file.into(), + row_position: pos, + distance: 1.0, + }; + // Deliberately unsorted input across two files at the same distance. + let ann = FakeAnnSearcher { + result: vec![ + hit("data-2", 0), + hit("data-1", 2), + hit("data-1", 0), + hit("data-2", 1), + hit("data-1", 1), + ], + }; + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let results = bucket_search( + Some(&ann), + &[segment], + &[active("data-1", 3)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 3, + &HashMap::new(), + ) + .unwrap(); + // Top-3 BEST_FIRST: (data-1,0), (data-1,1), (data-1,2) — the larger + // data_file_name "data-2" entries are evicted despite equal distance. + assert_eq!( + results + .iter() + .map(|r| (r.data_file_name.as_str(), r.row_position)) + .collect::>(), + vec![("data-1", 0), ("data-1", 1), ("data-1", 2)] + ); + } + + #[test] + fn test_merges_ann_and_exact_without_rescanning_covered_files() { + // data-1 is ANN-covered; data-2 is exact fallback. Factory must never be + // called for data-1. + let segment = BucketAnnSegment { + source_meta: meta(&[("data-1", 2)]), + }; + let ann = FakeAnnSearcher { + result: vec![PkVectorSearchResult { + data_file_name: "data-1".into(), + row_position: 1, + distance: 0.5, + }], + }; + let calls = RefCell::new(Vec::::new()); + let mut factory = |f: &BucketActiveFile| -> crate::Result> { + calls.borrow_mut().push(f.file_name.clone()); + // data-2 vectors: pos0 {1,0} dist 1.0, pos1 {3,0} dist 9.0 + Ok(Box::new(ArrayReader::new( + 2, + vec![Some(vec![1.0, 0.0]), Some(vec![3.0, 0.0])], + ))) + }; + let results = bucket_search( + Some(&ann), + &[segment], + &[active("data-1", 2), active("data-2", 2)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &HashMap::new(), + ) + .unwrap(); + assert_eq!( + results, + vec![ + PkVectorSearchResult { + data_file_name: "data-1".into(), + row_position: 1, + distance: 0.5 + }, + PkVectorSearchResult { + data_file_name: "data-2".into(), + row_position: 0, + distance: 1.0 + }, + ] + ); + assert_eq!(calls.borrow().as_slice(), &["data-2".to_string()]); + } + + #[test] + fn test_exact_fallback_merges_files_and_applies_deletion_vectors() { + // No ANN. data-1 pos0 {0,0} deleted; remaining candidates merge across files. + let calls = RefCell::new(0); + let mut factory = |f: &BucketActiveFile| -> crate::Result> { + *calls.borrow_mut() += 1; + let vectors = match f.file_name.as_str() { + "data-1" => vec![Some(vec![0.0, 0.0]), Some(vec![2.0, 0.0])], + "data-2" => vec![Some(vec![1.0, 0.0]), None], + _ => unreachable!(), + }; + Ok(Box::new(ArrayReader::new(2, vectors))) + }; + let mut dvs: HashMap> = HashMap::new(); + let mut bm = RoaringBitmap::new(); + bm.insert(0); // data-1 position 0 deleted + dvs.insert("data-1".into(), Arc::new(DeletionVector::from_bitmap(bm))); + + let results = bucket_search( + None, + &[], + &[active("data-1", 2), active("data-2", 2)], + &dvs, + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &HashMap::new(), + ) + .unwrap(); + // Candidates: data-2 pos0 {1,0} dist 1.0; data-1 pos1 {2,0} dist 4.0. + // (data-1 pos0 deleted, data-2 pos1 null.) + assert_eq!( + results, + vec![ + PkVectorSearchResult { + data_file_name: "data-2".into(), + row_position: 0, + distance: 1.0 + }, + PkVectorSearchResult { + data_file_name: "data-1".into(), + row_position: 1, + distance: 4.0 + }, + ] + ); + } + + #[test] + fn test_rejects_duplicate_active_file_name() { + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + None, + &[], + &[active("dup", 1), active("dup", 1)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 1, + &HashMap::new(), + ) + .unwrap_err(); + assert!(err.to_string().contains("duplicate") || err.to_string().contains("Duplicate")); + } + + #[test] + fn test_rejects_ann_source_row_count_mismatch_for_active_file() { + let ann = FakeAnnSearcher { result: vec![] }; + // Segment references data-1 with 2 rows, but the active file has 3 rows. + // An active source with a mismatched row count is still a hard error. + let segment = BucketAnnSegment { + source_meta: meta(&[("data-1", 2)]), + }; + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + Some(&ann), + &[segment], + &[active("data-1", 3)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 1, + &HashMap::new(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("does not match") || err.to_string().contains("ANN source") + ); + } + + #[test] + fn test_skips_inactive_ann_source_and_searches_active_ones() { + // Segment covers [data-1, data-2] but only data-1 is still active + // (data-2 was compacted away). Java master skips the inactive source + // instead of failing the whole query; data-2 is neither covered (so it + // is not treated as ANN-covered) nor an active file (so it is not exact + // scanned). The ANN searcher still runs for the segment. + let segment = BucketAnnSegment { + source_meta: meta(&[("data-1", 2), ("data-2", 2)]), + }; + let ann = FakeAnnSearcher { + result: vec![PkVectorSearchResult { + data_file_name: "data-1".into(), + row_position: 0, + distance: 0.5, + }], + }; + let calls = RefCell::new(Vec::::new()); + let mut factory = |f: &BucketActiveFile| -> crate::Result> { + calls.borrow_mut().push(f.file_name.clone()); + unreachable!("only data-1 is active and it is ANN-covered") + }; + let results = bucket_search( + Some(&ann), + &[segment], + &[active("data-1", 2)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &HashMap::new(), + ) + .unwrap(); + assert_eq!( + results, + vec![PkVectorSearchResult { + data_file_name: "data-1".into(), + row_position: 0, + distance: 0.5 + }] + ); + // No exact fallback ran: data-1 is ANN-covered, data-2 is not active. + assert!(calls.borrow().is_empty()); + } + + #[test] + fn test_rejects_segments_without_ann_searcher() { + let segment = BucketAnnSegment { + source_meta: meta(&[("data-1", 2)]), + }; + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + None, + &[segment], + &[active("data-1", 2)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 1, + &HashMap::new(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("ANN search is not configured") + || err.to_string().contains("not configured") + ); + } + + #[test] + fn test_negative_active_row_count_rejected() { + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + None, + &[], + &[active("data-1", -1)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 1, + &HashMap::new(), + ) + .unwrap_err(); + assert!(err.to_string().contains("row count") || err.to_string().contains("-1")); + } +} diff --git a/crates/paimon/src/vindex/pkvector/exact.rs b/crates/paimon/src/vindex/pkvector/exact.rs new file mode 100644 index 00000000..c9c04805 --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/exact.rs @@ -0,0 +1,289 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cmp::Ordering; +use std::collections::BinaryHeap; + +use super::data_invalid; +use super::metric::VectorSearchMetric; +use super::reader::PkVectorReader; +use super::result::PkVectorSearchResult; + +/// A candidate wrapped so a max-heap keeps the WORST candidate on top: +/// worst = largest distance, ties broken by largest row_position. Popping the +/// top therefore evicts the least-wanted candidate. Uses `total_cmp` for a +/// deterministic total order over f32 (NaN-safe, no panic). +struct WorstFirst(PkVectorSearchResult); + +impl PartialEq for WorstFirst { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} +impl Eq for WorstFirst {} +impl PartialOrd for WorstFirst { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for WorstFirst { + fn cmp(&self, other: &Self) -> Ordering { + self.0 + .distance + .total_cmp(&other.0.distance) + .then_with(|| self.0.row_position.cmp(&other.0.row_position)) + } +} + +/// True if `candidate` ranks strictly better (BEST_FIRST) than the current +/// worst-on-heap `weakest`: smaller distance, ties broken by smaller position. +fn is_better_than(candidate: &PkVectorSearchResult, weakest: &PkVectorSearchResult) -> bool { + candidate + .distance + .total_cmp(&weakest.distance) + .then_with(|| candidate.row_position.cmp(&weakest.row_position)) + == Ordering::Less +} + +/// Exact Top-K over one sequential physical-row vector source. Mirrors Java +/// `PkVectorExactSearcher.search`. Results are sorted BEST_FIRST: distance ASC, +/// then row_position ASC (single file, so data_file_name is constant). +pub(crate) fn exact_search( + data_file_name: &str, + reader: &mut dyn PkVectorReader, + query: &[f32], + metric: VectorSearchMetric, + limit: usize, + is_excluded: &dyn Fn(i64) -> bool, +) -> crate::Result> { + if query.len() != reader.dimension() { + return Err(data_invalid(format!( + "query vector dimension does not match: index expects {}, got {}", + reader.dimension(), + query.len() + ))); + } + if limit == 0 { + return Err(data_invalid("vector search limit must be positive")); + } + if let Some(i) = query.iter().position(|v| !v.is_finite()) { + return Err(data_invalid(format!( + "query vector element at position {i} must be finite" + ))); + } + let row_count = reader.row_count(); + if row_count < 0 { + return Err(data_invalid(format!( + "vector reader row count must not be negative: {row_count}" + ))); + } + + let mut reuse = vec![0.0f32; reader.dimension()]; + let mut heap: BinaryHeap = BinaryHeap::with_capacity(limit + 1); + for position in 0..row_count { + let has_vector = reader.read_next_vector(&mut reuse)?; + if !has_vector || is_excluded(position) { + continue; + } + let candidate = PkVectorSearchResult { + data_file_name: data_file_name.to_string(), + row_position: position, + distance: metric.compute_distance(query, &reuse), + }; + if heap.len() < limit { + heap.push(WorstFirst(candidate)); + } else if heap + .peek() + .is_some_and(|worst| is_better_than(&candidate, &worst.0)) + { + heap.pop(); + heap.push(WorstFirst(candidate)); + } + } + + let mut results: Vec = heap.into_iter().map(|w| w.0).collect(); + results.sort_by(|a, b| { + a.distance + .total_cmp(&b.distance) + .then_with(|| a.row_position.cmp(&b.row_position)) + }); + Ok(results) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::vindex::pkvector::metric::VectorSearchMetric; + use crate::vindex::pkvector::reader::test_support::ArrayReader; + + fn no_exclusion() -> impl Fn(i64) -> bool { + |_| false + } + + #[test] + fn test_distances_for_supported_metrics() { + // Java testDistancesForSupportedMetrics: q=[2,0] over stored [1,0]. + for (metric, expected) in [ + (VectorSearchMetric::L2, 1.0f32), + (VectorSearchMetric::Cosine, 0.0), + (VectorSearchMetric::InnerProduct, -2.0), + ] { + let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]); + let results = exact_search( + "data-file", + &mut reader, + &[2.0, 0.0], + metric, + 1, + &no_exclusion(), + ) + .unwrap(); + assert_eq!(results[0].distance, expected); + } + } + + #[test] + fn test_rejects_dimension_mismatch() { + let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]); + let err = exact_search( + "data-file", + &mut reader, + &[1.0], + VectorSearchMetric::L2, + 1, + &no_exclusion(), + ) + .unwrap_err(); + assert!(err.to_string().contains("dimension")); + } + + #[test] + fn test_rejects_non_positive_limit() { + let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]); + let err = exact_search( + "data-file", + &mut reader, + &[1.0, 0.0], + VectorSearchMetric::L2, + 0, + &no_exclusion(), + ) + .unwrap_err(); + assert!(err.to_string().contains("positive")); + } + + #[test] + fn test_rejects_non_finite_query() { + let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]); + let err = exact_search( + "data-file", + &mut reader, + &[f32::NAN, 0.0], + VectorSearchMetric::L2, + 1, + &no_exclusion(), + ) + .unwrap_err(); + assert!(err.to_string().contains("finite")); + } + + #[test] + fn test_rejects_negative_row_count() { + struct NegativeReader; + impl PkVectorReader for NegativeReader { + fn dimension(&self) -> usize { + 2 + } + fn row_count(&self) -> i64 { + -1 + } + fn read_next_vector(&mut self, _reuse: &mut [f32]) -> crate::Result { + unreachable!() + } + } + let mut reader = NegativeReader; + let err = exact_search( + "data-file", + &mut reader, + &[1.0, 0.0], + VectorSearchMetric::L2, + 1, + &no_exclusion(), + ) + .unwrap_err(); + assert!(err.to_string().contains("row count") || err.to_string().contains("-1")); + } + + #[test] + fn test_preserves_null_and_excluded_physical_positions() { + // Java testPreservesNullAndDeletedPhysicalPositions: + // vectors [{3,0}, null, {1,0}, {2,0}], q=[0,0], l2, limit 2, exclude pos==2. + let mut reader = ArrayReader::new( + 2, + vec![ + Some(vec![3.0, 0.0]), + None, + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + ], + ); + let results = exact_search( + "data-file", + &mut reader, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &|pos| pos == 2, + ) + .unwrap(); + assert_eq!( + results, + vec![ + PkVectorSearchResult { + data_file_name: "data-file".into(), + row_position: 3, + distance: 4.0 + }, + PkVectorSearchResult { + data_file_name: "data-file".into(), + row_position: 0, + distance: 9.0 + }, + ] + ); + } + + #[test] + fn test_tie_break_prefers_smaller_row_position() { + // Equal distances -> smaller row_position ranks first. + let mut reader = + ArrayReader::new(1, vec![Some(vec![1.0]), Some(vec![1.0]), Some(vec![1.0])]); + let results = exact_search( + "data-file", + &mut reader, + &[0.0], + VectorSearchMetric::L2, + 2, + &no_exclusion(), + ) + .unwrap(); + assert_eq!( + results.iter().map(|r| r.row_position).collect::>(), + vec![0, 1] + ); + } +} diff --git a/crates/paimon/src/vindex/pkvector/metric.rs b/crates/paimon/src/vindex/pkvector/metric.rs new file mode 100644 index 00000000..1dda56c7 --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/metric.rs @@ -0,0 +1,211 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::data_invalid; + +/// Normalize a metric name: lowercase and `-` → `_`. NO trim (deliberately +/// stricter than the build-side `vindex::normalize_metric`, to match Java +/// `VectorSearchMetric.normalize`). +pub(crate) fn normalize_metric(metric: &str) -> String { + metric.to_ascii_lowercase().replace('-', "_") +} + +/// True if the (normalized) metric is one of the three supported metrics. +pub(crate) fn is_supported_metric(metric: &str) -> bool { + matches!( + normalize_metric(metric).as_str(), + "l2" | "cosine" | "inner_product" + ) +} + +/// Numeric semantics for a supported vector search metric. Mirrors Java +/// `org.apache.paimon.globalindex.VectorSearchMetric`. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum VectorSearchMetric { + L2, + Cosine, + InnerProduct, +} + +impl VectorSearchMetric { + /// Normalize, validate, and map to the enum. Errors on an unsupported metric. + pub(crate) fn parse(metric: &str) -> crate::Result { + match normalize_metric(metric).as_str() { + "l2" => Ok(Self::L2), + "cosine" => Ok(Self::Cosine), + "inner_product" => Ok(Self::InnerProduct), + other => Err(data_invalid(format!( + "unsupported vector distance metric: {other}" + ))), + } + } + + /// Higher-is-better score for exact vector search. + pub(crate) fn compute_score(&self, query: &[f32], stored: &[f32]) -> f32 { + match self { + Self::L2 => 1.0 / (1.0 + squared_l2(query, stored)), + Self::Cosine => cosine_similarity(query, stored), + Self::InnerProduct => inner_product(query, stored), + } + } + + /// Lower-is-better distance for exact vector search. + pub(crate) fn compute_distance(&self, query: &[f32], stored: &[f32]) -> f32 { + match self { + Self::L2 => squared_l2(query, stored), + Self::Cosine => cosine_distance(cosine_similarity(query, stored)), + Self::InnerProduct => -inner_product(query, stored), + } + } + + /// Convert a higher-is-better standardized index score to a lower-is-better + /// distance. For L2 with `score == 0.0` this yields `inf` (natural f32 + /// behavior, matching Java — no clamp). + pub(crate) fn score_to_distance(&self, score: f32) -> f32 { + match self { + Self::L2 => 1.0 / score - 1.0, + Self::Cosine => cosine_distance(score), + Self::InnerProduct => -score, + } + } +} + +fn squared_l2(query: &[f32], stored: &[f32]) -> f32 { + let mut squared = 0.0f32; + for i in 0..query.len() { + let delta = query[i] - stored[i]; + squared += delta * delta; + } + squared +} + +fn cosine_similarity(query: &[f32], stored: &[f32]) -> f32 { + let mut dot = 0.0f32; + let mut query_norm = 0.0f32; + let mut stored_norm = 0.0f32; + for i in 0..query.len() { + dot += query[i] * stored[i]; + query_norm += query[i] * query[i]; + stored_norm += stored[i] * stored[i]; + } + let denominator = ((query_norm as f64).sqrt() * (stored_norm as f64).sqrt()) as f32; + if denominator == 0.0 { + 0.0 + } else { + dot / denominator + } +} + +fn inner_product(query: &[f32], stored: &[f32]) -> f32 { + let mut dot = 0.0f32; + for i in 0..query.len() { + dot += query[i] * stored[i]; + } + dot +} + +fn cosine_distance(similarity: f32) -> f32 { + 1.0 - similarity.clamp(-1.0, 1.0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_normalize_lowercases_and_replaces_hyphens_without_trimming() { + assert_eq!(normalize_metric("Inner-Product"), "inner_product"); + assert_eq!(normalize_metric("L2"), "l2"); + // No trim: surrounding whitespace is preserved (unlike the build-side helper). + assert_eq!(normalize_metric(" l2 "), " l2 "); + } + + #[test] + fn test_is_supported_only_for_three_metrics() { + assert!(is_supported_metric("L2")); + assert!(is_supported_metric("cosine")); + assert!(is_supported_metric("inner-product")); + assert!(!is_supported_metric("manhattan")); + assert!(!is_supported_metric(" l2 ")); + } + + #[test] + fn test_parse_rejects_unsupported_metric() { + assert!(VectorSearchMetric::parse("l2").is_ok()); + assert!(VectorSearchMetric::parse("cosine").is_ok()); + assert!(VectorSearchMetric::parse("inner_product").is_ok()); + assert!(VectorSearchMetric::parse("manhattan").is_err()); + } + + #[test] + fn test_compute_distance_matches_java_anchor() { + // Java PkVectorExactSearcherTest.testDistancesForSupportedMetrics: + // q=[2,0], s=[1,0] -> l2=1.0, cosine=0.0, inner_product=-2.0 + let q = [2.0f32, 0.0]; + let s = [1.0f32, 0.0]; + assert_eq!(VectorSearchMetric::L2.compute_distance(&q, &s), 1.0); + assert_eq!(VectorSearchMetric::Cosine.compute_distance(&q, &s), 0.0); + assert_eq!( + VectorSearchMetric::InnerProduct.compute_distance(&q, &s), + -2.0 + ); + } + + #[test] + fn test_compute_score_higher_is_better() { + let q = [2.0f32, 0.0]; + let s = [1.0f32, 0.0]; + assert_eq!(VectorSearchMetric::L2.compute_score(&q, &s), 0.5); // 1/(1+1) + assert_eq!(VectorSearchMetric::Cosine.compute_score(&q, &s), 1.0); // parallel + assert_eq!(VectorSearchMetric::InnerProduct.compute_score(&q, &s), 2.0); + } + + #[test] + fn test_cosine_zero_norm_similarity_is_zero() { + let zero = [0.0f32, 0.0]; + let s = [1.0f32, 0.0]; + assert_eq!(VectorSearchMetric::Cosine.compute_score(&zero, &s), 0.0); + assert_eq!(VectorSearchMetric::Cosine.compute_distance(&zero, &s), 1.0); + } + + #[test] + fn test_cosine_non_perfect_square_norm_uses_f64_sqrt() { + // query [0,3] -> norm 9, stored [1,2] -> norm 5; sqrt(5) is irrational + // so the f32-sqrt path (each sqrt taken in f32, then widened) and the + // f64-sqrt path (widen first, sqrt in f64) produce different f32 bits: + // buggy score 0.8944271 vs correct 0.8944272. Encode the f64 contract in + // the expected value (dot / (sqrt(9.0f64) * sqrt(5.0f64)) as f32) rather + // than a magic literal, so this pins the Java-matching f64 arithmetic. + let q = [0.0f32, 3.0]; + let s = [1.0f32, 2.0]; + let dot = 6.0f32; + let denominator = ((9.0f64).sqrt() * (5.0f64).sqrt()) as f32; + let expected = dot / denominator; + assert_eq!(VectorSearchMetric::Cosine.compute_score(&q, &s), expected); + } + + #[test] + fn test_score_to_distance_l2_zero_score_is_infinite() { + assert!(VectorSearchMetric::L2.score_to_distance(0.0).is_infinite()); + assert_eq!(VectorSearchMetric::L2.score_to_distance(0.5), 1.0); // 1/0.5 - 1 + assert_eq!( + VectorSearchMetric::InnerProduct.score_to_distance(2.0), + -2.0 + ); + assert_eq!(VectorSearchMetric::Cosine.score_to_distance(1.0), 0.0); + } +} diff --git a/crates/paimon/src/vindex/pkvector/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs new file mode 100644 index 00000000..3a685ca3 --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -0,0 +1,42 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Primary-key vector (bucket-local ANN) search kernel. +//! +//! Read-only bucket-local approximate-nearest-neighbour search over the +//! primary-key vector index. + +// The kernel is crate-private and has no production caller yet, so its items +// are unreachable outside their own tests. Suppress the resulting dead_code +// lint at the module boundary until the read path wires it in. +#![allow(dead_code)] + +pub(crate) mod ann; +pub(crate) mod bucket; +pub(crate) mod exact; +pub(crate) mod metric; +pub(crate) mod reader; +pub(crate) mod result; + +/// Shared constructor for validation failures in this module (mirrors Java +/// `checkArgument` / `IllegalArgumentException`). +pub(crate) fn data_invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} diff --git a/crates/paimon/src/vindex/pkvector/reader.rs b/crates/paimon/src/vindex/pkvector/reader.rs new file mode 100644 index 00000000..3298013a --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/reader.rs @@ -0,0 +1,80 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +/// Sequential exact-scan source of vectors for one data file. Mirrors Java +/// `org.apache.paimon.index.pkvector.PkVectorReader`. +pub(crate) trait PkVectorReader { + fn dimension(&self) -> usize; + + fn row_count(&self) -> i64; + + /// Read the next row's vector into `reuse` (`reuse.len() == dimension()`). + /// Returns `false` when the physical row is a NULL vector: it is not scored, + /// but the physical position still advances by one. Each call advances + /// exactly one physical row. + fn read_next_vector(&mut self, reuse: &mut [f32]) -> crate::Result; +} + +#[cfg(test)] +pub(crate) mod test_support { + use super::PkVectorReader; + + /// In-memory `PkVectorReader` for tests. `None` entries are NULL rows. + /// Mirrors Java's test `ArrayReader`. + pub(crate) struct ArrayReader { + dimension: usize, + vectors: Vec>>, + position: usize, + } + + impl ArrayReader { + pub(crate) fn new(dimension: usize, vectors: Vec>>) -> Self { + Self { + dimension, + vectors, + position: 0, + } + } + } + + impl PkVectorReader for ArrayReader { + fn dimension(&self) -> usize { + self.dimension + } + + fn row_count(&self) -> i64 { + self.vectors.len() as i64 + } + + fn read_next_vector(&mut self, reuse: &mut [f32]) -> crate::Result { + assert_eq!( + reuse.len(), + self.dimension, + "reuse buffer must equal dimension" + ); + let entry = &self.vectors[self.position]; + self.position += 1; + match entry { + Some(vector) => { + reuse.copy_from_slice(vector); + Ok(true) + } + None => Ok(false), + } + } + } +} diff --git a/crates/paimon/src/vindex/pkvector/result.rs b/crates/paimon/src/vindex/pkvector/result.rs new file mode 100644 index 00000000..97bea5ea --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/result.rs @@ -0,0 +1,27 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +/// One vector search result addressed by source data-file row position. +/// Mirrors Java `org.apache.paimon.index.pkvector.PkVectorSearchResult`. +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct PkVectorSearchResult { + pub data_file_name: String, + /// Zero-based physical row position within the source data file. + pub row_position: i64, + /// Lower-is-better distance. + pub distance: f32, +}