From f59ec1ada63f24a9308b00fbbd17e0624d7a0f2e Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 00:40:07 +0800 Subject: [PATCH 01/12] feat(deletion-vector): add is_deleted position accessor --- crates/paimon/src/deletion_vector/core.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) 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)); + } } From 07bcfe68ab7db496e9b08656a78287b9f0dcad8b Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 00:51:54 +0800 Subject: [PATCH 02/12] feat(vindex): add pkvector module scaffold and VectorSearchMetric --- crates/paimon/src/vindex/mod.rs | 2 + crates/paimon/src/vindex/pkvector/ann.rs | 16 ++ crates/paimon/src/vindex/pkvector/bucket.rs | 16 ++ crates/paimon/src/vindex/pkvector/exact.rs | 16 ++ crates/paimon/src/vindex/pkvector/metric.rs | 205 ++++++++++++++++++++ crates/paimon/src/vindex/pkvector/mod.rs | 38 ++++ crates/paimon/src/vindex/pkvector/reader.rs | 16 ++ crates/paimon/src/vindex/pkvector/result.rs | 16 ++ 8 files changed, 325 insertions(+) create mode 100644 crates/paimon/src/vindex/pkvector/ann.rs create mode 100644 crates/paimon/src/vindex/pkvector/bucket.rs create mode 100644 crates/paimon/src/vindex/pkvector/exact.rs create mode 100644 crates/paimon/src/vindex/pkvector/metric.rs create mode 100644 crates/paimon/src/vindex/pkvector/mod.rs create mode 100644 crates/paimon/src/vindex/pkvector/reader.rs create mode 100644 crates/paimon/src/vindex/pkvector/result.rs 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..b248758b --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -0,0 +1,16 @@ +// 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. diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs new file mode 100644 index 00000000..b248758b --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -0,0 +1,16 @@ +// 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. diff --git a/crates/paimon/src/vindex/pkvector/exact.rs b/crates/paimon/src/vindex/pkvector/exact.rs new file mode 100644 index 00000000..b248758b --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/exact.rs @@ -0,0 +1,16 @@ +// 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. diff --git a/crates/paimon/src/vindex/pkvector/metric.rs b/crates/paimon/src/vindex/pkvector/metric.rs new file mode 100644 index 00000000..64c9cdd5 --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/metric.rs @@ -0,0 +1,205 @@ +// 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..cfe8b257 --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -0,0 +1,38 @@ +// 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. +//! +//! Mirrors Java PR apache/paimon#8569 (`[core] Add primary-key vector bucket +//! search`) at commit `a50a36ff8`. Read-only, synthetically tested; real ANN +//! index-byte validation is deferred to PR4 (Java-written fixtures). + +pub(crate) mod metric; +pub(crate) mod result; +pub(crate) mod reader; +pub(crate) mod exact; +pub(crate) mod ann; +pub(crate) mod bucket; + +/// 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..b248758b --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/reader.rs @@ -0,0 +1,16 @@ +// 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. diff --git a/crates/paimon/src/vindex/pkvector/result.rs b/crates/paimon/src/vindex/pkvector/result.rs new file mode 100644 index 00000000..b248758b --- /dev/null +++ b/crates/paimon/src/vindex/pkvector/result.rs @@ -0,0 +1,16 @@ +// 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. From 63551c0d37073e1ae5757dc90d67329af6964eb3 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 00:56:58 +0800 Subject: [PATCH 03/12] feat(vindex): add PkVectorReader trait, result type, and exact Top-K search --- crates/paimon/src/vindex/pkvector/exact.rs | 273 ++++++++++++++++++++ crates/paimon/src/vindex/pkvector/mod.rs | 8 +- crates/paimon/src/vindex/pkvector/reader.rs | 64 +++++ crates/paimon/src/vindex/pkvector/result.rs | 11 + 4 files changed, 352 insertions(+), 4 deletions(-) diff --git a/crates/paimon/src/vindex/pkvector/exact.rs b/crates/paimon/src/vindex/pkvector/exact.rs index b248758b..c9c04805 100644 --- a/crates/paimon/src/vindex/pkvector/exact.rs +++ b/crates/paimon/src/vindex/pkvector/exact.rs @@ -14,3 +14,276 @@ // 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/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs index cfe8b257..8e416c4d 100644 --- a/crates/paimon/src/vindex/pkvector/mod.rs +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -21,12 +21,12 @@ //! search`) at commit `a50a36ff8`. Read-only, synthetically tested; real ANN //! index-byte validation is deferred to PR4 (Java-written fixtures). -pub(crate) mod metric; -pub(crate) mod result; -pub(crate) mod reader; -pub(crate) mod exact; 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`). diff --git a/crates/paimon/src/vindex/pkvector/reader.rs b/crates/paimon/src/vindex/pkvector/reader.rs index b248758b..3298013a 100644 --- a/crates/paimon/src/vindex/pkvector/reader.rs +++ b/crates/paimon/src/vindex/pkvector/reader.rs @@ -14,3 +14,67 @@ // 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 index b248758b..97bea5ea 100644 --- a/crates/paimon/src/vindex/pkvector/result.rs +++ b/crates/paimon/src/vindex/pkvector/result.rs @@ -14,3 +14,14 @@ // 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, +} From fed69a71231b872b69c0c40f6aeb209df35a3817 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:03:24 +0800 Subject: [PATCH 04/12] feat(vindex): add ANN live-row mask and ordinal->position mapping helpers --- crates/paimon/src/vindex/pkvector/ann.rs | 183 +++++++++++++++++++++++ 1 file changed, 183 insertions(+) diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index b248758b..c3334889 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -14,3 +14,186 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use super::data_invalid; +use super::metric::VectorSearchMetric; +use super::result::PkVectorSearchResult; +use crate::deletion_vector::DeletionVector; +use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + +/// 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`. +/// +/// Returns `None` when no deletion vector is relevant to these source files +/// (empty map, or no source file has a matching DV) — nothing to mask. When at +/// least one source file has a matching DV, returns the masked live ids. +pub(crate) fn build_live_row_ids( + source_files: &[PkVectorSourceFile], + deletion_vectors: &HashMap>, +) -> crate::Result> { + let has_relevant_dv = source_files + .iter() + .any(|f| deletion_vectors.contains_key(f.file_name())); + if !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"))?; + if row_count > 0 { + live.insert_range(file_offset..file_offset + row_count); + } + if let Some(dv) = deletion_vectors.get(source_file.file_name()) { + for position in dv.iter() { + deleted.insert(file_offset + position); + } + } + file_offset = file_offset + .checked_add(row_count) + .ok_or_else(|| data_invalid("vector source row counts overflow u64"))?; + } + live -= deleted; + Ok(Some(live)) +} + +/// Map ANN `(ordinal, score)` pairs to physical `(data file, position)` results, +/// validating ordinals against source metadata 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, + 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 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) +} + +#[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)) + } + + #[test] + fn test_build_live_row_ids_none_when_no_relevant_dv() { + let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; + // Empty map -> None. + assert!(build_live_row_ids(&files, &HashMap::new()) + .unwrap() + .is_none()); + // 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, &dvs).unwrap().is_none()); + } + + #[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, &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, &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, + &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_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, &dvs, VectorSearchMetric::L2).unwrap_err(); + assert!(err.to_string().contains("deleted")); + } +} From 4f4bc964138e3ea05a4a60f9a75db750bf5a56b6 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:11:30 +0800 Subject: [PATCH 05/12] feat(vindex): add PkVectorAnnSearcher trait and structural vindex adapter --- crates/paimon/src/vindex/pkvector/ann.rs | 171 ++++++++++++++++++++ crates/paimon/src/vindex/pkvector/bucket.rs | 18 +++ 2 files changed, 189 insertions(+) diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index c3334889..02a82531 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -18,11 +18,13 @@ use std::collections::HashMap; 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 @@ -103,6 +105,69 @@ pub(crate) fn map_ann_results( 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 { + fn search( + &self, + segment: &BucketAnnSegment, + query: &[f32], + metric: VectorSearchMetric, + limit: usize, + deletion_vectors: &HashMap>, + search_options: &HashMap, + ) -> crate::Result>; +} + +/// Scorer seam: drives the underlying vindex ANN reader. Returns `ordinal -> +/// score` (higher-is-better), with negative labels already dropped by the reader. +/// The real implementation (driving `VindexVectorGlobalIndexReader::visit_vector_search` +/// with a segment's index bytes) is supplied by PR4; PR2 tests inject a synthetic +/// scorer. This is the fixture-deferred boundary — the adapter's own logic +/// (live-row masking, ordinal mapping, deletion checks, ordering) is fully tested. +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, + 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, 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, deletion_vectors, metric) + } +} + #[cfg(test)] mod tests { use super::*; @@ -196,4 +261,110 @@ mod tests { let err = map_ann_results(&[(1u64, 0.5)], &meta, &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, + &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, + &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, + &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 index b248758b..a0c0d87d 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -14,3 +14,21 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + +use crate::spec::PkVectorSourceMeta; + +/// One ANN segment to be searched by the bucket kernel: the vindex index bytes +/// plus the source metadata resolving segment ordinals back to physical +/// `(data file, position)`. The index-byte identity (which physical index the +/// segment reads) is a PR4 concern — PR2 only needs `source_meta` for ordinal +/// mapping and live-row masking; the reader wiring is added in PR4. +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 (Task 6) to plan exact vs. ANN search over active files. +pub(crate) struct BucketActiveFile { + pub file_name: String, + pub row_count: i64, +} From 11a3b283420e6da0b1d083889befaef8ff8495ba Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:15:30 +0800 Subject: [PATCH 06/12] style(vindex): rustfmt metric assert lines to satisfy fmt --check The verbatim single-line assert_eq! calls in metric.rs tests exceed rustfmt's fn_call_width and were reflowed to multi-line. No logic change. --- crates/paimon/src/vindex/pkvector/metric.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/vindex/pkvector/metric.rs b/crates/paimon/src/vindex/pkvector/metric.rs index 64c9cdd5..1dda56c7 100644 --- a/crates/paimon/src/vindex/pkvector/metric.rs +++ b/crates/paimon/src/vindex/pkvector/metric.rs @@ -159,7 +159,10 @@ mod tests { 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); + assert_eq!( + VectorSearchMetric::InnerProduct.compute_distance(&q, &s), + -2.0 + ); } #[test] @@ -199,7 +202,10 @@ mod tests { 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::InnerProduct.score_to_distance(2.0), + -2.0 + ); assert_eq!(VectorSearchMetric::Cosine.score_to_distance(1.0), 0.0); } } From d365e40719c8c056e5452ca3df0c8813889c3c6b Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:21:25 +0800 Subject: [PATCH 07/12] feat(vindex): add primary-key vector bucket search (ANN + exact merge) --- crates/paimon/src/vindex/pkvector/bucket.rs | 410 ++++++++++++++++++++ 1 file changed, 410 insertions(+) diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index a0c0d87d..fe5014e3 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -15,6 +15,16 @@ // specific language governing permissions and limitations // under the License. +use std::collections::{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 vindex index bytes @@ -32,3 +42,403 @@ pub(crate) struct BucketActiveFile { pub file_name: String, pub row_count: i64, } + +/// True if `candidate` ranks strictly better (BEST_FIRST) than `weakest`: +/// distance ASC, then data_file_name ASC, then row_position ASC. +fn is_better_than(candidate: &PkVectorSearchResult, weakest: &PkVectorSearchResult) -> bool { + candidate + .distance + .total_cmp(&weakest.distance) + .then_with(|| candidate.data_file_name.cmp(&weakest.data_file_name)) + .then_with(|| candidate.row_position.cmp(&weakest.row_position)) + == std::cmp::Ordering::Less +} + +fn add_candidate( + heap: &mut Vec, + candidate: PkVectorSearchResult, + limit: usize, +) { + if heap.len() < limit { + heap.push(candidate); + return; + } + // Find the current worst (BEST_FIRST-largest) and replace if the candidate beats it. + let worst_idx = heap + .iter() + .enumerate() + .max_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)) + }) + .map(|(i, _)| i); + if let Some(i) = worst_idx { + if is_better_than(&candidate, &heap[i]) { + heap[i] = 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: Vec = Vec::with_capacity(limit + 1); + let mut covered: HashSet = HashSet::new(); + + for segment in ann_segments { + for source in segment.source_meta.source_files() { + match files_by_name.get(source.file_name()) { + Some(active) if active.row_count == source.row_count() => { + covered.insert(source.file_name().to_string()); + } + _ => { + return Err(data_invalid(format!( + "ANN source {} does not match the active data file", + source.file_name() + ))); + } + } + } + let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is not configured"))?; + for result in searcher.search( + segment, + query, + metric, + limit, + 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); + } + } + + heap.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(heap) +} + +#[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, + _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_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_missing_or_mismatched_active_file() { + let ann = FakeAnnSearcher { result: vec![] }; + // Segment references data-1 with 2 rows, but active file has 3 rows. + 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_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")); + } +} From c0b6b81c6d4fb86da422deb34ffc7db3a0a9ac19 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:32:17 +0800 Subject: [PATCH 08/12] chore(vindex): allow(dead_code) on crate-private pkvector kernel The pkvector bucket-search kernel is ported ahead of the read-path wiring (PR3/PR4). Its items are crate-private with no production caller yet, so clippy -D warnings flags them as dead_code. Suppress at the module boundary; remove this allow once PR3/PR4 lands production callers. --- crates/paimon/src/vindex/pkvector/mod.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/crates/paimon/src/vindex/pkvector/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs index 8e416c4d..7b73c9ae 100644 --- a/crates/paimon/src/vindex/pkvector/mod.rs +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -21,6 +21,13 @@ //! search`) at commit `a50a36ff8`. Read-only, synthetically tested; real ANN //! index-byte validation is deferred to PR4 (Java-written fixtures). +// PR2 ports the PK-vector bucket-local search kernel before PR3/PR4 wires it +// into the read path. The kernel is crate-private with no production caller +// yet, so its items are unreachable outside their own tests. Suppress the +// resulting dead_code lint at the module boundary; remove this allow once +// PR3/PR4 lands production callers. +#![allow(dead_code)] + pub(crate) mod ann; pub(crate) mod bucket; pub(crate) mod exact; From 1dea0b6e3a0469b433dc6b003cb051d60a7ab8a6 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:38:26 +0800 Subject: [PATCH 09/12] refactor(vindex): compute live-row end offset once with checked_add --- crates/paimon/src/vindex/pkvector/ann.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 02a82531..5363945f 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -50,17 +50,18 @@ pub(crate) fn build_live_row_ids( 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"))?; if row_count > 0 { - live.insert_range(file_offset..file_offset + row_count); + live.insert_range(file_offset..end); } if let Some(dv) = deletion_vectors.get(source_file.file_name()) { for position in dv.iter() { deleted.insert(file_offset + position); } } - file_offset = file_offset - .checked_add(row_count) - .ok_or_else(|| data_invalid("vector source row counts overflow u64"))?; + file_offset = end; } live -= deleted; Ok(Some(live)) From da39ff153f342bf32a058a21b6df66076f048cf0 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 01:50:56 +0800 Subject: [PATCH 10/12] refactor(vindex): guard live-row deleted position with checked_add MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use checked_add for the global deleted-position offset (file_offset + position) to match the surrounding overflow-guarded arithmetic, and correct the scorer-seam doc: negative ANN labels are SKIPPED by the existing vindex reader (collect_results drops row_id < 0), not rejected — so the seam only yields non-negative u64 ordinals and no signed-label handling is needed. --- crates/paimon/src/vindex/pkvector/ann.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 5363945f..6b112063 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -58,7 +58,10 @@ pub(crate) fn build_live_row_ids( } if let Some(dv) = deletion_vectors.get(source_file.file_name()) { for position in dv.iter() { - deleted.insert(file_offset + position); + let global = file_offset + .checked_add(position) + .ok_or_else(|| data_invalid("vector source deleted position overflows u64"))?; + deleted.insert(global); } } file_offset = end; @@ -121,7 +124,10 @@ pub(crate) trait PkVectorAnnSearcher { } /// Scorer seam: drives the underlying vindex ANN reader. Returns `ordinal -> -/// score` (higher-is-better), with negative labels already dropped by the reader. +/// 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 real implementation (driving `VindexVectorGlobalIndexReader::visit_vector_search` /// with a segment's index bytes) is supplied by PR4; PR2 tests inject a synthetic /// scorer. This is the fixture-deferred boundary — the adapter's own logic From 3188073c0f9e489bda387a521d88567e6f4aae89 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 16:58:39 +0800 Subject: [PATCH 11/12] feat(vindex): skip inactive ANN sources and use a BinaryHeap for the bucket Top-K Address both review comments on PR #516 (apache/paimon-rust). 1. Inactive ANN sources (correctness/interop with current Java master). Adopt current master PrimaryKeyVectorBucketSearch / PkVectorAnnSegmentSearcher behavior: an ANN segment source that is no longer an active file (e.g. compacted away) is skipped, not rejected, so the bucket still searches the remaining active sources. - build_live_row_ids takes active_source_files and masks out the ordinal ranges of inactive sources (they contribute no live ids; deletion vectors apply only to active sources). Returns None only when all sources are active and no deletion vector is relevant. - map_ann_results defensively rejects a hit resolving to an inactive source (mirrors the Java checkArgument), since the live-row mask should already have excluded it. - bucket_search skips ANN sources missing from the active set; active sources with a mismatched row count are still a hard error. - PkVectorAnnSearcher::search and VindexAnnSearcher thread the active source set through. 2. Bucket Top-K merge complexity. Replace the Vec + max_by linear worst-element scan with a BinaryHeap keeping the BEST_FIRST-worst on top, reducing the merge to O(candidate_count log K) and matching the Java PriorityQueue and the exact.rs heap. Final output stays deterministic via a BEST_FIRST sort. Tests: split the old missing-or-mismatched reject into a row-count mismatch reject (still errors) plus skip-inactive coverage; add inactive ordinal-range masking, inactive-hit rejection, and a bounded-heap tie-break eviction test. All 34 pkvector tests pass; fmt and clippy clean. --- crates/paimon/src/vindex/pkvector/ann.rs | 133 ++++++++++--- crates/paimon/src/vindex/pkvector/bucket.rs | 200 +++++++++++++++----- 2 files changed, 263 insertions(+), 70 deletions(-) diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 6b112063..9ebbf3cd 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use super::bucket::BucketAnnSegment; @@ -30,17 +30,25 @@ use crate::vector_search::VectorSearch; /// segment-ordinal space (source files concatenated in order). Mirrors Java /// `PkVectorAnnSegmentSearcher.liveRowPositions`. /// -/// Returns `None` when no deletion vector is relevant to these source files -/// (empty map, or no source file has a matching DV) — nothing to mask. When at -/// least one source file has a matching DV, returns the masked live ids. +/// 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 !has_relevant_dv { + if all_active && !has_relevant_dv { return Ok(None); } @@ -53,15 +61,18 @@ pub(crate) fn build_live_row_ids( let end = file_offset .checked_add(row_count) .ok_or_else(|| data_invalid("vector source row counts overflow u64"))?; - if row_count > 0 { + let active = active_source_files.contains(source_file.file_name()); + if active && row_count > 0 { live.insert_range(file_offset..end); } - 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); + 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; @@ -71,12 +82,14 @@ pub(crate) fn build_live_row_ids( } /// Map ANN `(ordinal, score)` pairs to physical `(data file, position)` results, -/// validating ordinals against source metadata and rejecting hits on -/// snapshot-deleted rows. Mirrors the post-processing loop of Java -/// `PkVectorAnnSegmentSearcher.search`. Results are sorted BEST_FIRST. +/// 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> { @@ -85,6 +98,11 @@ pub(crate) fn map_ann_results( 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"))?; @@ -112,12 +130,14 @@ pub(crate) fn map_ann_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>; @@ -154,6 +174,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { query: &[f32], metric: VectorSearchMetric, limit: usize, + active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, ) -> crate::Result> { @@ -163,7 +184,8 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { 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, deletion_vectors)? { + 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)? { @@ -171,7 +193,13 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { None => return Ok(Vec::new()), }; let scored: Vec<(u64, f32)> = scored.into_iter().collect(); - map_ann_results(&scored, &segment.source_meta, deletion_vectors, metric) + map_ann_results( + &scored, + &segment.source_meta, + active_source_files, + deletion_vectors, + metric, + ) } } @@ -196,17 +224,36 @@ mod tests { 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_no_relevant_dv() { + fn test_build_live_row_ids_none_when_all_active_and_no_relevant_dv() { let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; - // Empty map -> None. - assert!(build_live_row_ids(&files, &HashMap::new()) + let active = active_set(&["f0"]); + // All active + empty map -> None. + assert!(build_live_row_ids(&files, &active, &HashMap::new()) .unwrap() .is_none()); - // Non-empty map but no matching file name -> 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, &dvs).unwrap().is_none()); + 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] @@ -219,7 +266,9 @@ mod tests { 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, &dvs).unwrap().unwrap(); + let live = build_live_row_ids(&files, &active_set(&["f0", "f1"]), &dvs) + .unwrap() + .unwrap(); assert_eq!(live.iter().collect::>(), vec![0, 2, 4]); } @@ -228,8 +277,14 @@ mod tests { 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, &HashMap::new(), VectorSearchMetric::L2).unwrap(); + let results = map_ann_results( + &scored, + &meta, + &active_set(&["f0", "f1"]), + &HashMap::new(), + VectorSearchMetric::L2, + ) + .unwrap(); assert_eq!( results, vec![ @@ -253,6 +308,7 @@ mod tests { let err = map_ann_results( &[(3u64, 0.5)], &meta, + &active_set(&["f0"]), &HashMap::new(), VectorSearchMetric::L2, ) @@ -260,12 +316,34 @@ mod tests { 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, &dvs, VectorSearchMetric::L2).unwrap_err(); + let err = map_ann_results( + &[(1u64, 0.5)], + &meta, + &active_set(&["f0"]), + &dvs, + VectorSearchMetric::L2, + ) + .unwrap_err(); assert!(err.to_string().contains("deleted")); } @@ -309,6 +387,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 2, + &active_set(&["f0", "f1"]), &dvs, &HashMap::new(), ) @@ -342,6 +421,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 0, + &active_set(&["f0"]), &HashMap::new(), &HashMap::new(), ) @@ -368,6 +448,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 2, + &active_set(&["f0"]), &HashMap::new(), &HashMap::new(), ) diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index fe5014e3..eea7998e 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -use std::collections::{HashMap, HashSet}; +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet}; use std::sync::Arc; use super::ann::PkVectorAnnSearcher; @@ -43,41 +44,50 @@ pub(crate) struct BucketActiveFile { pub row_count: i64, } -/// True if `candidate` ranks strictly better (BEST_FIRST) than `weakest`: -/// distance ASC, then data_file_name ASC, then row_position ASC. -fn is_better_than(candidate: &PkVectorSearchResult, weakest: &PkVectorSearchResult) -> bool { - candidate - .distance - .total_cmp(&weakest.distance) - .then_with(|| candidate.data_file_name.cmp(&weakest.data_file_name)) - .then_with(|| candidate.row_position.cmp(&weakest.row_position)) - == std::cmp::Ordering::Less +/// 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)) } -fn add_candidate( - heap: &mut Vec, - candidate: PkVectorSearchResult, - limit: usize, -) { - if heap.len() < limit { - heap.push(candidate); - return; +/// 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 } - // Find the current worst (BEST_FIRST-largest) and replace if the candidate beats it. - let worst_idx = heap - .iter() - .enumerate() - .max_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)) - }) - .map(|(i, _)| i); - if let Some(i) = worst_idx { - if is_better_than(&candidate, &heap[i]) { - heap[i] = candidate; - } +} +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)); } } @@ -123,21 +133,29 @@ pub(crate) fn bucket_search( } } - let mut heap: Vec = Vec::with_capacity(limit + 1); + 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"))?; @@ -146,6 +164,7 @@ pub(crate) fn bucket_search( query, metric, limit, + &active_source_files, deletion_vectors, search_options, )? { @@ -179,13 +198,9 @@ pub(crate) fn bucket_search( } } - heap.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(heap) + let mut results: Vec = heap.into_iter().map(|w| w.0).collect(); + results.sort_by(best_first); + Ok(results) } #[cfg(test)] @@ -225,6 +240,7 @@ mod tests { _query: &[f32], _metric: VectorSearchMetric, _limit: usize, + _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, ) -> crate::Result> { @@ -251,6 +267,55 @@ mod tests { 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 @@ -373,9 +438,10 @@ mod tests { } #[test] - fn test_rejects_ann_source_missing_or_mismatched_active_file() { + fn test_rejects_ann_source_row_count_mismatch_for_active_file() { let ann = FakeAnnSearcher { result: vec![] }; - // Segment references data-1 with 2 rows, but active file has 3 rows. + // 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)]), }; @@ -398,6 +464,52 @@ mod tests { ); } + #[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 { From 36ee093af31e03af6d6938dc37c6d167295ce7a6 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 13 Jul 2026 20:38:22 +0800 Subject: [PATCH 12/12] chore(vindex): drop internal-process references from comments Rewrite doc/comments to describe the code itself rather than the porting process, removing references to PR phases, task numbers, upstream commit hashes, and fixture-deferral notes. --- crates/paimon/src/spec/pk_vector_source.rs | 4 ++-- crates/paimon/src/vindex/pkvector/ann.rs | 9 +++++---- crates/paimon/src/vindex/pkvector/bucket.rs | 10 ++++------ crates/paimon/src/vindex/pkvector/mod.rs | 13 +++++-------- 4 files changed, 16 insertions(+), 20 deletions(-) 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/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 9ebbf3cd..f87e7f98 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -148,10 +148,11 @@ pub(crate) trait PkVectorAnnSearcher { /// `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 real implementation (driving `VindexVectorGlobalIndexReader::visit_vector_search` -/// with a segment's index bytes) is supplied by PR4; PR2 tests inject a synthetic -/// scorer. This is the fixture-deferred boundary — the adapter's own logic -/// (live-row masking, ordinal mapping, deletion checks, ordering) is fully tested. +/// +/// 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 diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index eea7998e..23a2b1cb 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -28,17 +28,15 @@ use super::result::PkVectorSearchResult; use crate::deletion_vector::DeletionVector; use crate::spec::PkVectorSourceMeta; -/// One ANN segment to be searched by the bucket kernel: the vindex index bytes -/// plus the source metadata resolving segment ordinals back to physical -/// `(data file, position)`. The index-byte identity (which physical index the -/// segment reads) is a PR4 concern — PR2 only needs `source_meta` for ordinal -/// mapping and live-row masking; the reader wiring is added in PR4. +/// 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 (Task 6) to plan exact vs. ANN search over active files. +/// the bucket kernel to plan exact vs. ANN search over active files. pub(crate) struct BucketActiveFile { pub file_name: String, pub row_count: i64, diff --git a/crates/paimon/src/vindex/pkvector/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs index 7b73c9ae..3a685ca3 100644 --- a/crates/paimon/src/vindex/pkvector/mod.rs +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -17,15 +17,12 @@ //! Primary-key vector (bucket-local ANN) search kernel. //! -//! Mirrors Java PR apache/paimon#8569 (`[core] Add primary-key vector bucket -//! search`) at commit `a50a36ff8`. Read-only, synthetically tested; real ANN -//! index-byte validation is deferred to PR4 (Java-written fixtures). +//! Read-only bucket-local approximate-nearest-neighbour search over the +//! primary-key vector index. -// PR2 ports the PK-vector bucket-local search kernel before PR3/PR4 wires it -// into the read path. The kernel is crate-private with no production caller -// yet, so its items are unreachable outside their own tests. Suppress the -// resulting dead_code lint at the module boundary; remove this allow once -// PR3/PR4 lands production callers. +// 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;