This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new f5cb23f feat: support vindex vector search (#399)
f5cb23f is described below
commit f5cb23fe3b7c055e945185537139fb6ba7f30199
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Jun 19 23:13:01 2026 +0800
feat: support vindex vector search (#399)
---
.../testdata/test_java_vindex_vector.tar.gz | Bin 0 -> 3125 bytes
.../integrations/datafusion/tests/read_tables.rs | 38 ++-
crates/paimon/Cargo.toml | 1 +
crates/paimon/src/lib.rs | 2 +
crates/paimon/src/lumina/mod.rs | 262 -----------------
crates/paimon/src/lumina/reader.rs | 146 +++++-----
crates/paimon/src/table/vector_search_builder.rs | 90 +++++-
crates/paimon/src/vector_search.rs | 285 ++++++++++++++++++
crates/paimon/src/vindex/mod.rs | 47 +++
crates/paimon/src/vindex/reader.rs | 322 +++++++++++++++++++++
10 files changed, 845 insertions(+), 348 deletions(-)
diff --git
a/crates/integrations/datafusion/testdata/test_java_vindex_vector.tar.gz
b/crates/integrations/datafusion/testdata/test_java_vindex_vector.tar.gz
new file mode 100644
index 0000000..4a88084
Binary files /dev/null and
b/crates/integrations/datafusion/testdata/test_java_vindex_vector.tar.gz differ
diff --git a/crates/integrations/datafusion/tests/read_tables.rs
b/crates/integrations/datafusion/tests/read_tables.rs
index 396e951..30760d2 100644
--- a/crates/integrations/datafusion/tests/read_tables.rs
+++ b/crates/integrations/datafusion/tests/read_tables.rs
@@ -1021,9 +1021,10 @@ mod vector_search_tests {
use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
use paimon_datafusion::{register_vector_search, SQLContext};
- fn extract_test_warehouse() -> (tempfile::TempDir, String) {
+ fn extract_test_warehouse(archive_name: &str) -> (tempfile::TempDir,
String) {
let archive_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
- .join("testdata/test_lumina_vector.tar.gz");
+ .join("testdata")
+ .join(archive_name);
let file = std::fs::File::open(&archive_path)
.unwrap_or_else(|e| panic!("Failed to open {}: {e}",
archive_path.display()));
let decoder = flate2::read::GzDecoder::new(file);
@@ -1038,8 +1039,8 @@ mod vector_search_tests {
(tmp, warehouse)
}
- async fn create_vector_search_context() -> (SQLContext, tempfile::TempDir)
{
- let (tmp, warehouse) = extract_test_warehouse();
+ async fn create_vector_search_context(archive_name: &str) -> (SQLContext,
tempfile::TempDir) {
+ let (tmp, warehouse) = extract_test_warehouse(archive_name);
let mut options = Options::new();
options.set(CatalogOptions::WAREHOUSE, warehouse);
let catalog = FileSystemCatalog::new(options).expect("Failed to create
catalog");
@@ -1053,6 +1054,14 @@ mod vector_search_tests {
(ctx, tmp)
}
+ async fn create_lumina_vector_search_context() -> (SQLContext,
tempfile::TempDir) {
+ create_vector_search_context("test_lumina_vector.tar.gz").await
+ }
+
+ async fn create_java_vindex_vector_search_context() -> (SQLContext,
tempfile::TempDir) {
+ create_vector_search_context("test_java_vindex_vector.tar.gz").await
+ }
+
fn extract_ids(batches: &[datafusion::arrow::record_batch::RecordBatch])
-> Vec<i32> {
let mut ids = Vec::new();
for batch in batches {
@@ -1070,7 +1079,7 @@ mod vector_search_tests {
#[tokio::test]
async fn test_vector_search_top3() {
- let (ctx, _tmp) = create_vector_search_context().await;
+ let (ctx, _tmp) = create_lumina_vector_search_context().await;
let batches = ctx
.sql("SELECT id FROM
vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0,
0.0, 0.0]', 3)")
.await
@@ -1086,7 +1095,7 @@ mod vector_search_tests {
#[tokio::test]
async fn test_vector_search_top6_returns_all() {
- let (ctx, _tmp) = create_vector_search_context().await;
+ let (ctx, _tmp) = create_lumina_vector_search_context().await;
let batches = ctx
.sql("SELECT id FROM
vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0,
0.0, 0.0]', 6)")
.await
@@ -1101,7 +1110,7 @@ mod vector_search_tests {
#[tokio::test]
async fn test_vector_search_without_matching_index_returns_empty() {
- let (ctx, _tmp) = create_vector_search_context().await;
+ let (ctx, _tmp) = create_lumina_vector_search_context().await;
let batches = ctx
.sql("SELECT id FROM
vector_search('paimon.default.test_lumina_vector', 'missing_embedding',
'[1.0]', 10)")
.await
@@ -1116,4 +1125,19 @@ mod vector_search_tests {
"vector_search without a matching Lumina index should not fall
back to a full table scan"
);
}
+
+ #[tokio::test]
+ async fn test_vector_search_java_vindex_table() {
+ let (ctx, _tmp) = create_java_vindex_vector_search_context().await;
+ let batches = ctx
+ .sql("SELECT id FROM
vector_search('paimon.default.test_java_vindex_vector', 'embedding', '[1.0,
0.0, 0.0, 0.0]', 3)")
+ .await
+ .expect("SQL should parse")
+ .collect()
+ .await
+ .expect("query should execute");
+
+ let ids = extract_ids(&batches);
+ assert_eq!(ids, vec![0, 1, 2]);
+ }
}
diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml
index 2a45ff5..9f2b08b 100644
--- a/crates/paimon/Cargo.toml
+++ b/crates/paimon/Cargo.toml
@@ -103,6 +103,7 @@ urlencoding = "2.1"
tantivy = { version = "0.22", optional = true }
tempfile = { version = "3", optional = true }
paimon-mosaic-core = { version = "0.1.0", optional = true }
+paimon-vindex-core = "0.1.0"
vortex = { version = "0.68", features = ["tokio"], optional = true }
libloading = "0.9"
# Keep CI on the dependency set that passed before unicode-segmentation 1.13.3.
diff --git a/crates/paimon/src/lib.rs b/crates/paimon/src/lib.rs
index 1593122..7b5007f 100644
--- a/crates/paimon/src/lib.rs
+++ b/crates/paimon/src/lib.rs
@@ -37,6 +37,8 @@ pub mod spec;
pub mod table;
#[cfg(feature = "fulltext")]
pub mod tantivy;
+pub mod vector_search;
+pub mod vindex;
pub use catalog::Catalog;
pub use catalog::CatalogFactory;
diff --git a/crates/paimon/src/lumina/mod.rs b/crates/paimon/src/lumina/mod.rs
index 7986a53..69fb3d2 100644
--- a/crates/paimon/src/lumina/mod.rs
+++ b/crates/paimon/src/lumina/mod.rs
@@ -223,74 +223,6 @@ pub fn strip_lumina_options(paimon_options:
&HashMap<String, String>) -> HashMap
result
}
-#[derive(Clone)]
-pub struct VectorSearch {
- pub vector: Vec<f32>,
- pub limit: usize,
- pub field_name: String,
- pub include_row_ids: Option<roaring::RoaringTreemap>,
-}
-
-impl VectorSearch {
- pub fn new(vector: Vec<f32>, limit: usize, field_name: String) ->
crate::Result<Self> {
- if vector.is_empty() {
- return Err(crate::Error::DataInvalid {
- message: "Search vector cannot be empty".to_string(),
- source: None,
- });
- }
- if limit == 0 || limit > i32::MAX as usize {
- return Err(crate::Error::DataInvalid {
- message: format!("Limit must be between 1 and {}, got: {}",
i32::MAX, limit),
- source: None,
- });
- }
- if field_name.is_empty() {
- return Err(crate::Error::DataInvalid {
- message: "Field name cannot be null or empty".to_string(),
- source: None,
- });
- }
- Ok(Self {
- vector,
- limit,
- field_name,
- include_row_ids: None,
- })
- }
-
- pub fn with_include_row_ids(mut self, include_row_ids:
roaring::RoaringTreemap) -> Self {
- self.include_row_ids = Some(include_row_ids);
- self
- }
-}
-
-impl std::fmt::Display for VectorSearch {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- write!(
- f,
- "VectorSearch(field_name={}, limit={})",
- self.field_name, self.limit
- )
- }
-}
-
-pub struct GlobalIndexIOMeta {
- pub file_path: String,
- pub file_size: u64,
- pub metadata: Vec<u8>,
-}
-
-impl GlobalIndexIOMeta {
- pub fn new(file_path: String, file_size: u64, metadata: Vec<u8>) -> Self {
- Self {
- file_path,
- file_size,
- metadata,
- }
- }
-}
-
pub const KEY_DIMENSION: &str = "index.dimension";
pub const KEY_DISTANCE_METRIC: &str = "distance.metric";
pub const KEY_INDEX_TYPE: &str = "index.type";
@@ -375,126 +307,6 @@ impl LuminaIndexMeta {
}
}
-#[derive(Debug, Clone)]
-pub struct SearchResult {
- pub row_ids: Vec<u64>,
- pub scores: Vec<f32>,
-}
-
-impl SearchResult {
- pub fn new(row_ids: Vec<u64>, scores: Vec<f32>) -> Self {
- assert_eq!(row_ids.len(), scores.len());
- Self { row_ids, scores }
- }
-
- pub fn empty() -> Self {
- Self {
- row_ids: Vec::new(),
- scores: Vec::new(),
- }
- }
-
- pub fn from_scored_map(map: HashMap<u64, f32>) -> Self {
- let mut row_ids = Vec::with_capacity(map.len());
- let mut scores = Vec::with_capacity(map.len());
- for (id, score) in map {
- row_ids.push(id);
- scores.push(score);
- }
- Self { row_ids, scores }
- }
-
- pub fn len(&self) -> usize {
- self.row_ids.len()
- }
-
- pub fn is_empty(&self) -> bool {
- self.row_ids.is_empty()
- }
-
- pub fn offset(&self, offset: i64) -> Self {
- if offset == 0 {
- return self.clone();
- }
- let row_ids = self
- .row_ids
- .iter()
- .map(|&id| {
- if offset >= 0 {
- id.saturating_add(offset as u64)
- } else {
- id.saturating_sub(offset.unsigned_abs())
- }
- })
- .collect();
- Self {
- row_ids,
- scores: self.scores.clone(),
- }
- }
-
- pub fn or(&self, other: &SearchResult) -> Self {
- let mut row_ids = self.row_ids.clone();
- let mut scores = self.scores.clone();
- row_ids.extend_from_slice(&other.row_ids);
- scores.extend_from_slice(&other.scores);
- Self { row_ids, scores }
- }
-
- pub fn top_k(&self, k: usize) -> Self {
- if self.row_ids.len() <= k {
- return self.clone();
- }
- let mut indices: Vec<usize> = (0..self.row_ids.len()).collect();
- indices.sort_by(|&a, &b| {
- self.scores[b]
- .partial_cmp(&self.scores[a])
- .unwrap_or(std::cmp::Ordering::Equal)
- });
- indices.truncate(k);
- let row_ids = indices.iter().map(|&i| self.row_ids[i]).collect();
- let scores = indices.iter().map(|&i| self.scores[i]).collect();
- Self { row_ids, scores }
- }
-
- pub fn to_row_ranges(&self) -> crate::Result<Vec<crate::table::RowRange>> {
- if self.row_ids.is_empty() {
- return Ok(Vec::new());
- }
-
- let mut sorted = self
- .row_ids
- .iter()
- .copied()
- .map(|id| {
- i64::try_from(id).map_err(|_| crate::Error::DataInvalid {
- message: format!(
- "Lumina search row id {id} exceeds i64::MAX and cannot
be converted to RowRange"
- ),
- source: None,
- })
- })
- .collect::<crate::Result<Vec<_>>>()?;
-
- sorted.sort_unstable();
- sorted.dedup();
- let mut ranges = Vec::new();
- let mut start = sorted[0];
- let mut end = start;
- for &id in &sorted[1..] {
- if end.checked_add(1) == Some(id) {
- end = id;
- } else {
- ranges.push(crate::table::RowRange::new(start, end));
- start = id;
- end = id;
- }
- }
- ranges.push(crate::table::RowRange::new(start, end));
- Ok(ranges)
- }
-}
-
#[cfg(test)]
mod tests {
use super::*;
@@ -654,78 +466,4 @@ mod tests {
assert_eq!(lumina_opts.get("encoding.pq.m").unwrap(), "64");
assert_eq!(lumina_opts.get("search.parallel_number").unwrap(), "5");
}
-
- #[test]
- fn test_vector_search_clone_preserves_include_row_ids() {
- let mut include_row_ids = roaring::RoaringTreemap::new();
- include_row_ids.insert(1);
- include_row_ids.insert(3);
-
- let vector_search = VectorSearch::new(vec![1.0, 2.0], 10,
"embedding".to_string())
- .unwrap()
- .with_include_row_ids(include_row_ids.clone());
-
- let cloned = vector_search.clone();
- assert_eq!(cloned.vector, vector_search.vector);
- assert_eq!(cloned.limit, vector_search.limit);
- assert_eq!(cloned.field_name, vector_search.field_name);
- assert_eq!(cloned.include_row_ids.as_ref(), Some(&include_row_ids));
- }
-
- #[test]
- fn test_search_result_from_scored_map() {
- let mut map = HashMap::new();
- map.insert(1u64, 0.9f32);
- map.insert(2, 0.5);
- let result = SearchResult::from_scored_map(map);
- assert_eq!(result.len(), 2);
- }
-
- #[test]
- fn test_search_result_top_k() {
- let result = SearchResult::new(vec![1, 2, 3, 4, 5], vec![0.1, 0.9,
0.5, 0.8, 0.3]);
- let top = result.top_k(2);
- assert_eq!(top.len(), 2);
- assert!(top.row_ids.contains(&2));
- assert!(top.row_ids.contains(&4));
- }
-
- #[test]
- fn test_search_result_offset() {
- let result = SearchResult::new(vec![0, 1], vec![0.5, 0.6]);
- let offset = result.offset(100);
- assert_eq!(offset.row_ids, vec![100, 101]);
- assert_eq!(offset.scores, vec![0.5, 0.6]);
- }
-
- #[test]
- fn test_search_result_or() {
- let a = SearchResult::new(vec![1, 2], vec![0.5, 0.6]);
- let b = SearchResult::new(vec![3], vec![0.7]);
- let merged = a.or(&b);
- assert_eq!(merged.len(), 3);
- }
-
- #[test]
- fn test_search_result_to_row_ranges() {
- let result = SearchResult::new(vec![5, 1, 2, 3, 10], vec![0.1; 5]);
- let ranges = result.to_row_ranges().unwrap();
- assert_eq!(ranges.len(), 3);
- assert_eq!(ranges[0].from(), 1);
- assert_eq!(ranges[0].to(), 3);
- assert_eq!(ranges[1].from(), 5);
- assert_eq!(ranges[1].to(), 5);
- assert_eq!(ranges[2].from(), 10);
- assert_eq!(ranges[2].to(), 10);
- }
-
- #[test]
- fn test_search_result_to_row_ranges_rejects_i64_overflow() {
- let result = SearchResult::new(vec![i64::MAX as u64 + 1], vec![0.1]);
- let err = result.to_row_ranges().unwrap_err();
- assert!(
- err.to_string().contains("exceeds i64::MAX"),
- "unexpected error: {err}"
- );
- }
}
diff --git a/crates/paimon/src/lumina/reader.rs
b/crates/paimon/src/lumina/reader.rs
index 09acff8..331318b 100644
--- a/crates/paimon/src/lumina/reader.rs
+++ b/crates/paimon/src/lumina/reader.rs
@@ -16,9 +16,8 @@
// under the License.
use crate::lumina::ffi::LuminaSearcher;
-use crate::lumina::{
- strip_lumina_options, GlobalIndexIOMeta, LuminaIndexMeta,
LuminaVectorMetric, VectorSearch,
-};
+use crate::lumina::{strip_lumina_options, LuminaIndexMeta, LuminaVectorMetric};
+use crate::vector_search::{GlobalIndexIOMeta, VectorSearch};
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::io::{Read, Seek, SeekFrom};
@@ -146,71 +145,7 @@ impl LuminaVectorGlobalIndexReader {
source: None,
})?;
- let expected_dim = index_meta.dim()? as usize;
- if vector_search.vector.len() != expected_dim {
- return Err(crate::Error::DataInvalid {
- message: format!(
- "Query vector dimension mismatch: index expects {}, but
got {}",
- expected_dim,
- vector_search.vector.len()
- ),
- source: None,
- });
- }
-
- let limit = vector_search.limit;
- let index_metric = index_meta.metric()?;
- let count = searcher.get_count()? as usize;
- let effective_k = std::cmp::min(limit, count);
- if effective_k == 0 {
- return Ok(None);
- }
-
- let include_row_ids = &vector_search.include_row_ids;
-
- let (distances, labels) = if let Some(ref include_ids) =
include_row_ids {
- let filter_id_list: Vec<u64> = include_ids.iter().collect();
- if filter_id_list.is_empty() {
- return Ok(None);
- }
- let ek = std::cmp::min(effective_k, filter_id_list.len());
- let mut distances = vec![0.0f32; ek];
- let mut labels = vec![0u64; ek];
- let mut search_opts: HashMap<String, String> =
search_options_base.clone();
- search_opts.insert("search.thread_safe_filter".to_string(),
"true".to_string());
- ensure_search_list_size(&mut search_opts, ek);
- searcher.search_with_filter(
- &vector_search.vector,
- 1,
- ek as i32,
- &mut distances,
- &mut labels,
- &filter_id_list,
- &search_opts,
- )?;
- (distances, labels)
- } else {
- let mut distances = vec![0.0f32; effective_k];
- let mut labels = vec![0u64; effective_k];
- let mut search_opts: HashMap<String, String> =
search_options_base.clone();
- ensure_search_list_size(&mut search_opts, effective_k);
- searcher.search(
- &vector_search.vector,
- 1,
- effective_k as i32,
- &mut distances,
- &mut labels,
- &search_opts,
- )?;
- (distances, labels)
- };
-
- let id_to_scores = collect_results(&labels, &distances, effective_k,
index_metric);
- if id_to_scores.is_empty() {
- return Ok(None);
- }
-
- Ok(Some(id_to_scores))
+ search_lumina(searcher, index_meta, search_options_base, vector_search)
}
fn ensure_loaded<S: Read + Seek + Send + 'static>(
@@ -264,6 +199,79 @@ impl LuminaVectorGlobalIndexReader {
}
}
+fn search_lumina(
+ searcher: &LuminaSearcher,
+ index_meta: &LuminaIndexMeta,
+ search_options_base: &HashMap<String, String>,
+ vector_search: &VectorSearch,
+) -> crate::Result<Option<HashMap<u64, f32>>> {
+ let expected_dim = index_meta.dim()? as usize;
+ if vector_search.vector.len() != expected_dim {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Query vector dimension mismatch: index expects {}, but got
{}",
+ expected_dim,
+ vector_search.vector.len()
+ ),
+ source: None,
+ });
+ }
+
+ let limit = vector_search.limit;
+ let index_metric = index_meta.metric()?;
+ let count = searcher.get_count()? as usize;
+ let effective_k = std::cmp::min(limit, count);
+ if effective_k == 0 {
+ return Ok(None);
+ }
+
+ let include_row_ids = &vector_search.include_row_ids;
+
+ let (distances, labels) = if let Some(ref include_ids) = include_row_ids {
+ let filter_id_list: Vec<u64> = include_ids.iter().collect();
+ if filter_id_list.is_empty() {
+ return Ok(None);
+ }
+ let ek = std::cmp::min(effective_k, filter_id_list.len());
+ let mut distances = vec![0.0f32; ek];
+ let mut labels = vec![0u64; ek];
+ let mut search_opts: HashMap<String, String> =
search_options_base.clone();
+ search_opts.insert("search.thread_safe_filter".to_string(),
"true".to_string());
+ ensure_search_list_size(&mut search_opts, ek);
+ searcher.search_with_filter(
+ &vector_search.vector,
+ 1,
+ ek as i32,
+ &mut distances,
+ &mut labels,
+ &filter_id_list,
+ &search_opts,
+ )?;
+ (distances, labels)
+ } else {
+ let mut distances = vec![0.0f32; effective_k];
+ let mut labels = vec![0u64; effective_k];
+ let mut search_opts: HashMap<String, String> =
search_options_base.clone();
+ ensure_search_list_size(&mut search_opts, effective_k);
+ searcher.search(
+ &vector_search.vector,
+ 1,
+ effective_k as i32,
+ &mut distances,
+ &mut labels,
+ &search_opts,
+ )?;
+ (distances, labels)
+ };
+
+ let id_to_scores = collect_results(&labels, &distances, effective_k,
index_metric);
+ if id_to_scores.is_empty() {
+ return Ok(None);
+ }
+
+ Ok(Some(id_to_scores))
+}
+
fn write_temp_index_file<S: Read + Seek>(stream: &mut S) ->
crate::Result<PathBuf> {
stream
.seek(SeekFrom::Start(0))
@@ -312,7 +320,7 @@ impl Drop for LuminaVectorGlobalIndexReader {
#[cfg(test)]
mod tests {
use super::*;
- use crate::lumina::GlobalIndexIOMeta;
+ use crate::vector_search::GlobalIndexIOMeta;
use std::io::Cursor;
#[test]
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index 887b230..48271aa 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -15,16 +15,44 @@
// specific language governing permissions and limitations
// under the License.
+use crate::lumina::is_lumina_index_type;
use crate::lumina::reader::LuminaVectorGlobalIndexReader;
-use crate::lumina::{is_lumina_index_type, GlobalIndexIOMeta, SearchResult,
VectorSearch};
use crate::spec::{DataField, FileKind, IndexManifest};
use crate::table::snapshot_manager::SnapshotManager;
use crate::table::{find_field_id_by_name, RowRange, Table};
+use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch};
+use crate::vindex::is_vindex_index_type;
+use crate::vindex::reader::VindexVectorGlobalIndexReader;
use std::collections::HashMap;
use std::io::Cursor;
const INDEX_DIR: &str = "index";
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum VectorIndexBackend {
+ Lumina,
+ Vindex,
+}
+
+impl VectorIndexBackend {
+ fn from_index_type(index_type: &str) -> Option<Self> {
+ if is_lumina_index_type(index_type) {
+ Some(Self::Lumina)
+ } else if is_vindex_index_type(index_type) {
+ Some(Self::Vindex)
+ } else {
+ None
+ }
+ }
+
+ fn error_name(self) -> &'static str {
+ match self {
+ Self::Lumina => "Lumina",
+ Self::Vindex => "vindex",
+ }
+ }
+}
+
pub struct VectorSearchBuilder<'a> {
table: &'a Table,
vector_column: Option<String>,
@@ -126,11 +154,11 @@ async fn evaluate_vector_search(
None => return Ok(Vec::new()),
};
- let lumina_entries: Vec<_> = index_entries
+ let vector_entries: Vec<_> = index_entries
.iter()
.filter(|e| {
e.kind == FileKind::Add
- && is_lumina_index_type(&e.index_file.index_type)
+ &&
VectorIndexBackend::from_index_type(&e.index_file.index_type).is_some()
&& e.index_file
.global_index_meta
.as_ref()
@@ -138,14 +166,16 @@ async fn evaluate_vector_search(
})
.collect();
- if lumina_entries.is_empty() {
+ if vector_entries.is_empty() {
return Ok(Vec::new());
}
- let futures: Vec<_> = lumina_entries
+ let futures: Vec<_> = vector_entries
.into_iter()
.map(|entry| {
let global_meta =
entry.index_file.global_index_meta.as_ref().unwrap();
+ let backend =
VectorIndexBackend::from_index_type(&entry.index_file.index_type)
+ .expect("filtered vector index type");
let path = format!("{table_path}/{INDEX_DIR}/{}",
entry.index_file.file_name);
let file_name = entry.index_file.file_name.clone();
let file_size = entry.index_file.file_size as u64;
@@ -157,16 +187,30 @@ async fn evaluate_vector_search(
async move {
let input = input?;
let bytes = input.read().await.map_err(|e|
crate::Error::DataInvalid {
- message: format!("Failed to read Lumina index file '{}':
{}", file_name, e),
+ message: format!(
+ "Failed to read {} index file '{}': {}",
+ backend.error_name(),
+ file_name,
+ e
+ ),
source: None,
})?;
let io_meta =
GlobalIndexIOMeta::new(file_name.clone(), file_size,
index_meta_bytes);
- let mut reader = LuminaVectorGlobalIndexReader::new(io_meta,
options);
let data = bytes.to_vec();
- let result =
- reader.visit_vector_search(&vector_search_clone, |_|
Ok(Cursor::new(data)))?;
+ let result = match backend {
+ VectorIndexBackend::Lumina => {
+ let mut reader =
LuminaVectorGlobalIndexReader::new(io_meta, options);
+ reader
+ .visit_vector_search(&vector_search_clone, |_|
Ok(Cursor::new(data)))?
+ }
+ VectorIndexBackend::Vindex => {
+ let mut reader =
VindexVectorGlobalIndexReader::new(io_meta, options);
+ reader
+ .visit_vector_search(&vector_search_clone, |_|
Ok(Cursor::new(data)))?
+ }
+ };
match result {
Some(scored_map) => Ok::<_, crate::Error>(
@@ -192,6 +236,7 @@ mod tests {
use super::*;
use crate::lumina::{LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER,
LUMINA_IDENTIFIER};
use crate::spec::{DataType, GlobalIndexMeta, IndexFileMeta,
IndexManifestEntry, IntType};
+ use crate::vindex::IVF_FLAT_IDENTIFIER;
fn make_field(id: i32, name: &str) -> DataField {
DataField::new(id, name.to_string(), DataType::Int(IntType::default()))
@@ -239,7 +284,7 @@ mod tests {
}
#[tokio::test]
- async fn test_evaluate_ignores_non_lumina_index_type() {
+ async fn test_evaluate_ignores_non_vector_index_type() {
let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap();
let fields = vec![make_field(2, "embedding")];
let vs = VectorSearch::new(vec![1.0], 10,
"embedding".to_string()).unwrap();
@@ -366,6 +411,31 @@ mod tests {
);
}
+ #[tokio::test]
+ async fn test_evaluate_accepts_vindex_index_type() {
+ let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap();
+ let fields = vec![make_field(2, "embedding")];
+ let vs = VectorSearch::new(vec![1.0], 10,
"embedding".to_string()).unwrap();
+
+ let entry = make_lumina_entry("missing.idx", IVF_FLAT_IDENTIFIER,
FileKind::Add, 2);
+
+ let err = evaluate_vector_search(
+ &file_io,
+ "memory:///test_table",
+ &HashMap::new(),
+ &[entry],
+ &vs,
+ &fields,
+ )
+ .await
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("Failed to read vindex index file 'missing.idx'"),
+ "unexpected error: {err}"
+ );
+ }
+
fn make_lumina_entry(
file_name: &str,
index_type: &str,
diff --git a/crates/paimon/src/vector_search.rs
b/crates/paimon/src/vector_search.rs
new file mode 100644
index 0000000..491b65f
--- /dev/null
+++ b/crates/paimon/src/vector_search.rs
@@ -0,0 +1,285 @@
+// 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;
+
+#[derive(Clone)]
+pub struct VectorSearch {
+ pub vector: Vec<f32>,
+ pub limit: usize,
+ pub field_name: String,
+ pub include_row_ids: Option<roaring::RoaringTreemap>,
+}
+
+impl VectorSearch {
+ pub fn new(vector: Vec<f32>, limit: usize, field_name: String) ->
crate::Result<Self> {
+ if vector.is_empty() {
+ return Err(crate::Error::DataInvalid {
+ message: "Search vector cannot be empty".to_string(),
+ source: None,
+ });
+ }
+ if limit == 0 || limit > i32::MAX as usize {
+ return Err(crate::Error::DataInvalid {
+ message: format!("Limit must be between 1 and {}, got: {}",
i32::MAX, limit),
+ source: None,
+ });
+ }
+ if field_name.is_empty() {
+ return Err(crate::Error::DataInvalid {
+ message: "Field name cannot be null or empty".to_string(),
+ source: None,
+ });
+ }
+ Ok(Self {
+ vector,
+ limit,
+ field_name,
+ include_row_ids: None,
+ })
+ }
+
+ pub fn with_include_row_ids(mut self, include_row_ids:
roaring::RoaringTreemap) -> Self {
+ self.include_row_ids = Some(include_row_ids);
+ self
+ }
+}
+
+impl std::fmt::Display for VectorSearch {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "VectorSearch(field_name={}, limit={})",
+ self.field_name, self.limit
+ )
+ }
+}
+
+pub struct GlobalIndexIOMeta {
+ pub file_path: String,
+ pub file_size: u64,
+ pub metadata: Vec<u8>,
+}
+
+impl GlobalIndexIOMeta {
+ pub fn new(file_path: String, file_size: u64, metadata: Vec<u8>) -> Self {
+ Self {
+ file_path,
+ file_size,
+ metadata,
+ }
+ }
+}
+
+#[derive(Debug, Clone)]
+pub struct SearchResult {
+ pub row_ids: Vec<u64>,
+ pub scores: Vec<f32>,
+}
+
+impl SearchResult {
+ pub fn new(row_ids: Vec<u64>, scores: Vec<f32>) -> Self {
+ assert_eq!(row_ids.len(), scores.len());
+ Self { row_ids, scores }
+ }
+
+ pub fn empty() -> Self {
+ Self {
+ row_ids: Vec::new(),
+ scores: Vec::new(),
+ }
+ }
+
+ pub fn from_scored_map(map: HashMap<u64, f32>) -> Self {
+ let mut row_ids = Vec::with_capacity(map.len());
+ let mut scores = Vec::with_capacity(map.len());
+ for (id, score) in map {
+ row_ids.push(id);
+ scores.push(score);
+ }
+ Self { row_ids, scores }
+ }
+
+ pub fn len(&self) -> usize {
+ self.row_ids.len()
+ }
+
+ pub fn is_empty(&self) -> bool {
+ self.row_ids.is_empty()
+ }
+
+ pub fn offset(&self, offset: i64) -> Self {
+ if offset == 0 {
+ return self.clone();
+ }
+ let row_ids = self
+ .row_ids
+ .iter()
+ .map(|&id| {
+ if offset >= 0 {
+ id.saturating_add(offset as u64)
+ } else {
+ id.saturating_sub(offset.unsigned_abs())
+ }
+ })
+ .collect();
+ Self {
+ row_ids,
+ scores: self.scores.clone(),
+ }
+ }
+
+ pub fn or(&self, other: &SearchResult) -> Self {
+ let mut row_ids = self.row_ids.clone();
+ let mut scores = self.scores.clone();
+ row_ids.extend_from_slice(&other.row_ids);
+ scores.extend_from_slice(&other.scores);
+ Self { row_ids, scores }
+ }
+
+ pub fn top_k(&self, k: usize) -> Self {
+ if self.row_ids.len() <= k {
+ return self.clone();
+ }
+ let mut indices: Vec<usize> = (0..self.row_ids.len()).collect();
+ indices.sort_by(|&a, &b| {
+ self.scores[b]
+ .partial_cmp(&self.scores[a])
+ .unwrap_or(std::cmp::Ordering::Equal)
+ });
+ indices.truncate(k);
+ let row_ids = indices.iter().map(|&i| self.row_ids[i]).collect();
+ let scores = indices.iter().map(|&i| self.scores[i]).collect();
+ Self { row_ids, scores }
+ }
+
+ pub fn to_row_ranges(&self) -> crate::Result<Vec<crate::table::RowRange>> {
+ if self.row_ids.is_empty() {
+ return Ok(Vec::new());
+ }
+
+ let mut sorted = self
+ .row_ids
+ .iter()
+ .copied()
+ .map(|id| {
+ i64::try_from(id).map_err(|_| crate::Error::DataInvalid {
+ message: format!(
+ "Vector search row id {id} exceeds i64::MAX and cannot
be converted to RowRange"
+ ),
+ source: None,
+ })
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
+
+ sorted.sort_unstable();
+ sorted.dedup();
+ let mut ranges = Vec::new();
+ let mut start = sorted[0];
+ let mut end = start;
+ for &id in &sorted[1..] {
+ if end.checked_add(1) == Some(id) {
+ end = id;
+ } else {
+ ranges.push(crate::table::RowRange::new(start, end));
+ start = id;
+ end = id;
+ }
+ }
+ ranges.push(crate::table::RowRange::new(start, end));
+ Ok(ranges)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_vector_search_clone_preserves_include_row_ids() {
+ let mut include_row_ids = roaring::RoaringTreemap::new();
+ include_row_ids.insert(1);
+ include_row_ids.insert(3);
+
+ let vector_search = VectorSearch::new(vec![1.0, 2.0], 10,
"embedding".to_string())
+ .unwrap()
+ .with_include_row_ids(include_row_ids.clone());
+
+ let cloned = vector_search.clone();
+ assert_eq!(cloned.vector, vector_search.vector);
+ assert_eq!(cloned.limit, vector_search.limit);
+ assert_eq!(cloned.field_name, vector_search.field_name);
+ assert_eq!(cloned.include_row_ids.as_ref(), Some(&include_row_ids));
+ }
+
+ #[test]
+ fn test_search_result_from_scored_map() {
+ let mut map = HashMap::new();
+ map.insert(1u64, 0.9f32);
+ map.insert(2, 0.5);
+ let result = SearchResult::from_scored_map(map);
+ assert_eq!(result.len(), 2);
+ }
+
+ #[test]
+ fn test_search_result_top_k() {
+ let result = SearchResult::new(vec![1, 2, 3, 4, 5], vec![0.1, 0.9,
0.5, 0.8, 0.3]);
+ let top = result.top_k(2);
+ assert_eq!(top.len(), 2);
+ assert!(top.row_ids.contains(&2));
+ assert!(top.row_ids.contains(&4));
+ }
+
+ #[test]
+ fn test_search_result_offset() {
+ let result = SearchResult::new(vec![0, 1], vec![0.5, 0.6]);
+ let offset = result.offset(100);
+ assert_eq!(offset.row_ids, vec![100, 101]);
+ assert_eq!(offset.scores, vec![0.5, 0.6]);
+ }
+
+ #[test]
+ fn test_search_result_or() {
+ let a = SearchResult::new(vec![1, 2], vec![0.5, 0.6]);
+ let b = SearchResult::new(vec![3], vec![0.7]);
+ let merged = a.or(&b);
+ assert_eq!(merged.len(), 3);
+ }
+
+ #[test]
+ fn test_search_result_to_row_ranges() {
+ let result = SearchResult::new(vec![5, 1, 2, 3, 10], vec![0.1; 5]);
+ let ranges = result.to_row_ranges().unwrap();
+ assert_eq!(ranges.len(), 3);
+ assert_eq!(ranges[0].from(), 1);
+ assert_eq!(ranges[0].to(), 3);
+ assert_eq!(ranges[1].from(), 5);
+ assert_eq!(ranges[1].to(), 5);
+ assert_eq!(ranges[2].from(), 10);
+ assert_eq!(ranges[2].to(), 10);
+ }
+
+ #[test]
+ fn test_search_result_to_row_ranges_rejects_i64_overflow() {
+ let result = SearchResult::new(vec![i64::MAX as u64 + 1], vec![0.1]);
+ let err = result.to_row_ranges().unwrap_err();
+ assert!(
+ err.to_string().contains("exceeds i64::MAX"),
+ "unexpected error: {err}"
+ );
+ }
+}
diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs
new file mode 100644
index 0000000..aa4f7b8
--- /dev/null
+++ b/crates/paimon/src/vindex/mod.rs
@@ -0,0 +1,47 @@
+// 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.
+
+pub mod reader;
+
+pub const IVF_FLAT_IDENTIFIER: &str = "ivf-flat";
+pub const IVF_PQ_IDENTIFIER: &str = "ivf-pq";
+pub const IVF_HNSW_FLAT_IDENTIFIER: &str = "ivf-hnsw-flat";
+pub const IVF_HNSW_SQ_IDENTIFIER: &str = "ivf-hnsw-sq";
+
+pub fn is_vindex_index_type(index_type: &str) -> bool {
+ matches!(
+ index_type,
+ IVF_FLAT_IDENTIFIER | IVF_PQ_IDENTIFIER | IVF_HNSW_FLAT_IDENTIFIER |
IVF_HNSW_SQ_IDENTIFIER
+ )
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_vindex_index_type_identifier_helper() {
+ assert!(is_vindex_index_type(IVF_FLAT_IDENTIFIER));
+ assert!(is_vindex_index_type(IVF_PQ_IDENTIFIER));
+ assert!(is_vindex_index_type(IVF_HNSW_FLAT_IDENTIFIER));
+ assert!(is_vindex_index_type(IVF_HNSW_SQ_IDENTIFIER));
+ assert!(!is_vindex_index_type(""));
+ assert!(!is_vindex_index_type("btree"));
+ assert!(!is_vindex_index_type("lumina"));
+ assert!(!is_vindex_index_type("IVF-FLAT"));
+ }
+}
diff --git a/crates/paimon/src/vindex/reader.rs
b/crates/paimon/src/vindex/reader.rs
new file mode 100644
index 0000000..18a1fa7
--- /dev/null
+++ b/crates/paimon/src/vindex/reader.rs
@@ -0,0 +1,322 @@
+// 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 crate::vector_search::{GlobalIndexIOMeta, VectorSearch};
+use paimon_vindex_core::distance::MetricType;
+use paimon_vindex_core::index::{
+ VectorIndexMetadata, VectorIndexReader as VIndexReader, VectorSearchParams,
+};
+use std::collections::BinaryHeap;
+use std::collections::HashMap;
+use std::io::{Cursor, Read, Seek, SeekFrom};
+
+const DEFAULT_NPROBE: usize = 16;
+const DEFAULT_EF_SEARCH: usize = 0;
+const NPROBE_PARAMETER: &str = "ivf.nprobe";
+const EF_SEARCH_PARAMETER: &str = "hnsw.ef_search";
+
+pub struct VindexVectorGlobalIndexReader {
+ io_meta: GlobalIndexIOMeta,
+ options: HashMap<String, String>,
+ reader: Option<VIndexReader<Cursor<Vec<u8>>>>,
+ metadata: Option<VectorIndexMetadata>,
+}
+
+impl VindexVectorGlobalIndexReader {
+ pub fn new(io_meta: GlobalIndexIOMeta, options: HashMap<String, String>)
-> Self {
+ Self {
+ io_meta,
+ options,
+ reader: None,
+ metadata: None,
+ }
+ }
+
+ pub fn visit_vector_search<S: Read + Seek + Send + 'static>(
+ &mut self,
+ vector_search: &VectorSearch,
+ stream_fn: impl FnOnce(&str) -> crate::Result<S>,
+ ) -> crate::Result<Option<HashMap<u64, f32>>> {
+ self.ensure_loaded(stream_fn)?;
+ self.search(vector_search)
+ }
+
+ fn search(&mut self, vector_search: &VectorSearch) ->
crate::Result<Option<HashMap<u64, f32>>> {
+ let reader = self
+ .reader
+ .as_mut()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "vindex reader not initialized".to_string(),
+ source: None,
+ })?;
+ let metadata = self
+ .metadata
+ .as_ref()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "vindex metadata not initialized".to_string(),
+ source: None,
+ })?;
+
+ search_vindex(reader, metadata, &self.options, vector_search)
+ }
+
+ fn ensure_loaded<S: Read + Seek + Send + 'static>(
+ &mut self,
+ stream_fn: impl FnOnce(&str) -> crate::Result<S>,
+ ) -> crate::Result<()> {
+ if self.reader.is_some() {
+ return Ok(());
+ }
+
+ let mut stream = stream_fn(&self.io_meta.file_path)?;
+ stream
+ .seek(SeekFrom::Start(0))
+ .map_err(|e| crate::Error::UnexpectedError {
+ message: format!("Failed to seek vindex stream to start: {}",
e),
+ source: Some(Box::new(e)),
+ })?;
+ let mut bytes = Vec::with_capacity(self.io_meta.file_size as usize);
+ stream
+ .read_to_end(&mut bytes)
+ .map_err(|e| crate::Error::UnexpectedError {
+ message: format!("Failed to read vindex stream: {}", e),
+ source: Some(Box::new(e)),
+ })?;
+
+ let mut reader =
+ VIndexReader::open(Cursor::new(bytes)).map_err(|e|
crate::Error::DataInvalid {
+ message: format!("Failed to open paimon-vindex-core reader:
{}", e),
+ source: Some(Box::new(e)),
+ })?;
+ let metadata = reader.metadata();
+ reader
+ .optimize_for_search()
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("Failed to optimize paimon-vindex-core
reader: {}", e),
+ source: Some(Box::new(e)),
+ })?;
+
+ self.reader = Some(reader);
+ self.metadata = Some(metadata);
+ Ok(())
+ }
+}
+
+fn search_vindex(
+ reader: &mut VIndexReader<Cursor<Vec<u8>>>,
+ metadata: &VectorIndexMetadata,
+ options: &HashMap<String, String>,
+ vector_search: &VectorSearch,
+) -> crate::Result<Option<HashMap<u64, f32>>> {
+ let expected_dim = metadata.dimension;
+ if vector_search.vector.len() != expected_dim {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Query vector dimension mismatch: index expects {}, but got
{}",
+ expected_dim,
+ vector_search.vector.len()
+ ),
+ source: None,
+ });
+ }
+
+ let count = usize::try_from(metadata.total_vectors).unwrap_or(0);
+ let effective_k = std::cmp::min(vector_search.limit, count);
+ if effective_k == 0 {
+ return Ok(None);
+ }
+
+ let params = VectorSearchParams::with_ef_search(
+ effective_k,
+ int_parameter(options, NPROBE_PARAMETER, DEFAULT_NPROBE)?,
+ int_parameter(options, EF_SEARCH_PARAMETER, DEFAULT_EF_SEARCH)?,
+ );
+
+ let (labels, distances) = if let Some(include_ids) =
&vector_search.include_row_ids {
+ if include_ids.is_empty() {
+ return Ok(None);
+ }
+ let ek = std::cmp::min(effective_k, include_ids.len() as usize);
+ let params = VectorSearchParams::with_ef_search(
+ params.top_k.min(ek),
+ params.nprobe,
+ params.ef_search,
+ );
+ let mut filter_bytes = Vec::new();
+ include_ids
+ .serialize_into(&mut filter_bytes)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("Failed to serialize vector search row-id
filter: {}", e),
+ source: Some(Box::new(e)),
+ })?;
+ reader
+ .search_with_roaring_filter(&vector_search.vector, params,
&filter_bytes)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core filtered search failed:
{}", e),
+ source: Some(Box::new(e)),
+ })?
+ } else {
+ reader
+ .search(&vector_search.vector, params)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core search failed: {}", e),
+ source: Some(Box::new(e)),
+ })?
+ };
+
+ let id_to_scores = collect_results(&labels, &distances, effective_k,
metadata.metric);
+ if id_to_scores.is_empty() {
+ return Ok(None);
+ }
+
+ Ok(Some(id_to_scores))
+}
+
+fn collect_results(
+ labels: &[i64],
+ distances: &[f32],
+ top_k: usize,
+ metric: MetricType,
+) -> HashMap<u64, f32> {
+ #[derive(PartialEq)]
+ struct ScoredRow {
+ row_id: u64,
+ score: f32,
+ }
+ impl Eq for ScoredRow {}
+ impl PartialOrd for ScoredRow {
+ fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+ Some(self.cmp(other))
+ }
+ }
+ impl Ord for ScoredRow {
+ fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+ other.score.total_cmp(&self.score)
+ }
+ }
+
+ let mut min_heap: BinaryHeap<ScoredRow> = BinaryHeap::with_capacity(top_k
+ 1);
+ for (&row_id, &distance) in labels.iter().zip(distances.iter()) {
+ if row_id < 0 {
+ continue;
+ }
+ let score = convert_distance_to_score(distance, metric);
+ let row_id = row_id as u64;
+ if min_heap.len() < top_k {
+ min_heap.push(ScoredRow { row_id, score });
+ } else if let Some(peek) = min_heap.peek() {
+ if score > peek.score {
+ min_heap.pop();
+ min_heap.push(ScoredRow { row_id, score });
+ }
+ }
+ }
+
+ let mut result = HashMap::with_capacity(min_heap.len());
+ for entry in min_heap {
+ result.insert(entry.row_id, entry.score);
+ }
+ result
+}
+
+fn convert_distance_to_score(distance: f32, metric: MetricType) -> f32 {
+ match metric {
+ MetricType::L2 => 1.0 / (1.0 + distance),
+ MetricType::Cosine => 1.0 - distance,
+ MetricType::InnerProduct => -distance,
+ }
+}
+
+fn int_parameter(
+ options: &HashMap<String, String>,
+ key: &str,
+ default_value: usize,
+) -> crate::Result<usize> {
+ match options.get(key) {
+ Some(value) => value
+ .parse::<usize>()
+ .map_err(|_| crate::Error::DataInvalid {
+ message: format!(
+ "Invalid value for '{}': {}. Must be a non-negative
integer.",
+ key, value
+ ),
+ source: None,
+ }),
+ None => Ok(default_value),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_convert_distance_to_score() {
+ assert_eq!(convert_distance_to_score(0.0, MetricType::L2), 1.0);
+ assert_eq!(convert_distance_to_score(1.0, MetricType::L2), 0.5);
+ assert_eq!(convert_distance_to_score(0.0, MetricType::Cosine), 1.0);
+ assert_eq!(convert_distance_to_score(1.0, MetricType::Cosine), 0.0);
+ assert_eq!(
+ convert_distance_to_score(-0.75, MetricType::InnerProduct),
+ 0.75
+ );
+ }
+
+ #[test]
+ fn test_collect_results_converts_inner_product_distance_to_similarity() {
+ let labels = vec![9, 5, 1];
+ let distances = vec![-0.9, -0.5, -0.1];
+
+ let result = collect_results(&labels, &distances, 2,
MetricType::InnerProduct);
+
+ assert_eq!(result.len(), 2);
+ assert!(result.contains_key(&9), "0.9 similarity should be retained");
+ assert!(result.contains_key(&5), "0.5 similarity should be retained");
+ assert!(!result.contains_key(&1), "0.1 similarity should be trimmed");
+ assert_eq!(result.get(&9), Some(&0.9));
+ assert_eq!(result.get(&5), Some(&0.5));
+ }
+
+ #[test]
+ fn test_collect_results_skips_negative_labels() {
+ let labels = vec![0, -1, 2, 3];
+ let distances = vec![0.5, 0.0, 0.1, 0.9];
+ let result = collect_results(&labels, &distances, 2, MetricType::L2);
+ assert_eq!(result.len(), 2);
+ assert!(result.contains_key(&2));
+ assert!(result.contains_key(&0));
+ assert!(!result.contains_key(&3));
+ }
+
+ #[test]
+ fn test_int_parameter() {
+ let mut options = HashMap::new();
+ options.insert(NPROBE_PARAMETER.to_string(), "32".to_string());
+
+ assert_eq!(
+ int_parameter(&options, NPROBE_PARAMETER, DEFAULT_NPROBE).unwrap(),
+ 32
+ );
+ assert_eq!(
+ int_parameter(&options, EF_SEARCH_PARAMETER,
DEFAULT_EF_SEARCH).unwrap(),
+ DEFAULT_EF_SEARCH
+ );
+
+ options.insert(EF_SEARCH_PARAMETER.to_string(), "abc".to_string());
+ assert!(int_parameter(&options, EF_SEARCH_PARAMETER,
DEFAULT_EF_SEARCH).is_err());
+ }
+}
