From cb4c349b645c74533e9937518f830503791398a4 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 13 Sep 2026 22:08:03 +0800 Subject: [PATCH 1/3] refactor(table): unify vector search scan and read APIs --- bindings/c/src/bucket_vector_search_split.rs | 78 + bindings/c/src/lib.rs | 3 + bindings/c/src/result.rs | 24 + bindings/c/src/tests.rs | 295 +- bindings/c/src/types.rs | 29 + bindings/c/src/vector_read.rs | 71 + bindings/c/src/vector_scan.rs | 149 + bindings/c/src/vector_search.rs | 272 +- .../datafusion/src/lateral_vector_search.rs | 14 +- .../datafusion/src/vector_search.rs | 5 +- .../src/table/batch_vector_search_builder.rs | 170 + .../batch_vector_search_builder/tests.rs | 157 + crates/paimon/src/table/de_vector_read.rs | 2091 ++++ .../paimon/src/table/de_vector_read/tests.rs | 1392 +++ crates/paimon/src/table/de_vector_scan.rs | 293 + .../paimon/src/table/de_vector_scan/tests.rs | 364 + .../paimon/src/table/hybrid_search_builder.rs | 120 +- crates/paimon/src/table/mod.rs | 19 +- crates/paimon/src/table/pk_full_text_read.rs | 2 +- crates/paimon/src/table/pk_search_position.rs | 41 +- .../src/table/pk_vector_data_file_reader.rs | 2 +- .../src/table/pk_vector_indexed_split_read.rs | 1 + .../src/table/pk_vector_orchestrator.rs | 4 + crates/paimon/src/table/pk_vector_read.rs | 1026 ++ .../residual_positions_tests.rs | 536 + .../paimon/src/table/pk_vector_read/tests.rs | 1256 +++ crates/paimon/src/table/pk_vector_scan.rs | 102 +- .../src/table/pk_vector_search_params.rs | 218 + .../table/pk_vector_search_params/tests.rs | 41 + crates/paimon/src/table/vector_read.rs | 117 + crates/paimon/src/table/vector_scan.rs | 201 + .../paimon/src/table/vector_search_builder.rs | 8590 +---------------- .../src/table/vector_search_builder/tests.rs | 678 ++ .../paimon/src/table/vector_search_common.rs | 401 + .../paimon/src/table/vector_search_result.rs | 231 + .../src/table/vector_search_result/tests.rs | 211 + .../src/table/vector_search_test_utils.rs | 241 + crates/paimon/src/vector_search.rs | 38 +- crates/paimon/src/vindex/executor.rs | 5 +- crates/paimon/src/vindex/pkvector/bucket.rs | 7 +- .../paimon/tests/pk_vector_baseline_test.rs | 109 +- crates/paimon/tests/pk_vector_batch_test.rs | 132 +- .../tests/pk_vector_bucket_split_read_test.rs | 133 +- .../tests/pk_vector_java_fixture_test.rs | 69 +- docs/src/c-binding.md | 312 + 45 files changed, 11117 insertions(+), 9133 deletions(-) create mode 100644 bindings/c/src/bucket_vector_search_split.rs create mode 100644 bindings/c/src/vector_read.rs create mode 100644 bindings/c/src/vector_scan.rs create mode 100644 crates/paimon/src/table/batch_vector_search_builder.rs create mode 100644 crates/paimon/src/table/batch_vector_search_builder/tests.rs create mode 100644 crates/paimon/src/table/de_vector_read.rs create mode 100644 crates/paimon/src/table/de_vector_read/tests.rs create mode 100644 crates/paimon/src/table/de_vector_scan.rs create mode 100644 crates/paimon/src/table/de_vector_scan/tests.rs create mode 100644 crates/paimon/src/table/pk_vector_read.rs create mode 100644 crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs create mode 100644 crates/paimon/src/table/pk_vector_read/tests.rs create mode 100644 crates/paimon/src/table/pk_vector_search_params.rs create mode 100644 crates/paimon/src/table/pk_vector_search_params/tests.rs create mode 100644 crates/paimon/src/table/vector_read.rs create mode 100644 crates/paimon/src/table/vector_scan.rs create mode 100644 crates/paimon/src/table/vector_search_builder/tests.rs create mode 100644 crates/paimon/src/table/vector_search_common.rs create mode 100644 crates/paimon/src/table/vector_search_result.rs create mode 100644 crates/paimon/src/table/vector_search_result/tests.rs create mode 100644 crates/paimon/src/table/vector_search_test_utils.rs diff --git a/bindings/c/src/bucket_vector_search_split.rs b/bindings/c/src/bucket_vector_search_split.rs new file mode 100644 index 000000000..b22eb565f --- /dev/null +++ b/bindings/c/src/bucket_vector_search_split.rs @@ -0,0 +1,78 @@ +// 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. + +//! Decodes the standalone Java PKVSPLIT format outside search execution. +use crate::error::{paimon_error, PaimonErrorCode}; +use crate::result::paimon_result_bucket_vector_search_split; +use crate::types::paimon_bucket_vector_search_split; +use paimon::table::BucketVectorSearchSplit; +use std::ffi::c_void; + +/// Decode one Java BucketVectorSearchSplit.serialize buffer. The returned handle +/// owns the decoded metadata; input bytes may be released immediately. This does +/// not accept Java ObjectOutputStream envelopes or DE index/raw split formats. +/// Free the handle with paimon_bucket_vector_search_split_free. +/// +/// # Safety +/// data must point to len readable bytes. Null/empty/oversized buffers return an error. +#[no_mangle] +pub unsafe extern "C" fn paimon_bucket_vector_search_split_deserialize( + data: *const u8, + len: usize, +) -> paimon_result_bucket_vector_search_split { + if data.is_null() || len == 0 || len > isize::MAX as usize { + return paimon_result_bucket_vector_search_split { + split: std::ptr::null_mut(), + error: paimon_error::new( + PaimonErrorCode::InvalidInput, + "null, empty or oversized bucket split buffer".to_string(), + ), + }; + } + match BucketVectorSearchSplit::deserialize(std::slice::from_raw_parts(data, len)) { + Ok(split) => paimon_result_bucket_vector_search_split { + split: Box::into_raw(Box::new(paimon_bucket_vector_search_split { + inner: Box::into_raw(Box::new(split)) as *mut c_void, + })), + error: std::ptr::null_mut(), + }, + Err(error) => paimon_result_bucket_vector_search_split { + split: std::ptr::null_mut(), + error: paimon_error::from_paimon(error), + }, + } +} + +/// Free a decoded split. Null is accepted. +/// # Safety +/// split must be a live handle returned by paimon_bucket_vector_search_split_deserialize, or null. +#[no_mangle] +pub unsafe extern "C" fn paimon_bucket_vector_search_split_free( + split: *mut paimon_bucket_vector_search_split, +) { + if !split.is_null() { + let wrapper = Box::from_raw(split); + if !wrapper.inner.is_null() { + drop(Box::from_raw(wrapper.inner as *mut BucketVectorSearchSplit)); + } + } +} + +const _: unsafe extern "C" fn(*const u8, usize) -> paimon_result_bucket_vector_search_split = + paimon_bucket_vector_search_split_deserialize; +const _: unsafe extern "C" fn(*mut paimon_bucket_vector_search_split) = + paimon_bucket_vector_search_split_free; diff --git a/bindings/c/src/lib.rs b/bindings/c/src/lib.rs index 0a5710ccc..1a3484541 100644 --- a/bindings/c/src/lib.rs +++ b/bindings/c/src/lib.rs @@ -20,6 +20,7 @@ #![allow(non_camel_case_types)] mod blob_reader; +mod bucket_vector_search_split; mod catalog; mod error; mod file_io; @@ -29,6 +30,8 @@ mod table; #[cfg(test)] mod tests; mod types; +mod vector_read; +mod vector_scan; mod vector_search; mod write; diff --git a/bindings/c/src/result.rs b/bindings/c/src/result.rs index 94317572b..3f3b7f98e 100644 --- a/bindings/c/src/result.rs +++ b/bindings/c/src/result.rs @@ -120,6 +120,30 @@ pub struct paimon_result_vector_search_builder { pub error: *mut paimon_error, } +#[repr(C)] +pub struct paimon_result_vector_scan { + pub scan: *mut paimon_vector_scan, + pub error: *mut paimon_error, +} + +#[repr(C)] +pub struct paimon_result_vector_plan { + pub plan: *mut paimon_vector_plan, + pub error: *mut paimon_error, +} + +#[repr(C)] +pub struct paimon_result_vector_read { + pub read: *mut paimon_vector_read, + pub error: *mut paimon_error, +} + +#[repr(C)] +pub struct paimon_result_bucket_vector_search_split { + pub split: *mut paimon_bucket_vector_search_split, + pub error: *mut paimon_error, +} + // === Write/Commit result types === #[repr(C)] diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index 6724e1640..dc66432fd 100644 --- a/bindings/c/src/tests.rs +++ b/bindings/c/src/tests.rs @@ -46,10 +46,13 @@ use paimon::spec::{ use paimon::table::{SnapshotManager, Table}; use crate::blob_reader::*; +use crate::bucket_vector_search_split::*; use crate::error::*; use crate::file_io::*; use crate::table::*; use crate::types::*; +use crate::vector_read::*; +use crate::vector_scan::*; use crate::vector_search::*; use crate::write::*; @@ -2513,7 +2516,7 @@ fn test_two_commits_same_builder() { // // Two storage shapes are exercised end-to-end through the C `execute_read` // terminal, each compared against an independent core Rust -// `VectorSearchBuilder::execute_read()` reference: +// `SearchResultReadBuilder::read()` reference: // // * A primary-key vector table backed by a real vindex IVF-flat ANN segment // built in-process (bucket-local ANN search, residual filter supported). @@ -2987,7 +2990,9 @@ fn rust_execute_read_rows( .with_vector_column(column) .with_query_vector(query) .with_limit(limit); - let mut stream = builder.execute_read().await.unwrap(); + let mut stream = async { builder.execute().await?.new_read_builder().read().await } + .await + .unwrap(); let (mut rows, mut has_score) = (0usize, false); while let Some(b) = stream.try_next().await.unwrap() { rows += b.num_rows(); @@ -3014,7 +3019,9 @@ fn rust_execute_read_pairs( if let Some(f) = filter { builder.with_filter(f); } - let mut stream = builder.execute_read().await.unwrap(); + let mut stream = async { builder.execute().await?.new_read_builder().read().await } + .await + .unwrap(); let mut pairs = Vec::new(); while let Some(b) = stream.try_next().await.unwrap() { pairs.extend(batch_id_score_pairs(&b)); @@ -3247,169 +3254,181 @@ fn stage_split_fixture() -> (tempfile::TempDir, Table, Vec>) { (tmp, table, splits) } -/// The happy path over the ABI, driven by bytes JAVA planned: marshal the split -/// array, read the rows out, free everything. -/// -/// The one C test that reads real Java-produced splits. What the read MEANS -- the -/// row ranges, the projection, the routing rules -- is asserted on the Rust side, -/// where a failure names the semantic that broke; repeating those here would only -/// re-test the same kernel through a thinner lens. This asserts that the marshalling -/// is right and that the rows arrive. +/// Consume a common vector read's Arrow stream in relevance order. +unsafe fn vector_plan_pairs( + result: crate::result::paimon_result_record_batch_reader, +) -> Vec<(i32, f32)> { + assert!( + result.error.is_null(), + "{}", + if result.error.is_null() { + String::new() + } else { + error_message(result.error) + } + ); + let mut pairs = Vec::new(); + loop { + let next = paimon_record_batch_reader_next(result.reader); + assert!(next.error.is_null()); + if next.batch.array.is_null() { + break; + } + pairs.extend(batch_id_score_pairs(&import_batch(&next.batch))); + paimon_arrow_batch_free(next.batch); + } + paimon_record_batch_reader_free(result.reader); + pairs +} + +/// Decode before planning. Each handle owns its state independently, including +/// the Arrow stream after every planning/query handle has been freed. #[cfg(not(target_os = "windows"))] #[test] fn vector_search_bucket_splits_read_the_java_planned_fixture() { - let (_tmp, table, splits) = stage_split_fixture(); + let (_tmp, table, bytes) = stage_split_fixture(); let handle = unsafe { wrap_table(table) }; - unsafe { - let ptrs: Vec<*const u8> = splits.iter().map(|s| s.as_ptr()).collect(); - let lens: Vec = splits.iter().map(Vec::len).collect(); + let mut splits = Vec::new(); + for bytes in &bytes { + let decoded = + paimon_bucket_vector_search_split_deserialize(bytes.as_ptr(), bytes.len()); + assert!(decoded.error.is_null()); + splits.push(decoded.split); + } + drop(bytes); let builder = c_vector_builder(handle, "embedding", &[0.0, 0.0], 3, ptr::null_mut()); - let result = paimon_vector_search_builder_execute_read_for_bucket_splits( - builder, - ptrs.as_ptr(), - lens.as_ptr(), - splits.len(), - ); + let scan = paimon_vector_search_builder_new_scan(builder); + let read = paimon_vector_search_builder_new_read(builder); + assert!(scan.error.is_null() && read.error.is_null()); paimon_vector_search_builder_free(builder); - assert!(result.error.is_null(), "the fixture read must succeed"); - assert!(!result.reader.is_null()); - - let mut pairs = Vec::new(); - loop { - let next = paimon_record_batch_reader_next(result.reader); - assert!(next.error.is_null()); - if next.batch.array.is_null() { - break; - } - let batch = import_batch(&next.batch); - pairs.extend(batch_id_score_pairs(&batch)); - paimon_arrow_batch_free(next.batch); + unwrap_table(handle); + let ptrs: Vec<*const paimon_bucket_vector_search_split> = + splits.iter().map(|&p| p as *const _).collect(); + let plan = paimon_vector_scan_plan_from_bucket_splits(scan.scan, ptrs.as_ptr(), ptrs.len()); + assert!(plan.error.is_null()); + for split in splits { + paimon_bucket_vector_search_split_free(split); } - paimon_record_batch_reader_free(result.reader); - - assert_eq!( - pairs.iter().map(|p| p.0).collect::>(), - vec![0, 1, 2], - "the fixture's top-3 for query [0,0]" - ); + paimon_vector_scan_free(scan.scan); + // The same decoded plan can be read more than once without re-decoding. + let first = vector_plan_pairs(paimon_vector_read_read(read.read, plan.plan)); + let stream = paimon_vector_read_read(read.read, plan.plan); + paimon_vector_read_free(read.read); + paimon_vector_plan_free(plan.plan); + let pairs = vector_plan_pairs(stream); + assert_eq!(pairs, first); + assert_eq!(pairs.iter().map(|p| p.0).collect::>(), vec![0, 1, 2]); for (got, want) in pairs.iter().map(|p| p.1).zip([1.0f32, 0.5, 0.2]) { assert!((got - want).abs() < 1e-4, "score {got} != {want}"); } - - unwrap_table(handle); } } -/// The bucket-split terminal marshals an array of buffers, which the single-split -/// terminal does not: a caller passing a null array, a zero count, or a null -/// entry has made an input error, and it must be reported as one rather than -/// reaching the decoder as corrupt data. #[test] -fn vector_search_bucket_splits_reject_malformed_input() { - let path = "memory:/vsearch_pk_split_args"; - let (query, vectors) = pk_fixture_smoke(); - let table = build_pk_vector_table(path, &vectors); - let handle = unsafe { wrap_table(table) }; - +fn vector_search_de_uses_the_common_scan_plan_read_api() { + let table = build_append_vector_table("memory:/vector_plan_de"); + let expected = rust_execute_read_pairs(&table, "embedding", vec![1.0, 0.0], 2, None); unsafe { - // No splits at all. - let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); - let result = paimon_vector_search_builder_execute_read_for_bucket_splits( - builder, - ptr::null(), - ptr::null(), - 0, - ); - paimon_vector_search_builder_free(builder); - assert!(!result.error.is_null(), "a null split array must error"); - paimon_error_free(result.error); - - // A count that does not match the (absent) arrays. - let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); - let result = paimon_vector_search_builder_execute_read_for_bucket_splits( - builder, - ptr::null(), - ptr::null(), - 1, - ); + let handle = wrap_table(table); + let builder = c_vector_builder(handle, "embedding", &[1.0, 0.0], 2, ptr::null_mut()); + let scan = paimon_vector_search_builder_new_scan(builder); + let read = paimon_vector_search_builder_new_read(builder); + assert!(scan.error.is_null() && read.error.is_null()); paimon_vector_search_builder_free(builder); - assert!(!result.error.is_null(), "a null split array must error"); - paimon_error_free(result.error); - - // A null entry inside an otherwise valid array. - let bytes: Vec = vec![1, 2, 3, 4]; - let ptrs: [*const u8; 2] = [bytes.as_ptr(), ptr::null()]; - let lens: [usize; 2] = [bytes.len(), 0]; - let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); - let result = paimon_vector_search_builder_execute_read_for_bucket_splits( - builder, - ptrs.as_ptr(), - lens.as_ptr(), - 2, - ); - paimon_vector_search_builder_free(builder); - assert!(!result.error.is_null(), "a null split entry must error"); - paimon_error_free(result.error); - unwrap_table(handle); + let plan = paimon_vector_scan_plan(scan.scan); + assert!(plan.error.is_null()); + paimon_vector_scan_free(scan.scan); + let stream = paimon_vector_read_read(read.read, plan.plan); + paimon_vector_plan_free(plan.plan); + paimon_vector_read_free(read.read); + let mut actual = vector_plan_pairs(stream); + actual.sort_by_key(|p| p.0); + assert_eq!(actual, expected); + } +} + +#[test] +fn vector_search_split_decoder_rejects_invalid_buffers_without_a_builder() { + unsafe { + for (bytes, len) in [ + (ptr::null(), 0), + (ptr::null(), 1), + (b"x".as_ptr(), 0), + (b"x".as_ptr(), usize::MAX), + ] { + let result = paimon_bucket_vector_search_split_deserialize(bytes, len); + assert!(result.split.is_null()); + assert_eq!((*result.error).code, PaimonErrorCode::InvalidInput as i32); + paimon_error_free(result.error); + } + for bytes in [vec![0xAB; 64], vec![0; 3]] { + let result = paimon_bucket_vector_search_split_deserialize(bytes.as_ptr(), bytes.len()); + assert!(result.split.is_null() && !result.error.is_null()); + paimon_error_free(result.error); + } + paimon_bucket_vector_search_split_free(ptr::null_mut()); } } -/// A zero-initialized `#[repr(C)]` wrapper passes a null-POINTER check while carrying a -/// null `inner`, which the terminal dereferences. It has to be caught at the boundary, -/// not become a null dereference inside the library. #[test] -fn vector_search_bucket_splits_reject_an_uninitialized_builder() { +fn vector_search_plan_rejects_invalid_split_handles() { + let table = build_append_vector_table("memory:/vector_plan_invalid_handles"); unsafe { - let mut zero_builder = paimon_vector_search_builder { + let handle = wrap_table(table); + let builder = c_vector_builder(handle, "embedding", &[1.0, 0.0], 2, ptr::null_mut()); + let scan = paimon_vector_search_builder_new_scan(builder); + assert!(scan.error.is_null()); + let zero = paimon_bucket_vector_search_split { inner: ptr::null_mut(), }; - let bytes: Vec = vec![1, 2, 3, 4]; - let ptrs: [*const u8; 1] = [bytes.as_ptr()]; - let lens: [usize; 1] = [bytes.len()]; - let result = paimon_vector_search_builder_execute_read_for_bucket_splits( - &mut zero_builder, - ptrs.as_ptr(), - lens.as_ptr(), - 1, - ); - assert!(result.reader.is_null()); - assert!( - !result.error.is_null(), - "a zeroed builder must error, not crash" - ); - let msg = error_message(result.error); - assert!(msg.contains("not initialized"), "got: {msg}"); - paimon_error_free(result.error); + let null_entry = [ptr::null()]; + let zero_entry = [&zero as *const _]; + for (splits, count) in [ + (ptr::null(), 0), + (ptr::null(), 1), + (null_entry.as_ptr(), 1), + (zero_entry.as_ptr(), 1), + (zero_entry.as_ptr(), usize::MAX), + ] { + let result = paimon_vector_scan_plan_from_bucket_splits(scan.scan, splits, count); + assert!(result.plan.is_null() && !result.error.is_null()); + assert_eq!((*result.error).code, PaimonErrorCode::InvalidInput as i32); + paimon_error_free(result.error); + } + paimon_vector_scan_free(scan.scan); + paimon_vector_search_builder_free(builder); + unwrap_table(handle); } } -/// Split bytes come from outside the process, so a buffer that is not a split -/// must surface as an error, not a panic across the ABI boundary. #[test] -fn vector_search_bucket_splits_reject_corrupt_bytes() { - let path = "memory:/vsearch_pk_split_corrupt"; - let (query, vectors) = pk_fixture_smoke(); - let table = build_pk_vector_table(path, &vectors); - let handle = unsafe { wrap_table(table) }; - +fn vector_search_factories_and_reads_reject_uninitialized_handles() { unsafe { - let garbage: Vec = vec![0xAB; 64]; - let ptrs: [*const u8; 1] = [garbage.as_ptr()]; - let lens: [usize; 1] = [garbage.len()]; - let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut()); - let result = paimon_vector_search_builder_execute_read_for_bucket_splits( - builder, - ptrs.as_ptr(), - lens.as_ptr(), - 1, - ); - paimon_vector_search_builder_free(builder); - assert!(!result.error.is_null(), "garbage bytes must error"); - assert!(result.reader.is_null()); + let mut builder = paimon_vector_search_builder { + inner: ptr::null_mut(), + }; + let scan = paimon_vector_search_builder_new_scan(&builder); + let read = paimon_vector_search_builder_new_read(&builder); + let stream = paimon_vector_search_builder_execute_read(&mut builder); + for error in [scan.error, read.error, stream.error] { + assert!(!error.is_null()); + assert!(error_message(error).contains("not initialized")); + paimon_error_free(error); + } + let scan = paimon_vector_scan { + inner: ptr::null_mut(), + }; + let plan = paimon_vector_scan_plan(&scan); + assert!(plan.plan.is_null() && !plan.error.is_null()); + paimon_error_free(plan.error); + let read = paimon_vector_read { + inner: ptr::null_mut(), + }; + let result = paimon_vector_read_read(&read, ptr::null()); + assert!(result.reader.is_null() && !result.error.is_null()); paimon_error_free(result.error); - unwrap_table(handle); } } @@ -3743,10 +3762,12 @@ fn rust_execute_read_column_names( .with_vector_column(column) .with_query_vector(query) .with_limit(limit); + let result = builder.execute().await.unwrap(); + let mut reader = result.new_read_builder(); if let Some(cols) = projection { - builder.with_projection(cols); + reader.with_projection(cols); } - let mut stream = builder.execute_read().await.unwrap(); + let mut stream = reader.read().await.unwrap(); let mut names: Vec = Vec::new(); while let Some(b) = stream.try_next().await.unwrap() { names = b diff --git a/bindings/c/src/types.rs b/bindings/c/src/types.rs index 4e6d2709f..ede866648 100644 --- a/bindings/c/src/types.rs +++ b/bindings/c/src/types.rs @@ -246,6 +246,35 @@ pub struct paimon_vector_search_builder { pub inner: *mut c_void, } +/// Owned vector scan, usable after its builder is freed. +#[repr(C)] +pub struct paimon_vector_scan { + pub inner: *mut c_void, +} + +/// Owned snapshot-scoped vector plan, shared by DE and PK reads. +#[repr(C)] +pub struct paimon_vector_plan { + pub inner: *mut c_void, +} + +/// Owned vector query and output projection. +#[repr(C)] +pub struct paimon_vector_read { + pub inner: *mut c_void, +} + +/// A decoded Java PK bucket split. Free with paimon_bucket_vector_search_split_free. +#[repr(C)] +pub struct paimon_bucket_vector_search_split { + pub inner: *mut c_void, +} + +pub(crate) struct VectorReadState { + pub read: paimon::table::VectorRead, + pub projection: Option>, +} + /// Internal state for a vector-search builder: the table plus the query /// parameters accumulated by the setters before the search is run. pub(crate) struct VectorSearchState { diff --git a/bindings/c/src/vector_read.rs b/bindings/c/src/vector_read.rs new file mode 100644 index 000000000..e6c2339f3 --- /dev/null +++ b/bindings/c/src/vector_read.rs @@ -0,0 +1,71 @@ +// 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. + +//! Executes common vector plans through the C ABI. +use crate::error::check_non_null; +use crate::result::paimon_result_record_batch_reader; +use crate::runtime; +use crate::types::*; +use crate::vector_search::{materialize_search_result, wrap_vector_stream}; +use paimon::table::VectorScanPlan; + +/// Search a DE or PK plan and read projected rows plus __paimon_search_score. +/// The plan is borrowed and can be reused by other queries. Neither the plan nor +/// the reader needs to remain alive after the returned Arrow stream is created. +/// # Safety +/// read and plan must be live handles from vector API constructors, or null (error). +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_read_read( + read: *const paimon_vector_read, + plan: *const paimon_vector_plan, +) -> paimon_result_record_batch_reader { + let validation = check_non_null(read, "read") + .and_then(|_| check_non_null((*read).inner, "read is not initialized")) + .and_then(|_| check_non_null(plan, "plan")) + .and_then(|_| check_non_null((*plan).inner, "plan is not initialized")); + if let Err(error) = validation { + return paimon_result_record_batch_reader { + reader: std::ptr::null_mut(), + error, + }; + } + let state = &*((*read).inner as *const VectorReadState); + let plan = (&*((*plan).inner as *const VectorScanPlan)).clone(); + wrap_vector_stream(runtime().block_on(async { + let result = state.read.read(plan).await?; + materialize_search_result(result, state.projection.as_deref()).await + })) +} + +/// Free a vector reader. Null is accepted. +/// # Safety +/// read must be a live owned vector reader handle, or null. +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_read_free(read: *mut paimon_vector_read) { + if !read.is_null() { + let wrapper = Box::from_raw(read); + if !wrapper.inner.is_null() { + drop(Box::from_raw(wrapper.inner as *mut VectorReadState)); + } + } +} + +const _: unsafe extern "C" fn( + *const paimon_vector_read, + *const paimon_vector_plan, +) -> paimon_result_record_batch_reader = paimon_vector_read_read; +const _: unsafe extern "C" fn(*mut paimon_vector_read) = paimon_vector_read_free; diff --git a/bindings/c/src/vector_scan.rs b/bindings/c/src/vector_scan.rs new file mode 100644 index 000000000..fbc45eed0 --- /dev/null +++ b/bindings/c/src/vector_scan.rs @@ -0,0 +1,149 @@ +// 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. + +//! Common vector plans for DE and PK execution. +use crate::error::{check_non_null, paimon_error, PaimonErrorCode}; +use crate::result::paimon_result_vector_plan; +use crate::runtime; +use crate::types::*; +use paimon::table::{BucketVectorSearchSplit, VectorScan, VectorScanPlan}; +use std::ffi::c_void; + +fn wrap_plan(result: paimon::Result) -> paimon_result_vector_plan { + match result { + Ok(plan) => paimon_result_vector_plan { + plan: Box::into_raw(Box::new(paimon_vector_plan { + inner: Box::into_raw(Box::new(plan)) as *mut c_void, + })), + error: std::ptr::null_mut(), + }, + Err(error) => paimon_result_vector_plan { + plan: std::ptr::null_mut(), + error: paimon_error::from_paimon(error), + }, + } +} + +unsafe fn scan_ref<'a>( + scan: *const paimon_vector_scan, +) -> Result<&'a VectorScan, *mut paimon_error> { + check_non_null(scan, "scan")?; + check_non_null((*scan).inner, "scan is not initialized")?; + Ok(&*((*scan).inner as *const VectorScan)) +} + +/// Resolve a snapshot and plan vector search for either a DE or PK table. +/// The returned plan is independent of the scan; free it with paimon_vector_plan_free. +/// # Safety +/// scan must be a live handle from paimon_vector_search_builder_new_scan, or null (error). +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_scan_plan( + scan: *const paimon_vector_scan, +) -> paimon_result_vector_plan { + match scan_ref(scan) { + Ok(scan) => wrap_plan(runtime().block_on(scan.plan())), + Err(error) => paimon_result_vector_plan { + plan: std::ptr::null_mut(), + error, + }, + } +} + +/// Construct a common read plan from already-decoded Java PK bucket splits. +/// No snapshot or index manifest is read. Handles are borrowed and copied; +/// callers may free them as soon as this returns. Failure leaves inputs intact. +/// Empty input, DE scans and mixed snapshots are rejected. +/// # Safety +/// scan must be a live scan handle. splits must point to count live split pointers; +/// null pointers, uninitialized handles and zero/oversized counts return an error. +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_scan_plan_from_bucket_splits( + scan: *const paimon_vector_scan, + splits: *const *const paimon_bucket_vector_search_split, + count: usize, +) -> paimon_result_vector_plan { + let scan = match scan_ref(scan) { + Ok(scan) => scan, + Err(error) => { + return paimon_result_vector_plan { + plan: std::ptr::null_mut(), + error, + } + } + }; + if splits.is_null() + || count == 0 + || count + > isize::MAX as usize / std::mem::size_of::<*const paimon_bucket_vector_search_split>() + { + return paimon_result_vector_plan { + plan: std::ptr::null_mut(), + error: paimon_error::new( + PaimonErrorCode::InvalidInput, + "null, empty or oversized split array".to_string(), + ), + }; + } + let mut decoded = Vec::with_capacity(count); + for &split in std::slice::from_raw_parts(splits, count) { + if split.is_null() || (*split).inner.is_null() { + return paimon_result_vector_plan { + plan: std::ptr::null_mut(), + error: paimon_error::new( + PaimonErrorCode::InvalidInput, + "split is null or not initialized".to_string(), + ), + }; + } + decoded.push((&*((*split).inner as *const BucketVectorSearchSplit)).clone()); + } + wrap_plan(scan.plan_from_bucket_splits(decoded)) +} + +/// Free a vector scan. Null is accepted. +/// # Safety +/// scan must be a live owned scan handle, or null. +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_scan_free(scan: *mut paimon_vector_scan) { + if !scan.is_null() { + let wrapper = Box::from_raw(scan); + if !wrapper.inner.is_null() { + drop(Box::from_raw(wrapper.inner as *mut VectorScan)); + } + } +} + +/// Free a vector plan. Null is accepted. +/// # Safety +/// plan must be a live owned vector plan handle, or null. +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_plan_free(plan: *mut paimon_vector_plan) { + if !plan.is_null() { + let wrapper = Box::from_raw(plan); + if !wrapper.inner.is_null() { + drop(Box::from_raw(wrapper.inner as *mut VectorScanPlan)); + } + } +} + +const _: unsafe extern "C" fn(*const paimon_vector_scan) -> paimon_result_vector_plan = + paimon_vector_scan_plan; +const _: unsafe extern "C" fn( + *const paimon_vector_scan, + *const *const paimon_bucket_vector_search_split, + usize, +) -> paimon_result_vector_plan = paimon_vector_scan_plan_from_bucket_splits; diff --git a/bindings/c/src/vector_search.rs b/bindings/c/src/vector_search.rs index 2a04cba66..31c362f05 100644 --- a/bindings/c/src/vector_search.rs +++ b/bindings/c/src/vector_search.rs @@ -31,10 +31,14 @@ use std::collections::HashMap; use std::ffi::{c_char, c_void}; use paimon::spec::Predicate; -use paimon::table::Table; +use paimon::table::{ArrowRecordBatchStream, Table, VectorSearchBuilder}; +use paimon::vector_search::SearchResult; use crate::error::{check_non_null, paimon_error, validate_cstr, PaimonErrorCode}; -use crate::result::{paimon_result_record_batch_reader, paimon_result_vector_search_builder}; +use crate::result::{ + paimon_result_record_batch_reader, paimon_result_vector_read, paimon_result_vector_scan, + paimon_result_vector_search_builder, +}; use crate::runtime; use crate::types::*; @@ -300,46 +304,43 @@ pub unsafe extern "C" fn paimon_vector_search_builder_free(b: *mut paimon_vector pub unsafe extern "C" fn paimon_vector_search_builder_execute_read( b: *mut paimon_vector_search_builder, ) -> paimon_result_record_batch_reader { - if let Err(e) = check_non_null(b, "b") { - return paimon_result_record_batch_reader { - reader: std::ptr::null_mut(), - error: e, - }; - } - let state = &*((*b).inner as *const VectorSearchState); + let state = match vector_search_state(b) { + Ok(state) => state, + Err(error) => { + return paimon_result_record_batch_reader { + reader: std::ptr::null_mut(), + error, + } + } + }; + let builder = configured_builder(state); + wrap_vector_stream(runtime().block_on(async { + materialize_search_result(builder.execute().await?, state.projection.as_deref()).await + })) +} - let mut builder = state.table.new_vector_search_builder(); - if let Some(col) = &state.vector_column { - builder.with_vector_column(col); - } - if let Some(v) = &state.query_vector { - builder.with_query_vector(v.clone()); - } - if let Some(limit) = state.limit { - builder.with_limit(limit); - } - if !state.options.is_empty() { - builder.with_options(state.options.clone()); - } - if let Some(f) = &state.filter { - builder.with_filter(f.clone()); - } - if let Some(cols) = &state.projection { +pub(crate) async fn materialize_search_result( + result: SearchResult, + projection: Option<&[String]>, +) -> paimon::Result { + let mut reader = result.new_read_builder(); + if let Some(cols) = projection { let col_refs: Vec<&str> = cols.iter().map(String::as_str).collect(); - builder.with_projection(&col_refs); + reader.with_projection(&col_refs); } + reader.read().await +} - match runtime().block_on(builder.execute_read()) { - Ok(stream) => { - let reader = Box::new(stream); - let wrapper = Box::new(paimon_record_batch_reader { - inner: Box::into_raw(reader) as *mut c_void, - }); - paimon_result_record_batch_reader { - reader: Box::into_raw(wrapper), - error: std::ptr::null_mut(), - } - } +pub(crate) fn wrap_vector_stream( + result: paimon::Result, +) -> paimon_result_record_batch_reader { + match result { + Ok(stream) => paimon_result_record_batch_reader { + reader: Box::into_raw(Box::new(paimon_record_batch_reader { + inner: Box::into_raw(Box::new(stream)) as *mut c_void, + })), + error: std::ptr::null_mut(), + }, Err(e) => paimon_result_record_batch_reader { reader: std::ptr::null_mut(), error: paimon_error::from_paimon(e), @@ -347,129 +348,86 @@ pub unsafe extern "C" fn paimon_vector_search_builder_execute_read( } } -/// Run the vector search over bucket splits a Java planner produced, and stream -/// the materialized rows. -/// -/// This is the entry point for an engine that plans in Paimon Java and executes -/// here: the planner emits one `BucketVectorSearchSplit` per bucket and ships its -/// bytes to a worker, which calls this. `splits` points at `count` buffers and -/// `split_lens` at their lengths; both arrays must hold `count` entries. The -/// buffers are only read for the duration of the call. -/// -/// The splits are the plan -- their payload files, per-file row ranges and pinned -/// snapshot are used as given, and the table's index manifest is not read. Search, -/// optional refine, Top-K and materialization are the same as -/// `paimon_vector_search_builder_execute_read`, so the output is the projected -/// user columns plus `__paimon_search_score`, best-first. The Top-K is local to -/// the splits passed in; a caller distributing one call per bucket merges the -/// per-bucket results itself. -/// -/// Only a primary-key vector column can be read this way; a data-evolution table -/// returns an error rather than an answer from a different plan. Consume via -/// `paimon_record_batch_reader_next` and free with -/// `paimon_record_batch_reader_free`. -/// -/// # Safety -/// `b` must be a valid pointer from `paimon_table_new_vector_search_builder`, or -/// null (returns an error result). `splits` and `split_lens` must each point at -/// `count` valid entries, and each `splits[i]` at `split_lens[i]` readable bytes. -#[no_mangle] -pub unsafe extern "C" fn paimon_vector_search_builder_execute_read_for_bucket_splits( - b: *mut paimon_vector_search_builder, - splits: *const *const u8, - split_lens: *const usize, - count: usize, -) -> paimon_result_record_batch_reader { - if let Err(e) = check_non_null(b, "b") { - return paimon_result_record_batch_reader { - reader: std::ptr::null_mut(), - error: e, - }; - } - // A `#[repr(C)]` wrapper can arrive zero-initialized, which passes the null-POINTER - // check above while carrying a null `inner`. The state is dereferenced below, so - // that has to be caught here rather than as a null dereference inside the library. - if (*b).inner.is_null() { - return paimon_result_record_batch_reader { - reader: std::ptr::null_mut(), - error: paimon_error::new( - PaimonErrorCode::InvalidInput, - concat!( - "paimon_vector_search_builder_execute_read_for_bucket_splits: ", - "builder is not initialized" - ) - .to_string(), - ), - }; - } - if splits.is_null() || split_lens.is_null() || count == 0 { - return paimon_result_record_batch_reader { - reader: std::ptr::null_mut(), - error: paimon_error::new( - PaimonErrorCode::InvalidInput, - "paimon_vector_search_builder_execute_read_for_bucket_splits: null or empty splits" - .to_string(), - ), - }; - } - - let ptrs = std::slice::from_raw_parts(splits, count); - let lens = std::slice::from_raw_parts(split_lens, count); - let mut buffers: Vec<&[u8]> = Vec::with_capacity(count); - for (i, (&ptr, &len)) in ptrs.iter().zip(lens).enumerate() { - // A null or empty buffer cannot be a split, and reaching the decoder with - // one would report it as corrupt data rather than as the caller's error. - if ptr.is_null() || len == 0 { - return paimon_result_record_batch_reader { - reader: std::ptr::null_mut(), - error: paimon_error::new( - PaimonErrorCode::InvalidInput, - format!( - "paimon_vector_search_builder_execute_read_for_bucket_splits: \ - split {i} is null or empty" - ), - ), - }; - } - buffers.push(std::slice::from_raw_parts(ptr, len)); - } +unsafe fn vector_search_state<'a>( + builder: *const paimon_vector_search_builder, +) -> Result<&'a VectorSearchState, *mut paimon_error> { + check_non_null(builder, "builder")?; + check_non_null((*builder).inner, "builder is not initialized")?; + Ok(&*((*builder).inner as *const VectorSearchState)) +} - let state = &*((*b).inner as *const VectorSearchState); +fn configured_builder(state: &VectorSearchState) -> VectorSearchBuilder<'_> { let mut builder = state.table.new_vector_search_builder(); - if let Some(col) = &state.vector_column { - builder.with_vector_column(col); + if let Some(column) = &state.vector_column { + builder.with_vector_column(column); } - if let Some(v) = &state.query_vector { - builder.with_query_vector(v.clone()); + if let Some(vector) = &state.query_vector { + builder.with_query_vector(vector.clone()); } if let Some(limit) = state.limit { builder.with_limit(limit); } - if !state.options.is_empty() { - builder.with_options(state.options.clone()); - } - if let Some(f) = &state.filter { - builder.with_filter(f.clone()); + builder.with_options(state.options.clone()); + if let Some(filter) = &state.filter { + builder.with_filter(filter.clone()); } - if let Some(cols) = &state.projection { - let col_refs: Vec<&str> = cols.iter().map(String::as_str).collect(); - builder.with_projection(&col_refs); + builder +} + +/// Create an owned DE or PK vector scan. The vector column must be configured; +/// a query vector and limit are not required. Free with paimon_vector_scan_free. +/// # Safety +/// builder must be a live vector-search builder handle, or null (error). +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_search_builder_new_scan( + builder: *const paimon_vector_search_builder, +) -> paimon_result_vector_scan { + let result = vector_search_state(builder).and_then(|state| { + configured_builder(state) + .new_scan() + .map_err(paimon_error::from_paimon) + }); + match result { + Ok(scan) => paimon_result_vector_scan { + scan: Box::into_raw(Box::new(paimon_vector_scan { + inner: Box::into_raw(Box::new(scan)) as *mut c_void, + })), + error: std::ptr::null_mut(), + }, + Err(error) => paimon_result_vector_scan { + scan: std::ptr::null_mut(), + error, + }, } +} - match runtime().block_on(builder.execute_read_for_bucket_splits(&buffers)) { - Ok(stream) => { - let reader = Box::new(stream); - let wrapper = Box::new(paimon_record_batch_reader { - inner: Box::into_raw(reader) as *mut c_void, - }); - paimon_result_record_batch_reader { - reader: Box::into_raw(wrapper), - error: std::ptr::null_mut(), - } - } - Err(e) => paimon_result_record_batch_reader { - reader: std::ptr::null_mut(), - error: paimon_error::from_paimon(e), +/// Create an owned DE or PK reader from configured query parameters and projection. +/// The builder may then be freed. Free the reader with paimon_vector_read_free. +/// # Safety +/// builder must be a live vector-search builder handle, or null (error). +#[no_mangle] +pub unsafe extern "C" fn paimon_vector_search_builder_new_read( + builder: *const paimon_vector_search_builder, +) -> paimon_result_vector_read { + let result = vector_search_state(builder).and_then(|state| { + configured_builder(state) + .new_read() + .map(|read| VectorReadState { + read, + projection: state.projection.clone(), + }) + .map_err(paimon_error::from_paimon) + }); + match result { + Ok(read) => paimon_result_vector_read { + read: Box::into_raw(Box::new(paimon_vector_read { + inner: Box::into_raw(Box::new(read)) as *mut c_void, + })), + error: std::ptr::null_mut(), + }, + Err(error) => paimon_result_vector_read { + read: std::ptr::null_mut(), + error, }, } } @@ -517,10 +475,8 @@ const _: unsafe extern "C" fn(*mut paimon_vector_search_builder) = const _: unsafe extern "C" fn( *mut paimon_vector_search_builder, ) -> paimon_result_record_batch_reader = paimon_vector_search_builder_execute_read; -const _: unsafe extern "C" fn( - *mut paimon_vector_search_builder, - *const *const u8, - *const usize, - usize, -) -> paimon_result_record_batch_reader = - paimon_vector_search_builder_execute_read_for_bucket_splits; + +const _: unsafe extern "C" fn(*const paimon_vector_search_builder) -> paimon_result_vector_scan = + paimon_vector_search_builder_new_scan; +const _: unsafe extern "C" fn(*const paimon_vector_search_builder) -> paimon_result_vector_read = + paimon_vector_search_builder_new_read; diff --git a/crates/integrations/datafusion/src/lateral_vector_search.rs b/crates/integrations/datafusion/src/lateral_vector_search.rs index 2500d8b52..75a6250b0 100644 --- a/crates/integrations/datafusion/src/lateral_vector_search.rs +++ b/crates/integrations/datafusion/src/lateral_vector_search.rs @@ -59,7 +59,7 @@ use datafusion::prelude::SessionConfig; use futures::{Stream, StreamExt, TryStreamExt}; use paimon::spec::{Predicate, ROW_ID_FIELD_NAME}; use paimon::table::{PreparedVectorSearchFilter, RowRange, Table}; -use paimon::vector_search::SearchResult; +use paimon::vector_search::ScoredRowIds; use tokio::sync::OnceCell; use crate::error::to_datafusion_error; @@ -781,7 +781,13 @@ impl LateralVectorSearchExec { if let Some(prepared_filter) = prepared_filter { builder.with_prepared_filter(prepared_filter.clone()); } - let results = builder.execute().await.map_err(to_datafusion_error)?; + let results = builder + .execute() + .await + .map_err(to_datafusion_error)? + .into_iter() + .map(|result| result.into_row_ids().map_err(to_datafusion_error)) + .collect::>>()?; let (target_batch, target_row_id_to_index) = read_target_rows(target_table, &self.target_schema, &results).await?; @@ -989,7 +995,7 @@ fn collect_query_vectors(array: &ArrayRef) -> DFResult<(Vec>, Vec DFResult<(RecordBatch, HashMap)> { let mut row_ids = results .iter() @@ -1092,7 +1098,7 @@ async fn read_target_rows( fn row_ranges_from_row_ids(row_ids: &[u64]) -> DFResult> { let scores = vec![0.0; row_ids.len()]; - SearchResult::new(row_ids.to_vec(), scores) + ScoredRowIds::new(row_ids.to_vec(), scores) .to_row_ranges() .map_err(to_datafusion_error) } diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index d37d16f2c..715ce37ee 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -403,7 +403,10 @@ impl VectorSearchExec { builder.with_prepared_filter(prepared.clone()); } let mut results = builder.execute().await.map_err(to_datafusion_error)?; - Ok::<_, DataFusionError>(results.remove(0)) + results + .remove(0) + .into_row_ids() + .map_err(to_datafusion_error) }) .await?; diff --git a/crates/paimon/src/table/batch_vector_search_builder.rs b/crates/paimon/src/table/batch_vector_search_builder.rs new file mode 100644 index 000000000..4e1f9d9b5 --- /dev/null +++ b/crates/paimon/src/table/batch_vector_search_builder.rs @@ -0,0 +1,170 @@ +// 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. + +//! Configures batch vector queries and dispatches to DE or primary-key readers. + +use crate::spec::{CoreOptions, Predicate}; +use crate::table::de_vector_scan::PreparedVectorSearchFilter; +use crate::table::vector_read::BatchVectorRead; +use crate::table::vector_scan::{PlanContext, VectorScan}; +use crate::table::Table; +use crate::vector_search::SearchResult; +use roaring::RoaringTreemap; +use std::collections::HashMap; +use std::sync::Arc; + +pub struct BatchVectorSearchBuilder<'a> { + table: &'a Table, + vector_column: Option, + query_vectors: Option>>, + limit: Option, + options: HashMap, + filter: Option, + include_row_ids: Option>, + prepared_filter: Option, +} + +impl<'a> BatchVectorSearchBuilder<'a> { + pub(crate) fn new(table: &'a Table) -> Self { + Self { + table, + vector_column: None, + query_vectors: None, + limit: None, + options: HashMap::new(), + filter: None, + include_row_ids: None, + prepared_filter: None, + } + } + + pub fn with_vector_column(&mut self, name: &str) -> &mut Self { + self.vector_column = Some(name.to_string()); + self + } + + pub fn with_query_vectors(&mut self, vectors: Vec>) -> &mut Self { + self.query_vectors = Some(vectors); + self + } + + pub fn with_limit(&mut self, limit: usize) -> &mut Self { + self.limit = Some(limit); + self + } + + pub fn with_options(&mut self, options: HashMap) -> &mut Self { + self.options = options; + self + } + + /// Attach one scalar predicate shared by every query in the batch and applied + /// before vector Top-K. See [`crate::table::VectorSearchBuilder::with_filter`] for the + /// primary-key and data-evolution execution semantics. + pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { + self.filter = Some(filter); + self.include_row_ids = None; + self.prepared_filter = None; + self + } + + /// Attach a prepared scalar pre-filter together with the exact table + /// snapshot against which its row-ID allow-list was evaluated. + pub fn with_prepared_filter( + &mut self, + prepared_filter: PreparedVectorSearchFilter, + ) -> &mut Self { + self.prepared_filter = Some(prepared_filter); + self.filter = None; + self.include_row_ids = None; + self + } + + /// Attach a caller-managed row-ID allow-list. + /// + /// This low-level API does not bind the allow-list to a table snapshot. + /// Prefer [`Self::with_prepared_filter`] for scalar pre-filters. + pub fn with_include_row_ids(&mut self, include_row_ids: RoaringTreemap) -> &mut Self { + self.include_row_ids = Some(Arc::new(include_row_ids)); + self.filter = None; + self.prepared_filter = None; + self + } + + /// Create the same query-independent scan used by a single-vector builder. + pub fn new_scan(&self) -> crate::Result { + let column = self.column()?; + VectorScan::new( + self.table, + column, + self.filter.as_ref(), + self.include_row_ids.as_ref(), + self.prepared_filter.as_ref(), + ) + } + + /// Create an owned batch reader; result i belongs to input query i. + pub fn new_read(&self) -> crate::Result { + let column = self.column()?; + PlanContext::new( + self.table, + column, + self.filter.as_ref(), + self.include_row_ids.as_ref(), + self.prepared_filter.as_ref(), + )?; + let queries = self + .query_vectors + .as_deref() + .filter(|queries| !queries.is_empty()) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Query vectors must be set via with_query_vectors()".to_string(), + })?; + let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { + message: "Limit must be set via with_limit()".to_string(), + })?; + let query_refs: Vec<&[f32]> = queries.iter().map(Vec::as_slice).collect(); + BatchVectorRead::new( + self.table, + column, + &query_refs, + limit, + &self.options, + self.filter.as_ref(), + self.include_row_ids.as_ref(), + self.prepared_filter.as_ref(), + ) + } + + /// Search every query against one plan, including empty per-query results. + pub async fn execute(&self) -> crate::Result> { + let read = self.new_read()?; + read.read(self.new_scan()?.plan().await?).await + } + + fn column(&self) -> crate::Result<&str> { + CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + self.vector_column + .as_deref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Vector column must be set via with_vector_column()".to_string(), + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/batch_vector_search_builder/tests.rs b/crates/paimon/src/table/batch_vector_search_builder/tests.rs new file mode 100644 index 000000000..125e776f5 --- /dev/null +++ b/crates/paimon/src/table/batch_vector_search_builder/tests.rs @@ -0,0 +1,157 @@ +// 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::table::pk_vector_position_read::PKEY_VECTOR_POSITION_COLUMN; +use crate::table::vector_search_common::resolve_materialize_read_type; +use crate::table::vector_search_test_utils::{ + de_vector_table, pk_vector_table_with_extra_column, vector_test_table, +}; +use std::collections::HashMap; + +#[tokio::test] +async fn test_batch_vector_search_requires_vectors() { + let table = vector_test_table(); + let err = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(Vec::new()) + .with_limit(1) + .execute() + .await + .unwrap_err(); + + assert!( + err.to_string() + .contains("Query vectors must be set via with_query_vectors()"), + "unexpected error: {err}" + ); +} + +#[tokio::test] +async fn test_batch_vector_search_rejects_zero_limit() { + let table = vector_test_table(); + let err = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0]]) + .with_limit(0) + .execute() + .await + .unwrap_err(); + + assert!( + err.to_string().contains("Limit must be between 1"), + "unexpected error: {err}" + ); +} + +#[tokio::test] +async fn empty_prepared_filter_does_not_mask_an_invalid_query_in_the_batch() { + use crate::table::vector_search_test_utils::id_gt_filter; + + let table = vector_test_table(); + let prepared = table + .prepare_vector_search_filter(id_gt_filter(&table, 0)) + .await + .unwrap(); + assert!(prepared.include_row_ids().is_empty()); + let error = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![]]) + .with_limit(2) + .with_prepared_filter(prepared) + .execute() + .await + .unwrap_err(); + assert!( + error.to_string().contains("Search vector cannot be empty"), + "{error}" + ); +} + +#[tokio::test] +async fn single_and_batch_de_readers_preserve_query_option_precedence() { + let table = de_vector_table().await.copy_with_options(HashMap::from([( + "fields.embedding.ivf.refine-factor".to_string(), + "invalid".to_string(), + )])); + let options = HashMap::from([("refine_factor".to_string(), "1".to_string())]); + let mut single = table.new_vector_search_builder(); + single + .with_vector_column("embedding") + .with_query_vector(vec![0.0, 1.0]) + .with_limit(1); + assert!( + single.execute().await.is_err(), + "the invalid table option must require an override" + ); + let result = single + .with_options(options.clone()) + .execute() + .await + .unwrap(); + assert_eq!(result.row_ids().unwrap().row_ids, vec![1]); + + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_options(options) + .execute() + .await + .unwrap(); + assert_eq!(results.len(), 2); + for (result, expected) in results.iter().zip([vec![0, 2], vec![1, 2]]) { + assert_eq!(result.row_ids().unwrap().row_ids, expected); + } +} + +#[tokio::test] +async fn test_batch_execute_fails_closed_when_query_auth_enabled() { + // The batch scored entry returns data-derived row ids/scores outside + // `TableScan`/`TableRead`, so it must fail closed under + // `query-auth.enabled` exactly like the single-query builder. Its config + // is otherwise valid, so without the guard the empty-snapshot fast path + // would return empty results and silently bypass authorization. + let table = crate::table::query_auth_table(); + let err = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 2.0]]) + .with_limit(5) + .execute() + .await + .unwrap_err(); + assert!( + matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), + "batch vector search must fail closed for a query-auth table, got: {err:?}" + ); +} + +#[test] +fn batch_resolve_materialize_read_type_default_rejects_reserved_user_column() { + // Same guard on the batch resolver. + let table = pk_vector_table_with_extra_column(PKEY_VECTOR_POSITION_COLUMN); + let err = resolve_materialize_read_type(&table, None).unwrap_err(); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "batch default projection must reject reserved user column, got: {err:?}" + ); +} diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs new file mode 100644 index 000000000..328962ae0 --- /dev/null +++ b/crates/paimon/src/table/de_vector_read.rs @@ -0,0 +1,2091 @@ +// 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. + +//! Reads global vector indexes, scores raw vectors, and materializes global-row-ID results. + +use crate::io::{FileIO, FileRead}; +use crate::lumina::reader::LuminaVectorGlobalIndexReader; +use crate::lumina::{LuminaIndexMeta, LuminaVectorMetric}; +use crate::spec::{ + row_id_data_field, CoreOptions, DataField, DataType, FileKind, GlobalIndexSearchMode, + IndexFileMeta, IndexManifestEntry, ROW_ID_FIELD_NAME, +}; +use crate::table::de_vector_scan::DeVectorScanPlan; +use crate::table::global_index_scanner::{ + deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, + unindexed_ranges_for_global_index_entries, RowRangeIndex, +}; +use crate::table::index_file_path::IndexFileLocation; +use crate::table::pk_vector_position_read::SEARCH_SCORE_COLUMN; +use crate::table::row_id_predicate::intersect_sorted_ranges; +use crate::table::vector_read::Read; +use crate::table::vector_search_common::{ + configured_refine_factor, indexed_search_limit, log_vindex_range_io_stats, normalize_metric, + resolve_materialize_read_type, vindex_concurrency_limits, VectorIndexBackend, +}; +use crate::table::{ + find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, Table, +}; +use crate::vector_search::{GlobalIndexIOMeta, ScoredRowIds, SearchResult, VectorSearch}; +use crate::vindex::executor::{ + acquire_process_global_search_permit, drain_indexed_jobs, + ensure_global_index_executor_capacity, execute_global_index_with_guard, +}; +use crate::vindex::range_reader::{RangeReadLimiter, VindexFileReader}; +use crate::vindex::reader::VindexVectorGlobalIndexReader; +use crate::vindex::{is_vindex_index_type, vector_search_timing_enabled}; +use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; +use arrow_select::interleave::interleave_record_batch; +use futures::{stream, TryStreamExt}; +use paimon_vindex_core::blas::sgemm_a_bt; +use paimon_vindex_core::diskann_io::DISKANN_HEADER_SIZE; +use paimon_vindex_core::distance::MetricType; +use paimon_vindex_core::index::VectorIndexReader as VIndexReader; +use paimon_vindex_core::io::SeekRead; +use roaring::RoaringTreemap; +use std::borrow::Cow; +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::io::Cursor; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// Owns the queries and executes the snapshot-scoped plan from `DeVectorScan`. +pub(super) struct DeVectorRead { + vector_searches: Vec, +} + +impl DeVectorRead { + /// Validate query parameters before planning, including for empty tables and + /// filters. Snapshot-dependent row-ID filters are supplied by the plan. + pub(super) fn new( + vector_column: &str, + queries: &[&[f32]], + limit: usize, + options: &HashMap, + ) -> crate::Result { + if vector_column.is_empty() { + return Err(crate::Error::ConfigInvalid { + message: "Vector column must be set via with_vector_column()".to_string(), + }); + } + if queries.is_empty() { + return Err(crate::Error::ConfigInvalid { + message: "Query vectors must be set via with_query_vectors()".to_string(), + }); + } + let vector_searches = queries + .iter() + .map(|query| { + VectorSearch::new(query.to_vec(), limit, vector_column.to_string()) + .map(|search| search.with_options(options.clone())) + }) + .collect::>()?; + Ok(Self { vector_searches }) + } + + /// Validate a result read, including projections over an empty result. + pub(super) fn read_type( + table: &Table, + vector_column: &str, + projection: Option<&[String]>, + ) -> crate::Result> { + // Validate the target column exists and is a vector-bearing type before any + // work. The data-evolution search returns an empty result for an unknown + // field (its scored-path behavior), which would make a typo'd or scalar + // column look like a normal empty read here — violating the result reader's + // fail-loud contract (a C/Doris caller would see EOF, not an input error). + // Reject it up front instead. + let field = table + .schema() + .fields() + .iter() + .find(|f| f.name() == vector_column) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("vector search column '{vector_column}' does not exist"), + source: None, + })?; + // Require a FLOAT-element vector column: `ARRAY` or `VECTOR`, + // matching the element type the vector index/search operates on. An + // `ARRAY` (or any non-float element) is not a searchable vector column. + let is_float_vector = match field.data_type() { + DataType::Vector(t) => matches!(t.element_type(), DataType::Float(_)), + DataType::Array(t) => matches!(t.element_type(), DataType::Float(_)), + _ => false, + }; + if !is_float_vector { + return Err(crate::Error::DataInvalid { + message: format!( + "vector search column '{vector_column}' must be a FLOAT vector column \ + (ARRAY or VECTOR), got {:?}", + field.data_type() + ), + source: None, + }); + } + + resolve_materialize_read_type(table, projection) + } +} + +pub(super) async fn materialize_row_ids( + pinned_table: &Table, + sr: &ScoredRowIds, + mut read_type: Vec, +) -> crate::Result { + if sr.is_empty() { + return Ok(Box::pin(stream::empty())); + } + // rank = ordinal in the best-first scored result; score = the aligned score. + // Build ranges first (validates ids fit in i64::MAX) before constructing the map. + let ranges = sr.to_row_ranges()?; + let mut rank_score_of: HashMap = HashMap::new(); + for (rank, (&id, &score)) in sr.row_ids.iter().zip(sr.scores.iter()).enumerate() { + rank_score_of.insert(id as i64, (rank, score)); + } + + // Add _ROW_ID as the join key for score alignment; it is stripped before output. + if !read_type.iter().any(|f| f.name() == ROW_ID_FIELD_NAME) { + read_type.push(row_id_data_field()); + } + + let mut read_builder = pinned_table.new_read_builder(); + read_builder + .with_read_type(read_type) + .with_row_ranges(ranges); + let scan = read_builder.new_scan(); + let plan = scan.plan().await?; + let table_read = read_builder.new_read()?; + let mut stream = table_read.to_arrow(plan.splits())?; + + let mut batches: Vec = Vec::new(); + while let Some(batch) = stream.try_next().await? { + batches.push(batch); + } + let output = attach_scores_by_row_id(&batches, &rank_score_of, sr.len())?; + Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) +} + +impl Read for DeVectorRead { + type Plan = DeVectorScanPlan; + + async fn read(&self, plan: DeVectorScanPlan) -> crate::Result> { + let DeVectorScanPlan { + table, + index_entries, + include_row_ids, + next_row_id, + timing, + skip_search, + } = plan; + let table = Arc::new(table); + let make_results = |results: Vec| { + results + .into_iter() + .zip(&self.vector_searches) + .map(|(hits, query)| { + SearchResult::from_row_ids(table.clone(), query.field_name.clone(), hits) + }) + .collect() + }; + let pinned_table = &table; + let timing_enabled = timing.is_some(); + let (total_start, setup, snapshot_elapsed, manifest) = match timing { + Some(timing) => ( + Some(timing.total_start), + timing.setup, + timing.snapshot, + timing.manifest, + ), + None => (None, Duration::ZERO, Duration::ZERO, Duration::ZERO), + }; + let evaluate_start = (timing_enabled && !skip_search).then(Instant::now); + let results = if skip_search { + vec![ScoredRowIds::empty(); self.vector_searches.len()] + } else { + let mut vector_searches = Cow::Borrowed(self.vector_searches.as_slice()); + if let Some(include_row_ids) = include_row_ids { + for search in vector_searches.to_mut() { + search.set_shared_include_row_ids(Arc::clone(&include_row_ids)); + } + } + evaluate_batch_vector_search( + VectorSearchEvaluation { + table: Some(&pinned_table), + file_io: pinned_table.file_io(), + table_path: pinned_table.location(), + table_options: pinned_table.schema().options(), + schema_fields: pinned_table.schema().fields(), + next_row_id, + }, + &index_entries, + &vector_searches, + ) + .await? + }; + if let Some(total_start) = total_start { + let total = total_start.elapsed(); + let evaluate = evaluate_start.map_or(Duration::ZERO, |start| start.elapsed()); + let children = setup + .saturating_add(snapshot_elapsed) + .saturating_add(manifest) + .saturating_add(evaluate); + let result_count = results + .iter() + .map(|result| result.row_ids.len()) + .sum::(); + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_search_api nq={} index_entries={} result_count={} total_ms={:.3} setup_ms={:.3} snapshot_ms={:.3} manifest_ms={:.3} evaluate_ms={:.3} unattributed_ms={:.3}", + self.vector_searches.len(), + index_entries.len(), + result_count, + total.as_secs_f64() * 1000.0, + setup.as_secs_f64() * 1000.0, + snapshot_elapsed.as_secs_f64() * 1000.0, + manifest.as_secs_f64() * 1000.0, + evaluate.as_secs_f64() * 1000.0, + total.saturating_sub(children).as_secs_f64() * 1000.0, + ); + } + Ok(make_results(results)) + } +} + +const RAW_SCORE_MATRIX_MIN_QUERY_COUNT: usize = 4; + +const RAW_SCORE_MATRIX_TARGET_ELEMENTS: usize = 1 << 20; + +const RAW_TOP_K_MIN_PARTITION_SIZE: usize = 1 << 12; + +async fn execute_vindex_searches( + io_meta: GlobalIndexIOMeta, + options: HashMap, + vector_searches: Vec, + source: S, + file_name: String, + index_parallelism: usize, + guard: G, +) -> crate::Result>>> { + let panic_context = if vector_searches.len() > 1 { + "vindex global-index batch search task failed" + } else { + "vindex global-index search task failed" + }; + execute_global_index_with_guard(panic_context, guard, move || { + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options) + .with_batch_index_parallelism(index_parallelism); + reader + .visit_batch_vector_search(&vector_searches, |_| Ok(source)) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to read vindex index file '{}': {}", file_name, e), + source: Some(Box::new(e)), + }) + }) + .await +} + +#[derive(Clone, Copy)] +struct VectorSearchEvaluation<'a> { + table: Option<&'a Table>, + file_io: &'a FileIO, + table_path: &'a str, + table_options: &'a HashMap, + schema_fields: &'a [DataField], + next_row_id: Option, +} + +#[derive(Default)] +struct IndexSearchTiming { + permit_wait: Duration, + file_reader_open: Duration, +} + +#[cfg(test)] +async fn evaluate_vector_search( + evaluation: VectorSearchEvaluation<'_>, + index_entries: &[IndexManifestEntry], + vector_search: &VectorSearch, +) -> crate::Result> { + let results = evaluate_batch_vector_search( + evaluation, + index_entries, + std::slice::from_ref(vector_search), + ) + .await?; + crate::table::vector_search_common::take_only_result(results, "vector search")?.to_row_ranges() +} + +async fn evaluate_batch_vector_search( + evaluation: VectorSearchEvaluation<'_>, + index_entries: &[IndexManifestEntry], + vector_searches: &[VectorSearch], +) -> crate::Result> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); + if vector_searches.is_empty() { + return Ok(Vec::new()); + } + + let table_path = evaluation.table_path.trim_end_matches('/'); + let core_options = CoreOptions::new(evaluation.table_options); + let search_mode = core_options.vector_index_search_mode()?; + let field_name = &vector_searches[0].field_name; + if vector_searches + .iter() + .any(|vector_search| vector_search.field_name != *field_name) + { + return Err(crate::Error::DataInvalid { + message: "Batch vector search requires all query vectors to use the same field" + .to_string(), + source: None, + }); + } + let search_options = vector_searches[0].options.clone(); + if vector_searches + .iter() + .any(|vector_search| vector_search.options != search_options) + { + return Err(crate::Error::DataInvalid { + message: "Batch vector search requires all query vectors to use the same options" + .to_string(), + source: None, + }); + } + + let field_id = match find_field_id_by_name(evaluation.schema_fields, field_name) { + Some(id) => id, + None => return Ok(vec![ScoredRowIds::empty(); vector_searches.len()]), + }; + + let vector_entries: Vec<_> = index_entries + .iter() + .filter(|e| { + e.kind == FileKind::Add + && VectorIndexBackend::from_index_type(&e.index_file.index_type).is_some() + && e.index_file + .global_index_meta + .as_ref() + .is_some_and(|m| m.index_field_id == field_id) + }) + .collect(); + + if vector_entries.is_empty() && search_mode == GlobalIndexSearchMode::Fast { + return Ok(vec![ScoredRowIds::empty(); vector_searches.len()]); + } + + let deletion_vector_start = timing_enabled.then(Instant::now); + let deleted_row_index = if core_options.data_evolution_enabled() { + match evaluation.table { + Some(table) => { + let ranges = + deleted_row_ranges_for_data_evolution_dvs(table, index_entries).await?; + (!ranges.is_empty()).then(|| RowRangeIndex::create(ranges)) + } + None => None, + } + } else { + None + }; + let deletion_vector = deletion_vector_start.map_or(Duration::ZERO, |start| start.elapsed()); + + let max_limit = vector_searches + .iter() + .map(|vector_search| vector_search.limit) + .max() + .unwrap_or(0); + let refine_factor = match vector_entries.first() { + Some(entry) => configured_refine_factor( + &search_options, + evaluation.table_options, + field_name, + &entry.index_file.index_type, + )?, + None => 0, + }; + let index_search_limit = indexed_search_limit(max_limit, refine_factor)?; + + let vector_entry_count = vector_entries.len(); + let vector_search_plans = if let Some(include_row_ids) = + shared_batch_include_row_ids(vector_searches) + { + let ranges = vector_entries + .iter() + .map(|entry| { + let meta = entry.index_file.global_index_meta.as_ref().ok_or_else(|| { + crate::Error::DataInvalid { + message: format!( + "Vector index '{}' is missing global index metadata", + entry.index_file.file_name + ), + source: None, + } + })?; + Ok((meta.row_range_start, meta.row_range_end)) + }) + .collect::>>()?; + vector_entries + .iter() + .copied() + .zip(localize_shared_include_row_ids( + include_row_ids.as_ref(), + &ranges, + )?) + .filter_map(|(entry, local_filter)| local_filter.map(|filter| (entry, Some(filter)))) + .collect::>() + } else { + vector_entries + .iter() + .copied() + .map(|entry| (entry, None)) + .collect::>() + }; + let mut permit_wait = Duration::ZERO; + let mut file_reader_open = Duration::ZERO; + let mut index_search = Duration::ZERO; + let mut merge = Duration::ZERO; + let mut refine = Duration::ZERO; + let mut raw_fallback = Duration::ZERO; + let mut merged = vec![ScoredRowIds::empty(); vector_searches.len()]; + if !vector_entries.is_empty() { + let index_search_start = timing_enabled.then(Instant::now); + let concurrency = core_options.global_index_thread_num()?; + if concurrency > tokio::sync::Semaphore::MAX_PERMITS { + return Err(crate::Error::DataInvalid { + message: format!( + "Global index thread count must not exceed {}", + tokio::sync::Semaphore::MAX_PERMITS + ), + source: None, + }); + } + ensure_global_index_executor_capacity(concurrency); + let vindex_entry_count = vector_entries + .iter() + .filter(|entry| is_vindex_index_type(&entry.index_file.index_type)) + .count(); + let (batch_index_parallelism, range_read_limiter) = if vindex_entry_count == 0 { + (1, None) + } else { + let (index_parallelism, range_read_concurrency) = + vindex_concurrency_limits(&core_options, vindex_entry_count, concurrency)?; + ( + index_parallelism, + Some(RangeReadLimiter::new(range_read_concurrency)), + ) + }; + let futures: Vec<_> = vector_search_plans + .into_iter() + .map(|(entry, shared_local_filter)| { + let range_read_limiter = range_read_limiter.clone(); + 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 = IndexFileLocation::Global { table_path } + .resolve(&entry.index_file.file_name, entry.index_file.external_path.as_deref()); + let file_name = entry.index_file.file_name.clone(); + let file_size = entry.index_file.file_size as u64; + let index_meta_bytes = global_meta.index_meta.clone().unwrap_or_default(); + let row_range_start = global_meta.row_range_start; + let row_range_end = global_meta.row_range_end; + let index_limit = search_limit_with_deleted_rows( + index_search_limit, + row_range_start, + row_range_end, + deleted_row_index.as_ref(), + ) + .min(i32::MAX as usize); + let mut vector_searches = vector_searches.to_vec(); + for vector_search in &mut vector_searches { + vector_search.limit = index_limit; + } + let mut options = evaluation.table_options.clone(); + options.extend(search_options.clone()); + let input = evaluation.file_io.new_input(&path); + async move { + if let Some(local_filter) = shared_local_filter { + let local_filter = Arc::new(local_filter); + for vector_search in &mut vector_searches { + vector_search + .set_shared_include_row_ids(Arc::clone(&local_filter)); + } + } else { + for vector_search in &mut vector_searches { + if let Some(include_row_ids) = + vector_search.effective_include_row_ids() + { + vector_search.set_shared_include_row_ids(Arc::new( + localize_include_row_ids( + include_row_ids, + row_range_start, + row_range_end, + )?, + )); + } + } + } + if vector_searches.iter().all(|search| { + search + .effective_include_row_ids() + .is_some_and(|row_ids| row_ids.is_empty()) + }) { + return Ok(( + vec![ScoredRowIds::empty(); vector_searches.len()], + IndexSearchTiming::default(), + )); + } + let permit_start = timing_enabled.then(Instant::now); + let permit = acquire_process_global_search_permit(concurrency).await?; + let permit_wait = + permit_start.map_or(Duration::ZERO, |start| start.elapsed()); + let input = input?; + let query_count = vector_searches.len(); + let mut file_reader_open = Duration::ZERO; + let mut full_file_read = None; + let io_meta = + GlobalIndexIOMeta::new(file_name.clone(), file_size, index_meta_bytes); + let results = match backend { + VectorIndexBackend::Lumina => { + let read_start = timing_enabled.then(Instant::now); + let data = input.read().await.map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "Failed to read {} index file '{}': {}", + backend.error_name(), + file_name, + e + ), + source: None, + } + })?; + if let Some(start) = read_start { + full_file_read = Some((start.elapsed(), data.len())); + } + execute_global_index_with_guard( + "Lumina global-index batch search task failed", + permit, + move || { + let mut reader = + LuminaVectorGlobalIndexReader::new(io_meta, options); + reader.visit_batch_vector_search(&vector_searches, |_| { + Ok(Cursor::new(data)) + }) + }, + ) + .await? + } + VectorIndexBackend::Vindex => { + match tokio::runtime::Handle::try_current() { + Ok(runtime) => { + let file_reader_open_start = + timing_enabled.then(Instant::now); + let file_reader = input.reader().await.map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "Failed to open vindex file '{}' for range reads: {}", + file_name, e + ), + source: None, + } + })?; + file_reader_open = file_reader_open_start + .map_or(Duration::ZERO, |start| start.elapsed()); + let source = VindexFileReader::new_with_limiter( + Arc::new(file_reader), + runtime, + range_read_limiter.expect("Vindex range-read limiter"), + file_size, + file_name.clone(), + ); + let range_io_stats = source.range_io_stats(); + let results = execute_vindex_searches( + io_meta, + options, + vector_searches, + source, + file_name.clone(), + batch_index_parallelism, + permit, + ) + .await?; + if let Some(stats) = range_io_stats { + log_vindex_range_io_stats( + &file_name, + query_count, + &stats, + ); + } + results + } + Err(_) if query_count > 1 => { + let read_start = timing_enabled.then(Instant::now); + let data = input.read().await.map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "Failed to read vindex index file '{}': {}", + file_name, e + ), + source: None, + } + })?; + if let Some(start) = read_start { + full_file_read = Some((start.elapsed(), data.len())); + } + execute_vindex_searches( + io_meta, + options, + vector_searches, + Cursor::new(data), + file_name.clone(), + batch_index_parallelism, + permit, + ) + .await? + } + Err(error) => { + return Err(crate::Error::UnexpectedError { + message: + "Vector index range reader requires a Tokio runtime" + .to_string(), + source: Some(Box::new(error)), + }); + } + } + } + }; + if let Some((read, returned_bytes)) = full_file_read { + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_full_file_io backend={} file={} nq={} requested_bytes={} returned_bytes={} read_ms={:.3}", + backend.error_name(), + file_name, + query_count, + file_size, + returned_bytes, + read.as_secs_f64() * 1000.0, + ); + } + if results.len() != query_count { + return Err(crate::Error::DataInvalid { + message: format!( + "Batch vector search backend returned {} results for {} query vectors", + results.len(), + query_count + ), + source: None, + }); + } + + Ok::<_, crate::Error>(( + results + .into_iter() + .map(|result| match result { + Some(scored_map) => ScoredRowIds::from_scored_map(scored_map) + .offset(row_range_start), + None => ScoredRowIds::empty(), + }) + .collect::>(), + IndexSearchTiming { + permit_wait, + file_reader_open, + }, + )) + } + }) + .collect(); + + let results = drain_indexed_jobs(futures.into_iter(), concurrency).await?; + index_search = index_search_start.map_or(Duration::ZERO, |start| start.elapsed()); + let merge_start = timing_enabled.then(Instant::now); + for (per_entry, entry_timing) in &results { + permit_wait = permit_wait.saturating_add(entry_timing.permit_wait); + file_reader_open = file_reader_open.saturating_add(entry_timing.file_reader_open); + for (query_index, result) in per_entry.iter().enumerate() { + merged[query_index] = merged[query_index].or(result); + } + } + merge = merge_start.map_or(Duration::ZERO, |start| start.elapsed()); + } + + if refine_factor != 0 { + let refine_start = timing_enabled.then(Instant::now); + merged = maybe_rerank_indexed_batch_results( + evaluation, + index_entries, + field_id, + field_name, + vector_searches, + merged, + index_search_limit, + ) + .await?; + refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); + } + + if search_mode != GlobalIndexSearchMode::Fast { + let raw_fallback_start = timing_enabled.then(Instant::now); + let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail { + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Vector raw search in detail mode requires table context".to_string(), + source: None, + })?; + detail_data_ranges_for_table(table).await? + } else { + Vec::new() + }; + let field_ids = HashSet::from([field_id]); + let raw_ranges = unindexed_ranges_for_global_index_entries( + index_entries, + &field_ids, + search_mode, + evaluation.next_row_id, + &detail_ranges, + is_vector_global_index_file, + ); + if !raw_ranges.is_empty() { + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Vector raw search requires table context".to_string(), + source: None, + })?; + let metric_start = timing_enabled.then(Instant::now); + let metric = resolve_raw_vector_metric( + evaluation.file_io, + table_path, + evaluation.table_options, + index_entries, + field_id, + field_name, + ) + .await?; + let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); + let (raw_results, raw_timing) = + read_raw_batch_vector_search(table, vector_searches, &raw_ranges, metric).await?; + if let Some(raw_timing) = raw_timing { + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_raw_fallback nq={} row_ranges={} metric_resolve_ms={:.3} raw_plan_ms={:.3} split_count={} file_count={} raw_stream_wait_ms={:.3} raw_score_cpu_ms={:.3} arrow_batches={} arrow_rows={} total_raw_read_ms={:.3}", + vector_searches.len(), + raw_ranges.len(), + metric_resolve.as_secs_f64() * 1000.0, + raw_timing.plan.as_secs_f64() * 1000.0, + raw_timing.split_count, + raw_timing.file_count, + raw_timing.stream_wait.as_secs_f64() * 1000.0, + raw_timing.score_cpu.as_secs_f64() * 1000.0, + raw_timing.batch_count, + raw_timing.row_count, + raw_timing.total.as_secs_f64() * 1000.0, + ); + } + for (query_index, result) in raw_results.iter().enumerate() { + merged[query_index] = merged[query_index].or(result); + } + } + raw_fallback = raw_fallback_start.map_or(Duration::ZERO, |start| start.elapsed()); + } + + let finalize_start = timing_enabled.then(Instant::now); + let results = merged + .into_iter() + .zip(vector_searches) + .map(|(result, vector_search)| { + Ok(result + .without_deleted_row_ranges(deleted_row_index.as_ref())? + .top_k(vector_search.limit)) + }) + .collect::>>()?; + let finalize = finalize_start.map_or(Duration::ZERO, |start| start.elapsed()); + if let Some(total_start) = total_start { + let total = total_start.elapsed(); + let children = deletion_vector + .saturating_add(index_search) + .saturating_add(merge) + .saturating_add(refine) + .saturating_add(raw_fallback) + .saturating_add(finalize); + let result_count = results + .iter() + .map(|result| result.row_ids.len()) + .sum::(); + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_search_evaluate nq={} index_entries={} index_files={} result_count={} refine_factor={} total_ms={:.3} deletion_vector_ms={:.3} index_search_ms={:.3} global_permit_wait_sum_ms={:.3} file_reader_open_sum_ms={:.3} merge_ms={:.3} refine_ms={:.3} raw_fallback_ms={:.3} finalize_ms={:.3} unattributed_ms={:.3}", + vector_searches.len(), + index_entries.len(), + vector_entry_count, + result_count, + refine_factor, + total.as_secs_f64() * 1000.0, + deletion_vector.as_secs_f64() * 1000.0, + index_search.as_secs_f64() * 1000.0, + permit_wait.as_secs_f64() * 1000.0, + file_reader_open.as_secs_f64() * 1000.0, + merge.as_secs_f64() * 1000.0, + refine.as_secs_f64() * 1000.0, + raw_fallback.as_secs_f64() * 1000.0, + finalize.as_secs_f64() * 1000.0, + total.saturating_sub(children).as_secs_f64() * 1000.0, + ); + } + Ok(results) +} + +fn is_vector_global_index_file(index_file: &IndexFileMeta) -> bool { + VectorIndexBackend::from_index_type(&index_file.index_type).is_some() +} + +/// Collect materialized DE rows, join each row's `(rank, score)` by its global +/// `_ROW_ID`, reorder to the search rank order, append the `__paimon_search_score` +/// column, and drop `_ROW_ID`. Every row must map to a search candidate and the +/// total materialized count must equal `expected_len`; a miss or count mismatch +/// fails loud rather than silently dropping or NaN-scoring a row. Empty input +/// yields no batches. +fn attach_scores_by_row_id( + batches: &[RecordBatch], + rank_score_of: &HashMap, + expected_len: usize, +) -> crate::Result> { + // (rank, batch_index, row_index, score) per materialized row. + let mut ranked: Vec<(usize, usize, usize, f32)> = Vec::new(); + for (batch_index, batch) in batches.iter().enumerate() { + let row_id_idx = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|_| crate::Error::DataInvalid { + message: format!("materialized batch missing {ROW_ID_FIELD_NAME} column"), + source: None, + })?; + let col = batch.column(row_id_idx); + let ids = + col.as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("{ROW_ID_FIELD_NAME} column is not Int64"), + source: None, + })?; + for row_index in 0..batch.num_rows() { + if ids.is_null(row_index) { + return Err(crate::Error::DataInvalid { + message: format!( + "materialized DE vector row has null {ROW_ID_FIELD_NAME}; cannot align score" + ), + source: None, + }); + } + let id = ids.value(row_index); + let (rank, score) = + *rank_score_of + .get(&id) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "materialized DE vector row (row id {id}) has no matching search candidate" + ), + source: None, + })?; + ranked.push((rank, batch_index, row_index, score)); + } + } + + if ranked.len() != expected_len { + return Err(crate::Error::DataInvalid { + message: format!( + "DE vector materialization produced {} rows but search returned {expected_len}", + ranked.len() + ), + source: None, + }); + } + if ranked.is_empty() { + return Ok(Vec::new()); + } + + ranked.sort_by_key(|r| r.0); + let indices: Vec<(usize, usize)> = ranked.iter().map(|r| (r.1, r.2)).collect(); + let refs: Vec<&RecordBatch> = batches.iter().collect(); + let reordered = + interleave_record_batch(&refs, &indices).map_err(|e| crate::Error::DataInvalid { + message: format!("failed to reorder DE vector search rows: {e}"), + source: None, + })?; + + // Drop _ROW_ID. + let row_id_idx = reordered + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|_| crate::Error::DataInvalid { + message: format!("reordered batch missing {ROW_ID_FIELD_NAME} column"), + source: None, + })?; + let keep: Vec = (0..reordered.num_columns()) + .filter(|i| *i != row_id_idx) + .collect(); + let stripped = reordered + .project(&keep) + .map_err(|e| crate::Error::DataInvalid { + message: format!("failed to drop {ROW_ID_FIELD_NAME} column: {e}"), + source: None, + })?; + + // Append the score column in rank order. + let scores: Vec = ranked.iter().map(|r| r.3).collect(); + let score_array: Arc = Arc::new(Float32Array::from(scores)); + let mut fields: Vec> = + stripped.schema().fields().iter().cloned().collect(); + fields.push(Arc::new(arrow_schema::Field::new( + SEARCH_SCORE_COLUMN, + arrow_schema::DataType::Float32, + false, + ))); + let out_schema = Arc::new(arrow_schema::Schema::new(fields)); + let mut columns = stripped.columns().to_vec(); + columns.push(score_array); + let out = RecordBatch::try_new(out_schema, columns).map_err(|e| crate::Error::DataInvalid { + message: format!("failed to append DE vector score column: {e}"), + source: None, + })?; + Ok(vec![out]) +} + +async fn maybe_rerank_indexed_batch_results( + evaluation: VectorSearchEvaluation<'_>, + index_entries: &[IndexManifestEntry], + field_id: i32, + field_name: &str, + vector_searches: &[VectorSearch], + results: Vec, + index_search_limit: usize, +) -> crate::Result> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); + let mut candidate_searches = Vec::with_capacity(vector_searches.len()); + let mut candidate_results = Vec::with_capacity(vector_searches.len()); + let mut union_candidates = RoaringTreemap::new(); + let mut candidate_references = 0usize; + + for (result, vector_search) in results.into_iter().zip(vector_searches) { + let candidates = result.top_k(index_search_limit); + candidate_references = candidate_references.saturating_add(candidates.row_ids.len()); + let mut include_row_ids = RoaringTreemap::new(); + for &row_id in &candidates.row_ids { + include_row_ids.insert(row_id); + union_candidates.insert(row_id); + } + + let mut candidate_search = vector_search.clone(); + candidate_search.set_shared_include_row_ids(Arc::new(include_row_ids)); + candidate_searches.push(candidate_search); + candidate_results.push(candidates); + } + + if union_candidates.iter().next().is_none() { + return Ok(candidate_results); + } + + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Vector index rerank requires table context".to_string(), + source: None, + })?; + let unique_candidates = union_candidates.len(); + let raw_ranges = sorted_row_ids_to_row_ranges(union_candidates.iter())?; + let metric_start = timing_enabled.then(Instant::now); + let metric = resolve_raw_vector_metric( + evaluation.file_io, + evaluation.table_path.trim_end_matches('/'), + evaluation.table_options, + index_entries, + field_id, + field_name, + ) + .await?; + let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); + + let (results, raw_timing) = + read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric).await?; + if let (Some(total_start), Some(raw_timing)) = (total_start, raw_timing) { + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_refine nq={} candidate_references={} unique_candidates={} row_ranges={} metric_resolve_ms={:.3} raw_plan_ms={:.3} split_count={} file_count={} raw_stream_wait_ms={:.3} raw_score_cpu_ms={:.3} arrow_batches={} arrow_rows={} total_refine_ms={:.3}", + vector_searches.len(), + candidate_references, + unique_candidates, + raw_ranges.len(), + metric_resolve.as_secs_f64() * 1000.0, + raw_timing.plan.as_secs_f64() * 1000.0, + raw_timing.split_count, + raw_timing.file_count, + raw_timing.stream_wait.as_secs_f64() * 1000.0, + raw_timing.score_cpu.as_secs_f64() * 1000.0, + raw_timing.batch_count, + raw_timing.row_count, + total_start.elapsed().as_secs_f64() * 1000.0, + ); + } + Ok(results) +} + +fn sorted_row_ids_to_row_ranges( + row_ids: impl IntoIterator, +) -> crate::Result> { + let mut row_ids = row_ids.into_iter(); + let Some(first) = row_ids.next() else { + return Ok(Vec::new()); + }; + let mut start = row_id_to_i64_for_range(first)?; + let mut end = start; + let mut ranges = Vec::new(); + for row_id in row_ids { + let row_id = row_id_to_i64_for_range(row_id)?; + if end.checked_add(1) == Some(row_id) { + end = row_id; + } else { + ranges.push(RowRange::new(start, end)); + start = row_id; + end = row_id; + } + } + ranges.push(RowRange::new(start, end)); + Ok(ranges) +} + +fn row_id_to_i64_for_range(row_id: u64) -> crate::Result { + i64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { + message: format!( + "Vector search row id {row_id} exceeds i64::MAX and cannot be converted to RowRange" + ), + source: None, + }) +} + +fn shared_batch_include_row_ids(vector_searches: &[VectorSearch]) -> Option<&Arc> { + let first = vector_searches.first()?.shared_include_row_ids.as_ref()?; + vector_searches + .iter() + .skip(1) + .all(|search| { + search + .shared_include_row_ids + .as_ref() + .is_some_and(|include_row_ids| Arc::ptr_eq(first, include_row_ids)) + }) + .then_some(first) +} + +fn prune_raw_ranges_by_include_row_ids( + raw_ranges: &[RowRange], + vector_searches: &[VectorSearch], +) -> crate::Result> { + if vector_searches + .iter() + .any(|search| search.effective_include_row_ids().is_none()) + { + return Ok(raw_ranges.to_vec()); + } + + let include_ranges = + if let Some(include_row_ids) = shared_batch_include_row_ids(vector_searches) { + sorted_row_ids_to_row_ranges(include_row_ids.iter())? + } else { + let mut union = RoaringTreemap::new(); + for include_row_ids in vector_searches + .iter() + .filter_map(VectorSearch::effective_include_row_ids) + { + for row_id in include_row_ids.iter() { + union.insert(row_id); + } + } + sorted_row_ids_to_row_ranges(union.iter())? + }; + Ok(intersect_sorted_ranges(raw_ranges, &include_ranges)) +} + +fn localize_include_row_ids( + include_row_ids: &RoaringTreemap, + row_range_start: i64, + row_range_end: i64, +) -> crate::Result { + let start = u64::try_from(row_range_start).map_err(|_| crate::Error::DataInvalid { + message: format!("Negative vector index row range start: {row_range_start}"), + source: None, + })?; + let end = u64::try_from(row_range_end).map_err(|_| crate::Error::DataInvalid { + message: format!("Negative vector index row range end: {row_range_end}"), + source: None, + })?; + let mut localized = RoaringTreemap::new(); + for row_id in include_row_ids.iter() { + if row_id >= start && row_id <= end { + localized.insert(row_id - start); + } + } + Ok(localized) +} + +fn localize_shared_include_row_ids( + include_row_ids: &RoaringTreemap, + ranges: &[(i64, i64)], +) -> crate::Result>> { + let mut validated_ranges = Vec::with_capacity(ranges.len()); + for (index, &(start, end)) in ranges.iter().enumerate() { + if start < 0 || end < start { + return Err(crate::Error::DataInvalid { + message: format!("Invalid vector index row range [{start}, {end}]"), + source: None, + }); + } + validated_ranges.push((start as u64, end as u64, index)); + } + validated_ranges.sort_unstable_by_key(|(start, _, _)| *start); + + let mut localized = (0..ranges.len()) + .map(|_| RoaringTreemap::new()) + .collect::>(); + let mut active = Vec::::new(); + let mut next_range = 0usize; + for row_id in include_row_ids.iter() { + while next_range < validated_ranges.len() && validated_ranges[next_range].0 <= row_id { + active.push(next_range); + next_range += 1; + } + active.retain(|range_index| validated_ranges[*range_index].1 >= row_id); + for range_index in &active { + let (start, _, original_index) = validated_ranges[*range_index]; + localized[original_index].insert(row_id - start); + } + if next_range == validated_ranges.len() && active.is_empty() { + break; + } + } + + Ok(localized + .into_iter() + .map(|filter| (!filter.is_empty()).then_some(filter)) + .collect()) +} + +async fn detail_data_ranges_for_table(table: &Table) -> crate::Result> { + let plan = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan() + .await?; + let mut ranges = Vec::new(); + for split in plan.splits() { + for file in split.data_files() { + if let Some((from, to)) = file.row_id_range() { + ranges.push(RowRange::new(from, to)); + } + } + } + Ok(merge_row_ranges(ranges)) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RawVectorMetric { + L2, + Cosine, + InnerProduct, +} + +impl RawVectorMetric { + fn parse(value: &str) -> crate::Result { + Self::parse_normalized(&normalize_metric(value)).ok_or_else(|| crate::Error::DataInvalid { + message: format!("Unknown vector search metric: {value}"), + source: None, + }) + } + + fn parse_normalized(value: &str) -> Option { + match value { + "l2" => Some(Self::L2), + "cosine" => Some(Self::Cosine), + "inner_product" => Some(Self::InnerProduct), + _ => None, + } + } + + fn from_lumina(metric: LuminaVectorMetric) -> Self { + match metric { + LuminaVectorMetric::L2 => Self::L2, + LuminaVectorMetric::Cosine => Self::Cosine, + LuminaVectorMetric::InnerProduct => Self::InnerProduct, + } + } + + fn from_vindex(metric: MetricType) -> Self { + match metric { + MetricType::L2 => Self::L2, + MetricType::Cosine => Self::Cosine, + MetricType::InnerProduct => Self::InnerProduct, + } + } +} + +async fn resolve_raw_vector_metric( + file_io: &FileIO, + table_path: &str, + table_options: &HashMap, + index_entries: &[IndexManifestEntry], + field_id: i32, + field_name: &str, +) -> crate::Result { + for entry in index_entries { + if entry.kind != FileKind::Add { + continue; + } + let Some(global_meta) = entry.index_file.global_index_meta.as_ref() else { + continue; + }; + if global_meta.index_field_id != field_id { + continue; + } + let Some(backend) = VectorIndexBackend::from_index_type(&entry.index_file.index_type) + else { + continue; + }; + match backend { + VectorIndexBackend::Lumina => { + if let Some(index_meta) = global_meta.index_meta.as_ref() { + if !index_meta.is_empty() { + let metric = LuminaIndexMeta::deserialize(index_meta)?.metric()?; + return Ok(RawVectorMetric::from_lumina(metric)); + } + } + } + VectorIndexBackend::Vindex => { + if let Some(index_meta) = global_meta.index_meta.as_ref() { + if let Ok(options) = + serde_json::from_slice::>(index_meta) + { + if let Some(metric) = options.get("metric") { + if let Some(metric) = + RawVectorMetric::parse_normalized(&normalize_metric(metric)) + { + return Ok(metric); + } + } + } + } + let path = IndexFileLocation::Global { table_path }.resolve( + &entry.index_file.file_name, + entry.index_file.external_path.as_deref(), + ); + let input = file_io.new_input(&path)?; + let read_error = |e| crate::Error::DataInvalid { + message: format!( + "Failed to read vindex index file '{}' for raw search metric: {}", + entry.index_file.file_name, e + ), + source: Some(Box::new(e)), + }; + let header_size = if entry.index_file.file_size > 0 { + (entry.index_file.file_size as u64).min(DISKANN_HEADER_SIZE as u64) + } else { + input + .metadata() + .await + .map_err(&read_error)? + .size + .min(DISKANN_HEADER_SIZE as u64) + }; + let file_reader = input.reader().await.map_err(&read_error)?; + let bytes = file_reader.read(0..header_size).await.map_err(read_error)?; + let reader = VIndexReader::open(Cursor::new(bytes)).map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "Failed to open paimon-vindex-core reader for raw search metric: {}", + e + ), + source: Some(Box::new(e)), + } + })?; + return Ok(RawVectorMetric::from_vindex(reader.metadata().metric)); + } + } + } + + configured_raw_vector_metric(table_options, field_name) +} + +fn configured_raw_vector_metric( + options: &HashMap, + field_name: &str, +) -> crate::Result { + let direct_keys = [ + format!("fields.{field_name}.distance.metric"), + format!("fields.{field_name}.metric"), + "test.vector.metric".to_string(), + "lumina.distance.metric".to_string(), + "distance.metric".to_string(), + "metric".to_string(), + ]; + for key in direct_keys { + if let Some(value) = options.get(&key) { + return RawVectorMetric::parse(value); + } + } + + let mut inferred = None; + for (key, value) in options { + if !(key.ends_with(".distance.metric") || key.ends_with(".metric")) { + continue; + } + let normalized = normalize_metric(value); + let Some(metric) = RawVectorMetric::parse_normalized(&normalized) else { + continue; + }; + if let Some(existing) = inferred { + if existing != metric { + return Ok(RawVectorMetric::L2); + } + } else { + inferred = Some(metric); + } + } + Ok(inferred.unwrap_or(RawVectorMetric::L2)) +} + +#[derive(Default)] +struct RawVectorReadTiming { + plan: Duration, + stream_wait: Duration, + score_cpu: Duration, + total: Duration, + split_count: usize, + file_count: usize, + batch_count: usize, + row_count: usize, +} + +async fn read_raw_batch_vector_search( + table: &Table, + vector_searches: &[VectorSearch], + raw_ranges: &[RowRange], + metric: RawVectorMetric, +) -> crate::Result<(Vec, Option)> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); + if vector_searches.is_empty() { + return Ok((Vec::new(), None)); + } + if raw_ranges.is_empty() { + return Ok((vec![ScoredRowIds::empty(); vector_searches.len()], None)); + } + let raw_ranges = prune_raw_ranges_by_include_row_ids(raw_ranges, vector_searches)?; + if raw_ranges.is_empty() { + return Ok((vec![ScoredRowIds::empty(); vector_searches.len()], None)); + } + + let field_name = &vector_searches[0].field_name; + if vector_searches + .iter() + .any(|vector_search| vector_search.field_name != *field_name) + { + return Err(crate::Error::DataInvalid { + message: "Batch vector raw search requires all query vectors to use the same field" + .to_string(), + source: None, + }); + } + + let plan_start = timing_enabled.then(Instant::now); + let mut read_builder = table.new_read_builder(); + read_builder + .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])? + .with_row_ranges(raw_ranges); + let plan = read_builder.new_scan().plan().await?; + let plan_elapsed = plan_start.map_or(Duration::ZERO, |start| start.elapsed()); + let split_count = plan.splits().len(); + let file_count = plan + .splits() + .iter() + .map(|split| split.data_files().len()) + .sum(); + if plan.splits().is_empty() { + return Ok(( + vec![ScoredRowIds::empty(); vector_searches.len()], + total_start.map(|start| RawVectorReadTiming { + plan: plan_elapsed, + total: start.elapsed(), + ..RawVectorReadTiming::default() + }), + )); + } + let read = read_builder.new_read()?; + let mut stream = read.to_arrow(plan.splits())?; + + let scoring_plan = RawScoringPlan::new(vector_searches, metric); + let mut top_k = vector_searches + .iter() + .map(|vector_search| RawScoreTopK::new(vector_search.limit)) + .collect::>(); + let mut timing = timing_enabled.then(|| RawVectorReadTiming { + plan: plan_elapsed, + split_count, + file_count, + ..RawVectorReadTiming::default() + }); + loop { + let stream_wait_start = timing_enabled.then(Instant::now); + let batch = stream.try_next().await?; + if let (Some(timing), Some(stream_wait_start)) = (&mut timing, stream_wait_start) { + timing.stream_wait = timing + .stream_wait + .saturating_add(stream_wait_start.elapsed()); + } + let Some(batch) = batch else { + break; + }; + if let Some(timing) = &mut timing { + timing.batch_count += 1; + timing.row_count = timing.row_count.saturating_add(batch.num_rows()); + } + let score_start = timing_enabled.then(Instant::now); + collect_raw_batch_vector_batch(&batch, vector_searches, metric, &scoring_plan, &mut top_k)?; + if let (Some(timing), Some(score_start)) = (&mut timing, score_start) { + timing.score_cpu = timing.score_cpu.saturating_add(score_start.elapsed()); + } + } + + if let (Some(timing), Some(total_start)) = (&mut timing, total_start) { + timing.total = total_start.elapsed(); + } + Ok(( + top_k + .into_iter() + .map(RawScoreTopK::into_search_result) + .collect(), + timing, + )) +} + +struct RawScoringPlan { + all_query_indices: Vec, + shared_filter_groups: Vec, + candidate_query_indices: HashMap>, + query_l2_squared_norms: Vec, + dense_query_dimension: Option, + dense_query_matrix: Option>, +} + +struct SharedRawFilterGroup { + include_row_ids: Arc, + query_indices: Vec, +} + +impl RawScoringPlan { + fn new(vector_searches: &[VectorSearch], metric: RawVectorMetric) -> Self { + let mut all_query_indices = Vec::new(); + let mut shared_filter_groups = Vec::new(); + let mut candidate_query_indices: HashMap> = HashMap::new(); + let query_l2_squared_norms = vector_searches + .iter() + .map(|vector_search| match metric { + RawVectorMetric::L2 | RawVectorMetric::Cosine => vector_search + .vector + .iter() + .map(|value| value * value) + .sum::(), + RawVectorMetric::InnerProduct => 0.0, + }) + .collect(); + + if let Some(include_row_ids) = shared_batch_include_row_ids(vector_searches) { + shared_filter_groups.push(SharedRawFilterGroup { + include_row_ids: Arc::clone(include_row_ids), + query_indices: (0..vector_searches.len()).collect(), + }); + } else { + for (query_index, vector_search) in vector_searches.iter().enumerate() { + if let Some(include_row_ids) = vector_search.effective_include_row_ids() { + for row_id in include_row_ids.iter() { + candidate_query_indices + .entry(row_id) + .or_default() + .push(query_index); + } + } else { + all_query_indices.push(query_index); + } + } + } + + let dense_query_dimension = all_query_indices + .first() + .map(|&query_index| vector_searches[query_index].vector.len()); + let dense_query_matrix = dense_query_dimension.and_then(|dimension| { + all_query_indices + .iter() + .all(|&query_index| vector_searches[query_index].vector.len() == dimension) + .then(|| { + let mut matrix = + Vec::with_capacity(all_query_indices.len().saturating_mul(dimension)); + for &query_index in &all_query_indices { + matrix.extend_from_slice(&vector_searches[query_index].vector); + } + matrix + }) + }); + + Self { + all_query_indices, + shared_filter_groups, + candidate_query_indices, + query_l2_squared_norms, + dense_query_dimension, + dense_query_matrix, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +struct RawScoredRow { + row_id: u64, + score: f32, +} + +impl RawScoredRow { + fn strongest_first(a: &Self, b: &Self) -> Ordering { + b.score + .total_cmp(&a.score) + .then_with(|| a.row_id.cmp(&b.row_id)) + } +} + +struct RawScoreTopK { + limit: usize, + candidates: Vec, +} + +impl RawScoreTopK { + fn new(limit: usize) -> Self { + Self { + limit, + candidates: Vec::with_capacity(limit.min(1024).saturating_add(1)), + } + } + + fn offer(&mut self, row_id: u64, score: f32) { + if self.limit == 0 { + return; + } + self.candidates.push(RawScoredRow { row_id, score }); + if self.candidates.len() >= self.partition_size() { + self.reduce_to_limit(); + } + } + + fn offer_many(&mut self, candidates: I) + where + I: IntoIterator, + { + if self.limit == 0 { + return; + } + self.candidates.extend(candidates); + if self.candidates.len() >= self.partition_size() { + self.reduce_to_limit(); + } + } + + fn partition_size(&self) -> usize { + self.limit + .saturating_mul(2) + .max(RAW_TOP_K_MIN_PARTITION_SIZE) + } + + fn reduce_to_limit(&mut self) { + if self.candidates.len() <= self.limit { + return; + } + // Partition only after a substantial candidate block has accumulated. + // Each partition is linear in its input, so all reductions are O(n) + // amortized; only the final K survivors are fully sorted. + self.candidates + .select_nth_unstable_by(self.limit, RawScoredRow::strongest_first); + self.candidates.truncate(self.limit); + } + + fn into_search_result(mut self) -> ScoredRowIds { + self.reduce_to_limit(); + self.candidates + .sort_unstable_by(RawScoredRow::strongest_first); + let rows = self.candidates; + let mut row_ids = Vec::with_capacity(rows.len()); + let mut scores = Vec::with_capacity(rows.len()); + for row in rows { + row_ids.push(row.row_id); + scores.push(row.score); + } + ScoredRowIds::new(row_ids, scores) + } +} + +fn collect_raw_batch_vector_batch( + batch: &RecordBatch, + vector_searches: &[VectorSearch], + metric: RawVectorMetric, + scoring_plan: &RawScoringPlan, + top_k_out: &mut [RawScoreTopK], +) -> crate::Result<()> { + if vector_searches.is_empty() { + return Ok(()); + } + if top_k_out.len() != vector_searches.len() { + return Err(crate::Error::DataInvalid { + message: "Raw batch vector search output buffers must match query vector count" + .to_string(), + source: None, + }); + } + + let field_name = &vector_searches[0].field_name; + if vector_searches + .iter() + .any(|vector_search| vector_search.field_name != *field_name) + { + return Err(crate::Error::DataInvalid { + message: "Batch vector raw search requires all query vectors to use the same field" + .to_string(), + source: None, + }); + } + + let vector_index = + batch + .schema() + .index_of(field_name) + .map_err(|e| crate::Error::DataInvalid { + message: format!( + "Vector column '{}' not found in raw search batch: {}", + field_name, e + ), + source: None, + })?; + let row_id_index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|e| crate::Error::DataInvalid { + message: format!("_ROW_ID column not found in raw search batch: {e}"), + source: None, + })?; + + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: "Vector raw search requires non-null Int64 _ROW_ID".to_string(), + source: None, + })?; + + let column = batch.column(vector_index); + enum VectorLayout<'a> { + List(&'a ListArray), + Fixed(&'a FixedSizeListArray), + } + let layout = if let Some(a) = column.as_any().downcast_ref::() { + VectorLayout::List(a) + } else if let Some(a) = column.as_any().downcast_ref::() { + VectorLayout::Fixed(a) + } else { + return Err(crate::Error::DataInvalid { + message: "Vector raw search requires Arrow List or FixedSizeList" + .to_string(), + source: None, + }); + }; + let values = match layout { + VectorLayout::List(a) => a.values(), + VectorLayout::Fixed(a) => a.values(), + } + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: "Vector raw search requires Float32 vector elements".to_string(), + source: None, + })?; + + let use_dense_matrix = scoring_plan.all_query_indices.len() >= RAW_SCORE_MATRIX_MIN_QUERY_COUNT; + let dense_dimension = use_dense_matrix + .then_some(scoring_plan.dense_query_dimension) + .flatten(); + let mut dense_row_ids = Vec::with_capacity(batch.num_rows()); + let mut dense_vectors = Vec::with_capacity( + batch + .num_rows() + .saturating_mul(dense_dimension.unwrap_or_default()), + ); + for row in 0..batch.num_rows() { + if row_ids.is_null(row) { + return Err(crate::Error::DataInvalid { + message: "Vector raw search found null _ROW_ID".to_string(), + source: None, + }); + } + let row_id = row_id_to_u64(row_ids.value(row))?; + let is_null = match layout { + VectorLayout::List(a) => a.is_null(row), + VectorLayout::Fixed(a) => a.is_null(row), + }; + if is_null { + continue; + } + + let (start, end) = match layout { + VectorLayout::List(a) => { + let offsets = a.value_offsets(); + (offsets[row] as usize, offsets[row + 1] as usize) + } + VectorLayout::Fixed(a) => { + let len = a.value_length() as usize; + let start = a.value_offset(row) as usize; + (start, start + len) + } + }; + ensure_raw_vector_values_not_null(values, start, end)?; + + let raw_row = RawVectorRow { + row_id, + values, + start, + end, + }; + if let Some(dimension) = dense_dimension { + ensure_raw_vector_dimension(end - start, dimension)?; + if scoring_plan.dense_query_matrix.is_none() { + let &query_index = scoring_plan + .all_query_indices + .iter() + .find(|&&query_index| vector_searches[query_index].vector.len() != dimension) + .expect("a missing dense matrix requires inconsistent query dimensions"); + ensure_raw_vector_dimension(dimension, vector_searches[query_index].vector.len())?; + } + dense_row_ids.push(row_id); + dense_vectors.extend_from_slice(&values.values()[start..end]); + } else { + for &query_index in &scoring_plan.all_query_indices { + offer_raw_vector_score( + raw_row, + query_index, + metric, + vector_searches, + scoring_plan, + top_k_out, + )?; + } + } + if let Some(query_indices) = scoring_plan.candidate_query_indices.get(&row_id) { + for &query_index in query_indices { + offer_raw_vector_score( + raw_row, + query_index, + metric, + vector_searches, + scoring_plan, + top_k_out, + )?; + } + } + for group in &scoring_plan.shared_filter_groups { + if group.include_row_ids.contains(row_id) { + for &query_index in &group.query_indices { + offer_raw_vector_score( + raw_row, + query_index, + metric, + vector_searches, + scoring_plan, + top_k_out, + )?; + } + } + } + } + + if !dense_row_ids.is_empty() { + let query_matrix = scoring_plan + .dense_query_matrix + .as_deref() + .expect("dense query dimensions were validated above"); + let dimension = dense_dimension.expect("dense rows require dense queries"); + let queries_per_chunk = (RAW_SCORE_MATRIX_TARGET_ELEMENTS / dense_row_ids.len()) + .max(1) + .min(scoring_plan.all_query_indices.len()); + for (query_chunk_index, query_indices) in scoring_plan + .all_query_indices + .chunks(queries_per_chunk) + .enumerate() + { + let query_start = query_chunk_index * queries_per_chunk * dimension; + let query_end = query_start + query_indices.len() * dimension; + let scores = compute_raw_vector_score_matrix( + &dense_vectors, + dense_row_ids.len(), + &query_matrix[query_start..query_end], + query_indices.len(), + dimension, + &scoring_plan.query_l2_squared_norms, + query_indices, + metric, + )?; + for (matrix_query_index, &query_index) in query_indices.iter().enumerate() { + let query_scores = &scores[matrix_query_index * dense_row_ids.len() + ..(matrix_query_index + 1) * dense_row_ids.len()]; + top_k_out[query_index].offer_many( + dense_row_ids + .iter() + .zip(query_scores) + .map(|(&row_id, &score)| RawScoredRow { row_id, score }), + ); + } + } + } + + Ok(()) +} + +fn ensure_raw_vector_dimension(stored_len: usize, query_len: usize) -> crate::Result<()> { + if stored_len != query_len { + return Err(crate::Error::DataInvalid { + message: format!( + "Query vector dimension mismatch: raw row has {}, but query has {}", + stored_len, query_len + ), + source: None, + }); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn compute_raw_vector_score_matrix( + stored_vectors: &[f32], + row_count: usize, + query_vectors: &[f32], + query_count: usize, + dimension: usize, + query_l2_squared_norms: &[f32], + query_indices: &[usize], + metric: RawVectorMetric, +) -> crate::Result> { + let score_count = + row_count + .checked_mul(query_count) + .ok_or_else(|| crate::Error::DataInvalid { + message: "Vector raw search score matrix is too large".to_string(), + source: None, + })?; + debug_assert_eq!(stored_vectors.len(), row_count * dimension); + debug_assert_eq!(query_vectors.len(), query_count * dimension); + debug_assert_eq!(query_indices.len(), query_count); + + let mut scores = vec![0.0; score_count]; + // Query × stored-vector^T produces a query-major score matrix. Each query's + // scores are contiguous, which feeds partial Top-K without strided reads. + sgemm_a_bt( + query_count, + row_count, + dimension, + 1.0, + query_vectors, + stored_vectors, + 0.0, + &mut scores, + ); + if metric == RawVectorMetric::InnerProduct { + return Ok(scores); + } + + let stored_l2_squared_norms = stored_vectors + .chunks_exact(dimension) + .map(|vector| vector.iter().map(|value| value * value).sum::()) + .collect::>(); + for (matrix_query_index, &query_index) in query_indices.iter().enumerate() { + for (row_index, &stored_l2_squared_norm) in stored_l2_squared_norms.iter().enumerate() { + let score = &mut scores[matrix_query_index * row_count + row_index]; + let query_l2_squared_norm = query_l2_squared_norms[query_index]; + *score = match metric { + RawVectorMetric::L2 => { + let squared_distance = + stored_l2_squared_norm + query_l2_squared_norm - 2.0 * *score; + // The norm/dot reconstruction loses the low-order difference when two + // large vectors are close. Estimate a conservative accumulation-error + // bound and preserve the former scalar semantics inside that region. + let roundoff_bound = (stored_l2_squared_norm.abs() + + query_l2_squared_norm.abs() + + 2.0 * score.abs()) + * f32::EPSILON + * (dimension as f32 + 2.0) + * 4.0; + if !squared_distance.is_finite() || squared_distance <= roundoff_bound { + let stored = + &stored_vectors[row_index * dimension..(row_index + 1) * dimension]; + let query = &query_vectors + [matrix_query_index * dimension..(matrix_query_index + 1) * dimension]; + compute_raw_vector_l2_score(query, stored) + } else { + 1.0 / (1.0 + squared_distance) + } + } + RawVectorMetric::Cosine => { + let denominator = stored_l2_squared_norm.sqrt() * query_l2_squared_norm.sqrt(); + if denominator == 0.0 { + 0.0 + } else { + *score / denominator + } + } + RawVectorMetric::InnerProduct => unreachable!(), + }; + } + } + Ok(scores) +} + +fn ensure_raw_vector_values_not_null( + values: &Float32Array, + start: usize, + end: usize, +) -> crate::Result<()> { + if values.null_count() == 0 { + return Ok(()); + } + for value_index in start..end { + if values.is_null(value_index) { + return Err(crate::Error::DataInvalid { + message: "Vector raw search found null vector element".to_string(), + source: None, + }); + } + } + Ok(()) +} + +#[derive(Clone, Copy)] +struct RawVectorRow<'a> { + row_id: u64, + values: &'a Float32Array, + start: usize, + end: usize, +} + +fn offer_raw_vector_score( + row: RawVectorRow<'_>, + query_index: usize, + metric: RawVectorMetric, + vector_searches: &[VectorSearch], + scoring_plan: &RawScoringPlan, + top_k_out: &mut [RawScoreTopK], +) -> crate::Result<()> { + let vector_search = &vector_searches[query_index]; + let stored_len = row.end - row.start; + ensure_raw_vector_dimension(stored_len, vector_search.vector.len())?; + let score = compute_raw_vector_score_from_values( + &vector_search.vector, + scoring_plan.query_l2_squared_norms[query_index], + row.values, + row.start, + row.end, + metric, + ); + top_k_out[query_index].offer(row.row_id, score); + Ok(()) +} + +fn compute_raw_vector_score_from_values( + query: &[f32], + query_l2_squared_norm: f32, + values: &Float32Array, + start: usize, + end: usize, + metric: RawVectorMetric, +) -> f32 { + debug_assert_eq!(query.len(), end - start); + match metric { + RawVectorMetric::L2 => compute_raw_vector_l2_score(query, &values.values()[start..end]), + RawVectorMetric::Cosine => { + let mut dot = 0.0; + let mut norm_b = 0.0; + for (q, value_index) in query.iter().zip(start..end) { + let stored = values.value(value_index); + dot += q * stored; + norm_b += stored * stored; + } + let denominator = query_l2_squared_norm.sqrt() * norm_b.sqrt(); + if denominator == 0.0 { + 0.0 + } else { + dot / denominator + } + } + RawVectorMetric::InnerProduct => query + .iter() + .zip(start..end) + .map(|(q, value_index)| q * values.value(value_index)) + .sum(), + } +} + +fn compute_raw_vector_l2_score(query: &[f32], stored: &[f32]) -> f32 { + let squared_distance = query + .iter() + .zip(stored) + .map(|(query_value, stored_value)| { + let difference = query_value - stored_value; + difference * difference + }) + .sum::(); + 1.0 / (1.0 + squared_distance) +} + +fn row_id_to_u64(row_id: i64) -> crate::Result { + u64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { + message: format!("Negative _ROW_ID {row_id} cannot be used for global index search"), + source: None, + }) +} + +#[cfg(test)] +fn compute_raw_vector_score(query: &[f32], stored: &[f32], metric: RawVectorMetric) -> f32 { + match metric { + RawVectorMetric::L2 => compute_raw_vector_l2_score(query, stored), + RawVectorMetric::Cosine => { + let mut dot = 0.0; + let mut norm_a = 0.0; + let mut norm_b = 0.0; + for (q, s) in query.iter().zip(stored.iter()) { + dot += q * s; + norm_a += q * q; + norm_b += s * s; + } + let denominator = norm_a.sqrt() * norm_b.sqrt(); + if denominator == 0.0 { + 0.0 + } else { + dot / denominator + } + } + RawVectorMetric::InnerProduct => query.iter().zip(stored.iter()).map(|(q, s)| q * s).sum(), + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs new file mode 100644 index 000000000..63e1ff801 --- /dev/null +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -0,0 +1,1392 @@ +// 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::*; +use crate::io::{FileIO, FileIOBuilder}; +use crate::lumina::{LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, LUMINA_IDENTIFIER}; +use crate::spec::{ + BinaryRow, CoreOptions, DataField, DataType, FileKind, GlobalIndexMeta, IndexFileMeta, + IndexManifestEntry, IntType, ROW_ID_FIELD_NAME, +}; +use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; +use crate::table::vector_search_common::VectorIndexBackend; +use crate::table::vector_search_common::{collect_ranked_rows, reorder_and_strip_position}; +use crate::table::vector_search_test_utils::{ + build_vindex_segment_bytes, de_vector_table, id_gt_filter, pk_vector_table, +}; +use crate::table::{find_field_id_by_name, RowRange}; +use crate::vector_search::{ScoredRowIds, VectorSearch}; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::IVF_FLAT_IDENTIFIER; +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; +use arrow_array::{Array, ArrayRef, Float32Array, Int32Array, Int64Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use futures::TryStreamExt; +use roaring::RoaringTreemap; +use std::collections::HashMap; +use std::sync::Arc; + +fn l2_score(distance: f32) -> f32 { + VectorSearchMetric::L2.distance_to_score(distance) +} + +fn make_field(id: i32, name: &str) -> DataField { + DataField::new(id, name.to_string(), DataType::Int(IntType::default())) +} + +fn eval_context<'a>( + file_io: &'a FileIO, + options: &'a HashMap, + fields: &'a [DataField], + next_row_id: Option, +) -> VectorSearchEvaluation<'a> { + VectorSearchEvaluation { + table: None, + file_io, + table_path: "memory:///test_table", + table_options: options, + schema_fields: fields, + next_row_id, + } +} + +fn make_lumina_entry( + file_name: &str, + index_type: &str, + kind: FileKind, + index_field_id: i32, +) -> IndexManifestEntry { + IndexManifestEntry { + kind, + partition: vec![], + bucket: 0, + index_file: IndexFileMeta { + index_type: index_type.to_string(), + file_name: file_name.to_string(), + file_size: 100, + row_count: 10, + deletion_vectors_ranges: None, + external_path: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: 0, + row_range_end: 9, + index_field_id, + extra_field_ids: None, + source_meta: None, + index_meta: None, + }), + }, + version: 1, + } +} + +// ---- Task B: search-and-read (`SearchResultReadBuilder::read`) tests ---- + +/// Build a small materialization batch: user column `id: Int32`, the internal +/// `_PKEY_VECTOR_POSITION: Int64`, and `__paimon_search_score: Float32` (mirroring +/// what `PkVectorIndexedSplitRead` emits for a single file). +fn materialized_batch(rows: &[(i32, i64, f32)]) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new(PKEY_VECTOR_POSITION_COLUMN, ArrowDataType::Int64, false), + ArrowField::new(SEARCH_SCORE_COLUMN, ArrowDataType::Float32, false), + ])); + let ids = Int32Array::from(rows.iter().map(|(id, _, _)| *id).collect::>()); + let positions = Int64Array::from(rows.iter().map(|(_, pos, _)| *pos).collect::>()); + let scores = Float32Array::from(rows.iter().map(|(_, _, s)| *s).collect::>()); + RecordBatch::try_new( + schema, + vec![Arc::new(ids), Arc::new(positions), Arc::new(scores)], + ) + .unwrap() +} + +fn i32_col(batch: &RecordBatch, name: &str) -> Vec { + let idx = batch.schema().index_of(name).unwrap(); + batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() +} + +fn f32_col(batch: &RecordBatch, name: &str) -> Vec { + let idx = batch.schema().index_of(name).unwrap(); + batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() +} + +#[test] +fn vindex_concurrency_limits_are_independent() { + let default_options = HashMap::new(); + let default_core = CoreOptions::new(&default_options); + assert_eq!( + vindex_concurrency_limits(&default_core, 1, 32).unwrap(), + (1, 64) + ); + assert_eq!( + vindex_concurrency_limits(&default_core, 8, 4).unwrap(), + (4, 64) + ); + + let options = HashMap::from([( + "global-index.vindex.read-thread-num".to_string(), + "48".to_string(), + )]); + let core = CoreOptions::new(&options); + assert_eq!(vindex_concurrency_limits(&core, 1, 32).unwrap(), (1, 48)); + assert_eq!(vindex_concurrency_limits(&core, 8, 4).unwrap(), (4, 48)); +} + +#[test] +fn test_find_field_id_by_name() { + let fields = vec![make_field(1, "id"), make_field(2, "embedding")]; + assert_eq!(find_field_id_by_name(&fields, "embedding"), Some(2)); + assert_eq!(find_field_id_by_name(&fields, "nonexistent"), None); +} + +#[test] +fn shared_include_filter_is_localized_once_per_index_shard() { + let include_row_ids = RoaringTreemap::from_iter([101, 205, 999]); + let localized = + localize_shared_include_row_ids(&include_row_ids, &[(100, 109), (200, 209), (300, 309)]) + .unwrap(); + + assert_eq!( + localized[0].as_ref().unwrap().iter().collect::>(), + vec![1] + ); + assert_eq!( + localized[1].as_ref().unwrap().iter().collect::>(), + vec![5] + ); + assert!(localized[2].is_none(), "an empty shard must be skipped"); +} + +#[test] +fn shared_batch_include_filter_requires_the_same_arc() { + let shared = Arc::new(RoaringTreemap::from_iter([1, 2, 3])); + let mut shared_searches = vec![ + VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(), + ]; + for search in &mut shared_searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + let detected = shared_batch_include_row_ids(&shared_searches).unwrap(); + assert!(Arc::ptr_eq(detected, &shared)); + + let mut equal_but_distinct = shared_searches.clone(); + equal_but_distinct[1] + .set_shared_include_row_ids(Arc::new(RoaringTreemap::from_iter([1, 2, 3]))); + assert!(shared_batch_include_row_ids(&equal_but_distinct).is_none()); + + let mut owned = shared_searches; + owned[1] = owned[1] + .clone() + .with_include_row_ids(RoaringTreemap::from_iter([1, 2, 3])); + assert!(shared_batch_include_row_ids(&owned).is_none()); +} + +#[test] +fn shared_raw_filter_does_not_expand_row_query_associations() { + let shared = Arc::new(RoaringTreemap::from_iter(0..1_000)); + let mut searches = (0..128) + .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap()) + .collect::>(); + for search in &mut searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + + let plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); + let expanded_associations = plan + .candidate_query_indices + .values() + .map(Vec::len) + .sum::(); + + assert_eq!( + expanded_associations, 0, + "one shared bitmap must stay O(B + Q), not expand to O(B * Q)" + ); + assert_eq!(plan.shared_filter_groups.len(), 1); + assert!(Arc::ptr_eq( + &plan.shared_filter_groups[0].include_row_ids, + &shared + )); + assert_eq!(plan.shared_filter_groups[0].query_indices.len(), 128); +} + +#[test] +fn shared_raw_filter_prunes_unindexed_ranges_before_reading() { + let shared = Arc::new(RoaringTreemap::from_iter([7, 1_000, 1_001, 900_000])); + let mut searches = (0..128) + .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap()) + .collect::>(); + for search in &mut searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + + let raw_ranges = vec![RowRange::new(0, 999_999)]; + assert_eq!( + prune_raw_ranges_by_include_row_ids(&raw_ranges, &searches).unwrap(), + vec![ + RowRange::new(7, 7), + RowRange::new(1_000, 1_001), + RowRange::new(900_000, 900_000), + ] + ); +} + +#[test] +fn test_raw_vector_score_matches_java_metric_semantics() { + let l2 = compute_raw_vector_score(&[1.0, 2.0], &[1.0, 4.0], RawVectorMetric::L2); + assert!((l2 - 0.2).abs() < 1e-6); + assert_eq!( + compute_raw_vector_score(&[1.0, 2.0], &[3.0, 4.0], RawVectorMetric::InnerProduct), + 11.0 + ); + let cosine = compute_raw_vector_score(&[1.0, 0.0], &[1.0, 1.0], RawVectorMetric::Cosine); + assert!((cosine - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-6); + assert_eq!( + compute_raw_vector_score(&[0.0, 0.0], &[1.0, 1.0], RawVectorMetric::Cosine), + 0.0 + ); +} + +#[test] +fn test_raw_vector_score_matrix_matches_scalar_metrics() { + let stored = vec![1.0, 2.0, 3.0, 4.0, 0.0, 0.0]; + let queries = vec![1.0, 1.0, -1.0, 2.0]; + let query_indices = vec![0, 1]; + let query_l2_squared_norms = vec![2.0, 5.0]; + + for metric in [ + RawVectorMetric::L2, + RawVectorMetric::Cosine, + RawVectorMetric::InnerProduct, + ] { + let matrix_scores = compute_raw_vector_score_matrix( + &stored, + 3, + &queries, + 2, + 2, + &query_l2_squared_norms, + &query_indices, + metric, + ) + .unwrap(); + for (row_index, stored_vector) in stored.as_chunks::<2>().0.iter().enumerate() { + for (query_index, query) in queries.as_chunks::<2>().0.iter().enumerate() { + let expected = compute_raw_vector_score(query, stored_vector, metric); + let actual = matrix_scores[query_index * 3 + row_index]; + assert!( + (actual - expected).abs() < 1e-5, + "metric={metric:?}, row={row_index}, query={query_index}: {actual} != {expected}" + ); + } + } + } + + let non_finite_score = compute_raw_vector_score_matrix( + &[f32::INFINITY, 0.0], + 1, + &[1.0, 0.0], + 1, + 2, + &[1.0], + &[0], + RawVectorMetric::L2, + ) + .unwrap()[0]; + assert_eq!(non_finite_score, 0.0); +} + +#[test] +fn test_raw_vector_score_matrix_l2_preserves_large_finite_distances() { + let dimension = 128; + let query = vec![1.0e10_f32; dimension]; + let mut nearby = query.clone(); + nearby[0] += 1024.0; + let mut stored = query.clone(); + stored.extend_from_slice(&nearby); + let queries = query.repeat(4); + let query_l2_squared_norm = query.iter().map(|value| value * value).sum::(); + let query_l2_squared_norms = vec![query_l2_squared_norm; 4]; + let query_indices = vec![0, 1, 2, 3]; + + let matrix_scores = compute_raw_vector_score_matrix( + &stored, + 2, + &queries, + 4, + dimension, + &query_l2_squared_norms, + &query_indices, + RawVectorMetric::L2, + ) + .unwrap(); + let exact_score = compute_raw_vector_score(&query, &query, RawVectorMetric::L2); + let nearby_score = compute_raw_vector_score(&query, &nearby, RawVectorMetric::L2); + + for query_index in 0..4 { + assert_eq!(matrix_scores[query_index * 2], exact_score); + assert_eq!(matrix_scores[query_index * 2 + 1], nearby_score); + assert!(matrix_scores[query_index * 2] > matrix_scores[query_index * 2 + 1]); + } +} + +#[test] +fn test_raw_vector_cosine_avoids_squared_norm_product_overflow() { + let query = vec![1.0e15_f32, 0.0]; + let query_l2_squared_norm = query.iter().map(|value| value * value).sum::(); + assert!(query_l2_squared_norm.is_finite()); + let values = Float32Array::from(query.clone()); + let scalar_score = compute_raw_vector_score_from_values( + &query, + query_l2_squared_norm, + &values, + 0, + 2, + RawVectorMetric::Cosine, + ); + assert!((scalar_score - 1.0).abs() < 1e-6); + + let queries = query.repeat(4); + let matrix_scores = compute_raw_vector_score_matrix( + &query, + 1, + &queries, + 4, + 2, + &[query_l2_squared_norm; 4], + &[0, 1, 2, 3], + RawVectorMetric::Cosine, + ) + .unwrap(); + assert!(matrix_scores + .iter() + .all(|score| (*score - 1.0).abs() < 1e-6)); +} + +#[test] +fn test_raw_score_top_k_matches_full_sort_with_linear_partial_selection() { + let limit = 7; + let mut top_k = RawScoreTopK::new(limit); + let mut batched_top_k = RawScoreTopK::new(limit); + let mut expected = Vec::new(); + for row_id in 0..10_000 { + let score = ((row_id * 37) % 101) as f32 / 10.0; + let candidate = RawScoredRow { row_id, score }; + expected.push(candidate); + top_k.offer(row_id, score); + batched_top_k.offer_many(std::iter::once(candidate)); + assert!(top_k.candidates.len() < top_k.partition_size()); + assert!(batched_top_k.candidates.len() < batched_top_k.partition_size()); + } + expected.sort_unstable_by(RawScoredRow::strongest_first); + expected.truncate(limit); + + let result = top_k.into_search_result(); + let batched_result = batched_top_k.into_search_result(); + assert_eq!( + result.row_ids, + expected.iter().map(|row| row.row_id).collect::>() + ); + assert_eq!( + result.scores, + expected.iter().map(|row| row.score).collect::>() + ); + assert_eq!(batched_result.row_ids, result.row_ids); + assert_eq!(batched_result.scores, result.scores); +} + +#[test] +fn test_configured_raw_vector_metric_precedence_and_conflict_default() { + let mut options = HashMap::new(); + options.insert( + "fields.embedding.distance.metric".to_string(), + "inner-product".to_string(), + ); + options.insert("metric".to_string(), "cosine".to_string()); + assert_eq!( + configured_raw_vector_metric(&options, "embedding").unwrap(), + RawVectorMetric::InnerProduct + ); + + options.clear(); + options.insert("foo.metric".to_string(), "cosine".to_string()); + options.insert("bar.distance.metric".to_string(), "l2".to_string()); + assert_eq!( + configured_raw_vector_metric(&options, "embedding").unwrap(), + RawVectorMetric::L2 + ); +} + +#[tokio::test] +async fn test_resolve_raw_vector_metric_uses_vindex_manifest_metadata() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let mut entry = make_lumina_entry("missing.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + let index_meta = serde_json::to_vec(&HashMap::from([( + "metric".to_string(), + "cosine".to_string(), + )])) + .unwrap(); + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .index_meta = Some(index_meta); + + let metric = resolve_raw_vector_metric( + &file_io, + "memory:///test_table", + &HashMap::new(), + &[entry], + 2, + "embedding", + ) + .await + .unwrap(); + + assert_eq!(metric, RawVectorMetric::Cosine); +} + +#[tokio::test] +async fn test_resolve_raw_vector_metric_falls_back_to_vindex_header() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let index = build_vindex_segment_bytes("inner_product"); + file_io + .new_output("memory:///test_table/index/test.idx") + .unwrap() + .write(bytes::Bytes::from(index.clone())) + .await + .unwrap(); + for (file_size, index_meta) in [ + (index.len() as i64, br#"{"metric":"euclidean"}"#.to_vec()), + (0, b"{}".to_vec()), + (-1, b"{}".to_vec()), + ] { + let mut entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + entry.index_file.file_size = file_size; + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .index_meta = Some(index_meta); + + let metric = resolve_raw_vector_metric( + &file_io, + "memory:///test_table", + &HashMap::new(), + &[entry], + 2, + "embedding", + ) + .await + .unwrap(); + + assert_eq!(metric, RawVectorMetric::InnerProduct); + } +} + +#[test] +fn test_configured_refine_factor_precedence_and_aliases() { + let table_options = HashMap::from([( + "fields.embedding.ivf.refine-factor".to_string(), + "3".to_string(), + )]); + let search_options = HashMap::from([( + "fields.embedding.ivf_flat.rerank_factor".to_string(), + "2".to_string(), + )]); + assert_eq!( + configured_refine_factor( + &search_options, + &table_options, + "embedding", + IVF_FLAT_IDENTIFIER, + ) + .unwrap(), + 2 + ); + + assert_eq!( + configured_refine_factor( + &HashMap::new(), + &table_options, + "embedding", + IVF_FLAT_IDENTIFIER, + ) + .unwrap(), + 3 + ); + + let global_options = HashMap::from([("rerank-factor".to_string(), "4".to_string())]); + assert_eq!( + configured_refine_factor( + &HashMap::new(), + &global_options, + "embedding", + LUMINA_IDENTIFIER, + ) + .unwrap(), + 4 + ); +} + +#[test] +fn test_configured_refine_factor_rejects_invalid_values() { + let zero_options = HashMap::from([("refine_factor".to_string(), "0".to_string())]); + let err = configured_refine_factor( + &zero_options, + &HashMap::new(), + "embedding", + LUMINA_IDENTIFIER, + ) + .unwrap_err(); + assert!(err.to_string().contains("must be positive")); + + let invalid_options = HashMap::from([("refine_factor".to_string(), "abc".to_string())]); + let err = configured_refine_factor( + &invalid_options, + &HashMap::new(), + "embedding", + LUMINA_IDENTIFIER, + ) + .unwrap_err(); + assert!(err.to_string().contains("Must be an integer")); + + assert!(indexed_search_limit(i32::MAX as usize, 2).is_err()); +} + +#[test] +fn test_collect_raw_batch_vector_batch_preserves_query_order() { + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); + for vector in [[1.0, 0.0], [0.0, 1.0], [0.8, 0.2]] { + builder.values().append_value(vector[0]); + builder.values().append_value(vector[1]); + builder.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + 2, + ), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(builder.finish()) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(10), Some(11), Some(12)])) as ArrayRef, + ], + ) + .unwrap(); + let searches = vec![ + VectorSearch::new(vec![1.0, 0.0], 1, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.0, 1.0], 1, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.8, 0.2], 1, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.5, 0.5], 1, "embedding".to_string()).unwrap(), + ]; + let scoring_plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); + let mut top_k = searches + .iter() + .map(|search| RawScoreTopK::new(search.limit)) + .collect::>(); + + collect_raw_batch_vector_batch( + &batch, + &searches, + RawVectorMetric::L2, + &scoring_plan, + &mut top_k, + ) + .unwrap(); + let results = top_k + .into_iter() + .map(RawScoreTopK::into_search_result) + .collect::>(); + + assert_eq!(results[0].row_ids, vec![10]); + assert_eq!(results[1].row_ids, vec![11]); + assert_eq!(results[2].row_ids, vec![12]); + assert_eq!(results[3].row_ids, vec![12]); +} + +#[test] +fn test_collect_raw_batch_vector_batch_respects_fixed_size_list_offset() { + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); + for vector in [[1.0, 0.0], [0.0, 1.0], [0.8, 0.2]] { + builder.values().append_value(vector[0]); + builder.values().append_value(vector[1]); + builder.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + 2, + ), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(builder.finish()) as ArrayRef, + Arc::new(Int64Array::from(vec![10, 11, 12])) as ArrayRef, + ], + ) + .unwrap() + .slice(1, 2); + let searches = vec![VectorSearch::new(vec![0.0, 1.0], 1, "embedding".to_string()).unwrap()]; + let scoring_plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); + let mut top_k = vec![RawScoreTopK::new(1)]; + + collect_raw_batch_vector_batch( + &batch, + &searches, + RawVectorMetric::L2, + &scoring_plan, + &mut top_k, + ) + .unwrap(); + + assert_eq!(top_k.pop().unwrap().into_search_result().row_ids, vec![11]); +} + +#[test] +fn test_collect_raw_batch_vector_batch_scores_only_include_row_ids() { + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); + for vector in [[1.0, 0.0], [0.0, 1.0], [0.8, 0.2]] { + builder.values().append_value(vector[0]); + builder.values().append_value(vector[1]); + builder.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + 2, + ), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(builder.finish()) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(10), Some(11), Some(12)])) as ArrayRef, + ], + ) + .unwrap(); + let mut include_row_ids = RoaringTreemap::new(); + include_row_ids.insert(12); + let searches = vec![ + VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()) + .unwrap() + .with_include_row_ids(include_row_ids), + ]; + let scoring_plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); + let mut top_k = searches + .iter() + .map(|search| RawScoreTopK::new(search.limit)) + .collect::>(); + + collect_raw_batch_vector_batch( + &batch, + &searches, + RawVectorMetric::L2, + &scoring_plan, + &mut top_k, + ) + .unwrap(); + let results = top_k + .into_iter() + .map(RawScoreTopK::into_search_result) + .collect::>(); + + assert_eq!(results[0].row_ids, vec![12]); + assert_eq!(results[0].scores.len(), 1); +} + +#[tokio::test] +async fn test_batch_evaluate_no_matching_field_returns_empty_per_query() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let fields = vec![make_field(1, "id")]; + let searches = vec![ + VectorSearch::new(vec![1.0], 10, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.0], 10, "embedding".to_string()).unwrap(), + ]; + let options = HashMap::new(); + + let entry = make_lumina_entry( + "test.idx", + LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, + FileKind::Add, + 99, + ); + + let results = evaluate_batch_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &searches, + ) + .await + .unwrap(); + + assert_eq!(results.len(), searches.len()); + assert!(results.iter().all(ScoredRowIds::is_empty)); +} + +#[tokio::test] +async fn test_evaluate_no_matching_entries() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let fields = vec![make_field(1, "id"), make_field(2, "embedding")]; + let vs = VectorSearch::new(vec![1.0, 2.0], 10, "embedding".to_string()).unwrap(); + let options = HashMap::new(); + + let entry = IndexManifestEntry { + kind: FileKind::Add, + partition: vec![], + bucket: 0, + index_file: IndexFileMeta { + index_type: "btree".to_string(), + file_name: "test.idx".to_string(), + file_size: 100, + row_count: 10, + deletion_vectors_ranges: None, + external_path: None, + global_index_meta: None, + }, + version: 1, + }; + + let result = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap(); + assert!(result.is_empty()); +} + +#[tokio::test] +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(); + let options = HashMap::new(); + + let entry = make_lumina_entry("test.idx", "btree", FileKind::Add, 2); + + let result = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap(); + assert!(result.is_empty()); +} + +#[tokio::test] +async fn test_evaluate_full_mode_without_vector_entries_uses_raw_path() { + 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 options = HashMap::from([("vector-index.search-mode".to_string(), "full".to_string())]); + + let err = evaluate_vector_search( + eval_context(&file_io, &options, &fields, Some(10)), + &[], + &vs, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Vector raw search requires table context"), + "unexpected error: {err}" + ); +} + +#[tokio::test] +async fn test_evaluate_no_matching_field() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let fields = vec![make_field(1, "id")]; + let vs = VectorSearch::new(vec![1.0], 10, "embedding".to_string()).unwrap(); + let options = HashMap::new(); + + let entry = make_lumina_entry( + "test.idx", + LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, + FileKind::Add, + 99, + ); + + let result = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap(); + assert!(result.is_empty()); +} + +#[tokio::test] +async fn test_evaluate_skips_delete_entries() { + 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 options = HashMap::new(); + + let entry = make_lumina_entry( + "test.idx", + LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, + FileKind::Delete, + 2, + ); + + let result = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap(); + assert!(result.is_empty()); +} + +#[tokio::test] +async fn test_evaluate_accepts_canonical_lumina_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 options = HashMap::new(); + + let entry = make_lumina_entry("missing.idx", LUMINA_IDENTIFIER, FileKind::Add, 2); + + let err = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Failed to read Lumina index file 'missing.idx'"), + "unexpected error: {err}" + ); +} + +#[tokio::test] +async fn test_evaluate_accepts_legacy_lumina_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 options = HashMap::new(); + + let entry = make_lumina_entry( + "missing.idx", + LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, + FileKind::Add, + 2, + ); + + let err = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Failed to read Lumina index file 'missing.idx'"), + "unexpected error: {err}" + ); +} + +#[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 options = HashMap::new(); + + let entry = make_lumina_entry("missing.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + + let err = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Failed to read vindex index file 'missing.idx'"), + "unexpected error: {err}" + ); + assert!( + std::error::Error::source(&err).is_some(), + "wrapped vindex read errors should retain their source: {err:?}" + ); +} + +#[test] +fn test_single_vindex_outside_tokio_returns_error() { + futures::executor::block_on(async { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + file_io + .new_output("memory:///test_table/index/test.idx") + .unwrap() + .write(bytes::Bytes::from_static(b"index")) + .await + .unwrap(); + let fields = vec![make_field(2, "embedding")]; + let vs = VectorSearch::new(vec![1.0], 10, "embedding".to_string()).unwrap(); + let options = HashMap::new(); + let entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + + let err = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &vs, + ) + .await + .expect_err("vindex range reads outside Tokio should fail without panicking"); + + assert!( + matches!(err, crate::Error::UnexpectedError { ref message, .. } + if message.contains("requires a Tokio runtime")), + "unexpected error: {err:?}" + ); + }); +} + +#[test] +fn test_batch_vindex_outside_tokio_uses_buffered_fallback() { + futures::executor::block_on(async { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let index = build_vindex_segment_bytes("l2"); + file_io + .new_output("memory:///test_table/index/test.idx") + .unwrap() + .write(bytes::Bytes::from(index.clone())) + .await + .unwrap(); + let fields = vec![make_field(2, "embedding")]; + let searches = vec![ + VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(), + ]; + let options = HashMap::new(); + let mut entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + entry.index_file.file_size = index.len() as i64; + entry.index_file.row_count = 3; + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .row_range_end = 2; + + let results = evaluate_batch_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &searches, + ) + .await + .expect("batch vindex search should fall back to buffered I/O outside Tokio"); + + assert_eq!(results.len(), searches.len()); + assert!(results.iter().all(|result| !result.is_empty())); + }); +} + +#[test] +fn from_index_type_classifies_lumina_and_vindex() { + assert_eq!( + VectorIndexBackend::from_index_type("lumina"), + Some(VectorIndexBackend::Lumina) + ); + assert_eq!( + VectorIndexBackend::from_index_type("lumina-vector-ann"), + Some(VectorIndexBackend::Lumina) + ); + assert_eq!( + VectorIndexBackend::from_index_type("ivf-flat"), + Some(VectorIndexBackend::Vindex) + ); + for index_type in ["ivf-sq", "ivf-rq", "diskann"] { + assert_eq!( + VectorIndexBackend::from_index_type(index_type), + Some(VectorIndexBackend::Vindex) + ); + } +} + +#[tokio::test] +async fn execute_filter_on_empty_de_path_returns_empty() { + // No PK-vector index and no snapshot: the request follows the + // data-evolution path. Scalar pre-filter support must not turn an empty + // table into an error. + let table = pk_vector_table(&[]); + let filter = id_gt_filter(&table, 2); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .with_filter(filter) + .execute() + .await + .expect("an empty data-evolution search with a filter should succeed"); + assert!(result.is_empty()); +} + +#[test] +fn reorder_and_strip_position_recovers_best_first_and_drops_position() { + // Single file, one bucket. The materialization reader emits rows in + // ascending physical position [pos0, pos1, pos2] -> ids [40,41,42]. The + // search candidates ranked them best-first as pos1(rank0), pos2(rank1), + // pos0(rank2), which is NEITHER position order nor score order-by-batch. + // The reorder must yield ids [41,42,40] and drop _PKEY_VECTOR_POSITION. + let batch = materialized_batch(&[ + (40, 0, l2_score(9.0)), + (41, 1, l2_score(1.0)), + (42, 2, l2_score(4.0)), + ]); + let batches = vec![batch]; + let part = BinaryRow::new(0).to_serialized_bytes(); + let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 1), 0); + rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 2), 1); + rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 0), 2); + + let mut ranked = Vec::new(); + collect_ranked_rows(&batches[0], 0, &part, 0, "o.mosaic", &rank_of, &mut ranked).unwrap(); + let out = reorder_and_strip_position(&batches, ranked).unwrap(); + assert_eq!(out.len(), 1); + let out = &out[0]; + + // Best-first row order, not ascending position order. + assert_eq!(i32_col(out, "id"), vec![41, 42, 40]); + // Score column preserved and aligned to the reordered rows. + assert_eq!( + f32_col(out, SEARCH_SCORE_COLUMN), + vec![l2_score(1.0), l2_score(4.0), l2_score(9.0)] + ); + // Position column dropped; _ROW_ID never present. + assert!(out.schema().index_of(PKEY_VECTOR_POSITION_COLUMN).is_err()); + assert!(out.schema().index_of("_ROW_ID").is_err()); +} + +#[test] +fn reorder_and_strip_position_merges_rows_across_files() { + // Two files (two materialization batches). Best-first interleaves them: + // file-b pos0 (rank0), file-a pos1 (rank1), file-a pos0 (rank2). The + // reorder must pull rows from both batches into one best-first output. + let batch_a = materialized_batch(&[(10, 0, l2_score(9.0)), (11, 1, l2_score(1.0))]); + let batch_b = materialized_batch(&[(20, 0, l2_score(0.5))]); + let batches = vec![batch_a, batch_b]; + let part = BinaryRow::new(0).to_serialized_bytes(); + let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + rank_of.insert((part.clone(), 0, "b".to_string(), 0), 0); + rank_of.insert((part.clone(), 0, "a".to_string(), 1), 1); + rank_of.insert((part.clone(), 0, "a".to_string(), 0), 2); + + let mut ranked = Vec::new(); + collect_ranked_rows(&batches[0], 0, &part, 0, "a", &rank_of, &mut ranked).unwrap(); + collect_ranked_rows(&batches[1], 1, &part, 0, "b", &rank_of, &mut ranked).unwrap(); + let out = reorder_and_strip_position(&batches, ranked).unwrap(); + assert_eq!(i32_col(&out[0], "id"), vec![20, 11, 10]); + assert_eq!( + f32_col(&out[0], SEARCH_SCORE_COLUMN), + vec![l2_score(0.5), l2_score(1.0), l2_score(9.0)] + ); +} + +#[test] +fn reorder_and_strip_position_empty_yields_no_batches() { + let out = reorder_and_strip_position(&[], Vec::new()).unwrap(); + assert!(out.is_empty()); +} + +#[test] +fn collect_ranked_rows_missing_candidate_fails_loud() { + // A materialized position with no candidate rank must fail loud rather than + // silently drop the row. + let batch = materialized_batch(&[(40, 7, l2_score(1.0))]); + let part = BinaryRow::new(0).to_serialized_bytes(); + let rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + let mut ranked = Vec::new(); + let err = collect_ranked_rows(&batch, 0, &part, 0, "f", &rank_of, &mut ranked) + .expect_err("missing candidate must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("no matching search candidate")), + "unexpected error: {err:?}" + ); +} + +#[test] +fn attach_scores_reorders_by_rank_not_score() { + use arrow_array::{Int32Array, Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + // Two rows materialized in row-id order [10, 20]; ranks say 20 is best (rank 0), + // 10 is rank 1. Scores tie at 0.5 to prove ordering follows rank, not score. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new(ROW_ID_FIELD_NAME, DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![100, 200])), + Arc::new(Int64Array::from(vec![10, 20])), + ], + ) + .unwrap(); + let mut map = HashMap::new(); + map.insert(20i64, (0usize, 0.5f32)); + map.insert(10i64, (1usize, 0.5f32)); + + let out = attach_scores_by_row_id(&[batch], &map, 2).unwrap(); + assert_eq!(out.len(), 1); + let b = &out[0]; + // _ROW_ID stripped, score appended. + assert!(b.schema().index_of(ROW_ID_FIELD_NAME).is_err()); + let score_idx = b.schema().index_of("__paimon_search_score").unwrap(); + assert_eq!( + b.schema().field(score_idx).data_type(), + &arrow_schema::DataType::Float32 + ); + // Row order is rank order: id 200 (rank 0) first, then id 100 (rank 1). + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + assert_eq!(ids.values(), &[200, 100]); +} + +#[test] +fn attach_scores_fails_on_unknown_row_id() { + use arrow_array::{Int32Array, Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new(ROW_ID_FIELD_NAME, DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int64Array::from(vec![99])), + ], + ) + .unwrap(); + let map: HashMap = HashMap::new(); // no entry for 99 + let err = attach_scores_by_row_id(&[batch], &map, 1).unwrap_err(); + assert!(matches!(err, crate::Error::DataInvalid { .. })); +} + +#[test] +fn attach_scores_fails_on_count_mismatch() { + use arrow_array::{Int32Array, Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new(ROW_ID_FIELD_NAME, DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int64Array::from(vec![10])), + ], + ) + .unwrap(); + let mut map = HashMap::new(); + map.insert(10i64, (0usize, 0.5f32)); + // expected_len 2 but only 1 row materialized. + let err = attach_scores_by_row_id(&[batch], &map, 2).unwrap_err(); + assert!(matches!(err, crate::Error::DataInvalid { .. })); +} + +#[test] +fn attach_scores_fails_on_null_row_id() { + use arrow_array::{Int32Array, Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + // _ROW_ID column has a NULL at row 1; the map contains the non-null id, so + // the failure is specifically the null (not an unknown id). + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new(ROW_ID_FIELD_NAME, DataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int64Array::from(vec![Some(10i64), None])), + ], + ) + .unwrap(); + let mut map = HashMap::new(); + map.insert(10i64, (0usize, 0.5f32)); + let err = attach_scores_by_row_id(&[batch], &map, 2).unwrap_err(); + assert!(matches!(err, crate::Error::DataInvalid { .. })); +} + +#[test] +fn attach_scores_fails_on_wrong_type_row_id() { + use arrow_array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + // _ROW_ID column is Int32, not Int64: the downcast fails loud. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new(ROW_ID_FIELD_NAME, DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![10])), + ], + ) + .unwrap(); + let mut map = HashMap::new(); + map.insert(10i64, (0usize, 0.5f32)); + let err = attach_scores_by_row_id(&[batch], &map, 1).unwrap_err(); + assert!(matches!(err, crate::Error::DataInvalid { .. })); +} + +#[tokio::test] +async fn de_result_read_materializes_rows_with_score() { + // A data-evolution vector table with a committed global index: result_read + // must materialize one row per scored hit and carry the unified score + // column, in best-first rank order. + let table = de_vector_table().await; + let query = vec![1.0, 0.0]; + + let scored = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(query.clone()) + .with_limit(3) + .execute() + .await + .unwrap(); + assert!(!scored.is_empty(), "DE search must return hits"); + + let mut stream = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(query) + .with_limit(3) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .unwrap(); + + let mut rows = 0usize; + let mut saw_score = false; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows(); + saw_score |= batch.schema().index_of(SEARCH_SCORE_COLUMN).is_ok(); + } + assert_eq!( + rows, + scored.len(), + "DE read must emit exactly the scored result count" + ); + assert!( + saw_score, + "DE read output must carry the search score column" + ); +} + +#[tokio::test] +async fn de_result_read_applies_scalar_filter_before_top_k() { + // Row id=1 is the closest vector to [1, 0], but the scalar filter excludes + // it. Filter-before-Top-K must return the best rows among ids > 1 instead + // of recalling id=1 first and filtering it after the search. + let table = de_vector_table().await; + let filter = id_gt_filter(&table, 1); + let mut stream = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(2) + .with_filter(filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .expect("DE vector search should support a scalar pre-filter"); + + let mut ids = Vec::new(); + while let Some(batch) = stream.try_next().await.unwrap() { + let id = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..id.len()).map(|row| id.value(row))); + } + + assert_eq!(ids, vec![3, 2]); +} diff --git a/crates/paimon/src/table/de_vector_scan.rs b/crates/paimon/src/table/de_vector_scan.rs new file mode 100644 index 000000000..c5830a65b --- /dev/null +++ b/crates/paimon/src/table/de_vector_scan.rs @@ -0,0 +1,293 @@ +// 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. + +//! Plans global-index vector searches against one snapshot and resolves scalar pre-filters. + +use crate::spec::{CoreOptions, IndexManifest, IndexManifestEntry, Predicate, ROW_ID_FIELD_NAME}; +use crate::table::vector_scan::Scan; +use crate::table::Table; +use crate::vindex::vector_search_timing_enabled; +use arrow_array::{Array, Int64Array}; +use futures::TryStreamExt; +use roaring::RoaringTreemap; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// A scalar vector pre-filter resolved once against one pinned snapshot. +/// +/// Reusing this value avoids repeating the same scalar-index/table read for +/// every input batch of a lateral vector query. +#[derive(Debug, Clone)] +pub struct PreparedVectorSearchFilter { + table: Table, + include_row_ids: Arc, +} + +impl PreparedVectorSearchFilter { + pub fn table(&self) -> &Table { + &self.table + } + + pub fn include_row_ids(&self) -> &Arc { + &self.include_row_ids + } +} + +fn same_vector_search_table(left: &Table, right: &Table) -> bool { + left.location().trim_end_matches('/') == right.location().trim_end_matches('/') + && left.branch() == right.branch() +} + +async fn matching_row_ids_for_filter( + table: &Table, + filter: &Predicate, +) -> crate::Result { + let mut read_builder = table.new_read_builder(); + read_builder + .with_projection(&[ROW_ID_FIELD_NAME])? + .with_filter(filter.clone()); + let plan = read_builder.new_scan().plan().await?; + let read = read_builder.new_read()?; + let mut stream = read.to_arrow(plan.splits())?; + let mut row_ids = RoaringTreemap::new(); + while let Some(batch) = stream.try_next().await? { + let index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|_| crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter read is missing {ROW_ID_FIELD_NAME}" + ), + source: None, + })?; + let values = batch + .column(index) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter {ROW_ID_FIELD_NAME} column is not Int64" + ), + source: None, + })?; + for row in 0..values.len() { + if values.is_null(row) { + return Err(crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter produced a null {ROW_ID_FIELD_NAME}" + ), + source: None, + }); + } + let row_id = values.value(row); + let row_id = u64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { + message: format!( + "scalar vector pre-filter produced a negative {ROW_ID_FIELD_NAME}: {row_id}" + ), + source: None, + })?; + row_ids.insert(row_id); + } + } + Ok(row_ids) +} + +impl Table { + /// Resolve a scalar predicate once and pin all later vector-search/read + /// stages to the same snapshot. + pub async fn prepare_vector_search_filter( + &self, + filter: Predicate, + ) -> crate::Result { + CoreOptions::new(self.schema().options()).ensure_read_authorized()?; + let Some(snapshot) = crate::table::time_travel::resolve_snapshot(self).await? else { + return Ok(PreparedVectorSearchFilter { + table: self.clone(), + include_row_ids: Arc::new(RoaringTreemap::new()), + }); + }; + let table = self.copy_with_resolved_snapshot(&snapshot).await?; + let include_row_ids = matching_row_ids_for_filter(&table, &filter).await?; + Ok(PreparedVectorSearchFilter { + table, + include_row_ids: Arc::new(include_row_ids), + }) + } +} + +pub(super) struct DeVectorScan { + table: Table, + filter: Option, + include_row_ids: Option>, + prepared_filter: Option, +} + +impl DeVectorScan { + pub(super) fn new( + table: &Table, + filter: Option<&Predicate>, + include_row_ids: Option<&Arc>, + prepared_filter: Option<&PreparedVectorSearchFilter>, + ) -> Self { + Self { + table: table.clone(), + filter: filter.cloned(), + include_row_ids: include_row_ids.cloned(), + prepared_filter: prepared_filter.cloned(), + } + } +} + +/// One pinned snapshot, its index manifest and resolved row-ID filter. +/// Query vectors and search options belong to the reader, not to this plan. +#[derive(Clone)] +pub(super) struct DeVectorScanPlan { + pub(super) table: Table, + pub(super) index_entries: Vec, + pub(super) include_row_ids: Option>, + pub(super) timing: Option, + pub(super) next_row_id: Option, + pub(super) skip_search: bool, +} + +impl DeVectorScanPlan { + fn empty(table: Table, timing: Option) -> Self { + Self { + table, + skip_search: true, + index_entries: Vec::new(), + include_row_ids: None, + next_row_id: None, + timing, + } + } +} + +#[derive(Clone)] +pub(super) struct DeVectorScanTiming { + pub(super) total_start: Instant, + pub(super) setup: Duration, + pub(super) snapshot: Duration, + pub(super) manifest: Duration, +} + +impl Scan for DeVectorScan { + type Plan = DeVectorScanPlan; + async fn plan(&self) -> crate::Result { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); + // The builder target is authoritative for current auth/type policy. + // A prepared filter only pins a snapshot and may carry older options. + CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + if let Some(prepared) = self.prepared_filter.as_ref() { + if !same_vector_search_table(&self.table, prepared.table()) { + return Err(crate::Error::DataInvalid { + message: format!( + "Prepared vector search filter belongs to a different table: builder target is '{}@{}', prepared filter target is '{}@{}'", + self.table.location(), + self.table.branch(), + prepared.table().location(), + prepared.table().branch(), + ), + source: None, + }); + } + } + // Check the pinned execution view as defense in depth before any fast + // path returns data-derived row ids/scores outside TableScan/TableRead. + let execution_table = self + .prepared_filter + .as_ref() + .map(PreparedVectorSearchFilter::table) + .unwrap_or(&self.table); + let core = CoreOptions::new(execution_table.schema().options()); + core.ensure_read_authorized()?; + let mut plan = DeVectorScanPlan::empty( + execution_table.clone(), + total_start.map(|total_start| DeVectorScanTiming { + total_start, + setup: total_start.elapsed(), + snapshot: Duration::ZERO, + manifest: Duration::ZERO, + }), + ); + + if self + .prepared_filter + .as_ref() + .is_some_and(|prepared| prepared.include_row_ids().is_empty()) + { + return Ok(plan); + } + + let snapshot_manager = execution_table.snapshot_manager(); + let snapshot_start = timing_enabled.then(Instant::now); + let snapshot = crate::table::time_travel::resolve_snapshot(execution_table).await?; + if let Some(timing) = &mut plan.timing { + timing.snapshot = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); + } + let Some(snapshot) = snapshot else { + return Ok(plan); + }; + plan.table = match self.prepared_filter.as_ref() { + Some(prepared) => prepared.table().clone(), + None => { + execution_table + .copy_with_resolved_snapshot(&snapshot) + .await? + } + }; + plan.next_row_id = snapshot.next_row_id(); + + plan.include_row_ids = if let Some(prepared) = self.prepared_filter.as_ref() { + Some(Arc::clone(prepared.include_row_ids())) + } else if let Some(include_row_ids) = self.include_row_ids.as_ref() { + Some(Arc::clone(include_row_ids)) + } else if let Some(filter) = self.filter.as_ref() { + Some(Arc::new( + matching_row_ids_for_filter(&plan.table, filter).await?, + )) + } else { + None + }; + if plan + .include_row_ids + .as_ref() + .is_some_and(|ids| ids.is_empty()) + { + return Ok(plan); + } + + let manifest_start = timing_enabled.then(Instant::now); + plan.index_entries = match snapshot.index_manifest() { + Some(index_manifest_name) => { + let manifest_path = snapshot_manager.manifest_path(index_manifest_name); + IndexManifest::read(execution_table.file_io(), &manifest_path).await? + } + None => Vec::new(), + }; + if let Some(timing) = &mut plan.timing { + timing.manifest = manifest_start.map_or(Duration::ZERO, |start| start.elapsed()); + } + plan.skip_search = false; + Ok(plan) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/de_vector_scan/tests.rs b/crates/paimon/src/table/de_vector_scan/tests.rs new file mode 100644 index 000000000..2016e6974 --- /dev/null +++ b/crates/paimon/src/table/de_vector_scan/tests.rs @@ -0,0 +1,364 @@ +// 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::*; +use crate::spec::CoreOptions; +use crate::table::de_vector_read::DeVectorRead; +use crate::table::vector_read::Read; +use crate::table::vector_search_test_utils::{ + de_vector_table, id_gt_filter, vector_test_table, vector_test_table_at, +}; +use crate::table::{TableCommit, TableWrite}; +use crate::vector_search::SearchResult; +use arrow_array::builder::{Float32Builder, ListBuilder}; +use arrow_array::{ArrayRef, Float32Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use roaring::RoaringTreemap; +use std::collections::HashMap; +use std::sync::Arc; + +#[tokio::test] +async fn reader_uses_planned_snapshot_after_snapshot_files_are_removed() { + let table = de_vector_table().await; + let options = HashMap::new(); + let scan = DeVectorScan::new(&table, None, None, None); + let read = DeVectorRead::new("embedding", &[&[1.0, 0.0]], 2, &options).unwrap(); + let plan = scan.plan().await.unwrap(); + let snapshot_id = plan.table.travel_snapshot().unwrap().id(); + let manager = table.snapshot_manager(); + for id in manager.list_all_ids().await.unwrap() { + table + .file_io() + .delete_file(&manager.snapshot_path(id)) + .await + .unwrap(); + } + assert!(manager.get_snapshot(snapshot_id).await.is_err()); + + // Re-resolving "latest" here would lose the scan's index and row-ID context. + let result = read.read(plan).await.unwrap(); + let result = result.into_iter().next().unwrap(); + assert_eq!(result.snapshot_id(), Some(snapshot_id)); + let batches: Vec = result + .new_read_builder() + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.len(), 1); + let ids = batches[0] + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = batches[0] + .column_by_name("__paimon_search_score") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values().as_ref(), &[1, 3]); + assert_eq!(scores.values().as_ref(), &[1.0, 1.0]); +} + +#[tokio::test] +async fn reader_reuses_queries_without_leaking_plan_filters() { + let table = de_vector_table().await; + let read = + DeVectorRead::new("embedding", &[&[1.0, 0.0], &[0.0, 1.0]], 2, &HashMap::new()).unwrap(); + let filter = id_gt_filter(&table, 1); + let empty_filter = id_gt_filter(&table, 99); + // One reader executes filtered, empty, and unfiltered plans. The plan owns + // the allow-list; neither it nor the empty result can alter later queries. + for (filter, expected) in [ + (Some(&filter), [vec![2, 1], vec![1, 2]]), + (Some(&empty_filter), [vec![], vec![]]), + (None, [vec![0, 2], vec![1, 2]]), + ] { + let plan = DeVectorScan::new(&table, filter, None, None) + .plan() + .await + .unwrap(); + let snapshot_id = plan.table.travel_snapshot().unwrap().id(); + let results = read.read(plan).await.unwrap(); + assert_eq!(results.len(), 2, "empty plans must preserve query count"); + for (result, row_ids) in results.iter().zip(expected) { + assert_eq!(result.snapshot_id(), Some(snapshot_id)); + assert_eq!(result.row_ids().unwrap().row_ids, row_ids); + } + } +} + +#[tokio::test] +async fn empty_prepared_filter_retains_its_snapshot_without_resolving_latest() { + let table = de_vector_table().await; + let prepared = table + .prepare_vector_search_filter(id_gt_filter(&table, 99)) + .await + .unwrap(); + let snapshot_id = prepared.table().travel_snapshot().unwrap().id(); + assert!(prepared.include_row_ids().is_empty()); + let manager = table.snapshot_manager(); + for id in manager.list_all_ids().await.unwrap() { + table + .file_io() + .delete_file(&manager.snapshot_path(id)) + .await + .unwrap(); + } + + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_prepared_filter(prepared) + .execute() + .await + .unwrap(); + assert_eq!(results.len(), 2); + for result in results { + assert!(result.is_empty()); + assert_eq!(result.snapshot_id(), Some(snapshot_id)); + assert!(result + .new_read_builder() + .read() + .await + .unwrap() + .try_next() + .await + .unwrap() + .is_none()); + } +} + +#[tokio::test] +async fn prepared_filter_cannot_bypass_builder_target_query_auth() { + let source = vector_test_table(); + let prepared = source + .prepare_vector_search_filter(id_gt_filter(&source, 0)) + .await + .unwrap(); + let target = source.copy_with_options(HashMap::from([( + "query-auth.enabled".to_string(), + "true".to_string(), + )])); + + let err = target + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0]]) + .with_limit(1) + .with_prepared_filter(prepared) + .execute() + .await + .expect_err("a stale prepared filter must not bypass current target authorization"); + + assert!( + matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), + "builder target authorization must remain authoritative, got: {err:?}" + ); +} + +#[tokio::test] +async fn de_vector_search_uses_time_travel_snapshot() { + let table = de_vector_table().await; + let latest = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(3) + .execute() + .await + .unwrap(); + assert!( + !latest.is_empty(), + "latest snapshot should contain the committed vector index" + ); + + let traveled = table + .copy_with_time_travel(HashMap::from([( + crate::spec::SCAN_VERSION_OPTION.to_string(), + "1".to_string(), + )])) + .await + .unwrap(); + assert_eq!( + traveled.travel_snapshot().map(|snapshot| snapshot.id()), + Some(1) + ); + + let historical = traveled + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(3) + .execute() + .await + .unwrap(); + assert!( + historical.is_empty(), + "snapshot 1 predates the vector index and should return no hits" + ); +} + +#[tokio::test] +async fn resolved_vector_snapshot_can_be_reused_by_all_read_stages() { + let table = de_vector_table().await; + let snapshot = crate::table::time_travel::resolve_snapshot(&table) + .await + .unwrap() + .unwrap(); + let pinned = table.copy_with_resolved_snapshot(&snapshot).await.unwrap(); + + assert_eq!( + pinned.travel_snapshot().map(|snapshot| snapshot.id()), + Some(snapshot.id()) + ); + let options = CoreOptions::new(pinned.schema().options()); + let selector = options.try_time_travel_selector().unwrap().unwrap(); + assert!(matches!( + selector, + crate::spec::TimeTravelSelector::SnapshotId { + value, + option_name: crate::spec::SCAN_SNAPSHOT_ID_OPTION, + } if value == snapshot.id().to_string() + )); +} + +#[tokio::test] +async fn de_scalar_filter_with_no_matching_rows_returns_empty() { + let table = de_vector_table().await; + let filter = id_gt_filter(&table, 99); + + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(2) + .with_filter(filter.clone()) + .execute() + .await + .unwrap(); + assert!(result.is_empty()); + + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_filter(filter) + .execute() + .await + .unwrap(); + assert_eq!(results.len(), 2); + assert!(results.iter().all(SearchResult::is_empty)); +} + +#[tokio::test] +async fn prepared_de_scalar_filter_can_be_reused_by_batch_search() { + let table = de_vector_table().await; + let prepared = table + .prepare_vector_search_filter(id_gt_filter(&table, 1)) + .await + .unwrap(); + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_prepared_filter(prepared) + .execute() + .await + .unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].row_ids().unwrap().row_ids, vec![2, 1]); + assert_eq!(results[1].row_ids().unwrap().row_ids, vec![1, 2]); +} + +#[tokio::test] +async fn prepared_filter_from_different_table_is_rejected() { + let prepared = PreparedVectorSearchFilter { + table: vector_test_table_at("memory:/prepared_filter_source"), + include_row_ids: Arc::new(RoaringTreemap::from_iter([1])), + }; + let target = vector_test_table_at("memory:/prepared_filter_target"); + + let error = target + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0]]) + .with_limit(1) + .with_prepared_filter(prepared) + .execute() + .await + .expect_err("a prepared filter must not retarget the builder to another table"); + + assert!( + error.to_string().contains("different table"), + "unexpected error: {error}" + ); +} + +#[tokio::test] +async fn de_scalar_filter_applies_to_unindexed_raw_fallback() { + let table = de_vector_table().await; + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = + ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); + vector_builder.values().append_value(1.0); + vector_builder.values().append_value(0.0); + vector_builder.append(true); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element_field), true), + ])), + vec![ + Arc::new(Int32Array::from(vec![4])) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap(); + let mut writer = TableWrite::new(&table, "test-user".to_string()).unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let table = table.copy_with_options(HashMap::from([ + ("vector-index.search-mode".to_string(), "full".to_string()), + ("scalar-index.search-mode".to_string(), "full".to_string()), + ])); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1) + .with_filter(id_gt_filter(&table, 3)) + .execute() + .await + .unwrap(); + + assert_eq!(result.row_ids().unwrap().row_ids, vec![3]); +} diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index af6dbb8c6..2ef9be1ae 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -29,13 +29,12 @@ use crate::table::data_file_reader::DataFileReader; use crate::table::pk_search_position::PrimaryKeySearchPosition; use crate::table::pk_search_ranker::{self, Ranking}; use crate::table::pk_vector_indexed_split_read::{PkVectorIndexedSplit, PkVectorIndexedSplitRead}; -use crate::table::pk_vector_orchestrator::build_indexed_splits; use crate::table::source::DataSplit; -use crate::table::vector_search_builder::{ +use crate::table::vector_search_common::{ collect_ranked_rows, ensure_no_reserved_read_columns, reorder_and_strip_position, RankedRow, }; use crate::table::{ArrowRecordBatchStream, RowRange, Table}; -use crate::vector_search::SearchResult; +use crate::vector_search::ScoredRowIds; #[cfg(feature = "fulltext")] use crate::spec::GlobalIndexSearchMode; @@ -285,7 +284,7 @@ impl<'a> HybridSearchBuilder<'a> { self.execute_scored().await?.to_row_ranges() } - pub async fn execute_scored(&self) -> crate::Result { + pub async fn execute_scored(&self) -> crate::Result { let core = CoreOptions::new(self.table.schema().options()); core.ensure_read_authorized()?; let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { @@ -324,7 +323,7 @@ impl<'a> HybridSearchBuilder<'a> { .with_query_vector(route.vector.clone().expect("validated vector route")) .with_limit(route.limit) .with_options(route.options.clone()); - builder.execute_scored().await? + builder.execute().await?.into_row_ids()? } HybridSearchRouteKind::FullText => { execute_full_text_route(self.table, route).await? @@ -437,9 +436,7 @@ impl<'a> HybridSearchBuilder<'a> { let mut routes: Vec = Vec::with_capacity(self.routes.len()); for route in &self.routes { let pk_route = match route.kind { - HybridSearchRouteKind::Vector => { - self.pk_vector_route(route_table, core, route).await? - } + HybridSearchRouteKind::Vector => self.pk_vector_route(route_table, route).await?, HybridSearchRouteKind::FullText => { self.pk_full_text_route(route_table, core, route).await? } @@ -562,13 +559,11 @@ impl<'a> HybridSearchBuilder<'a> { Ok(Some(pinned)) } - /// Run the vector route's primary-key candidate producer and convert its hits - /// into shared physical positions (distance → score via the resolved metric), - /// keeping the route's single-file source splits and pinned snapshot. + /// Consume the vector search's scored positions, retaining its source files + /// and pinned snapshot for fusion before materialization. async fn pk_vector_route( &self, table: &Table, - core: &CoreOptions<'_>, route: &HybridSearchRoute, ) -> crate::Result { let vector = route.vector.as_deref().expect("validated vector route"); @@ -578,24 +573,21 @@ impl<'a> HybridSearchBuilder<'a> { .with_query_vector(vector.to_vec()) .with_limit(route.limit) .with_options(route.options.clone()); - let result = builder - .search_pk_route(core, &route.field_name, vector, route.limit) - .await?; + let result = builder.execute().await?; let positions = result - .candidates + .positions()? .iter() - .map(|candidate| { - PrimaryKeySearchPosition::from_vector_candidate(candidate, result.metric) - }) + .map(PrimaryKeySearchPosition::from_vector_position) .collect::>>()?; - let source_splits = build_indexed_splits(result.candidates, &result.splits, result.metric)? - .into_iter() - .map(|split| split.split) + let source_splits = result + .indexed_splits()? + .iter() + .map(|split| split.split.clone()) .collect(); Ok(PkRoute { positions, source_splits, - snapshot_id: result.snapshot_id, + snapshot_id: result.snapshot_id().unwrap_or(0), weight: route.weight as f64, }) } @@ -901,7 +893,7 @@ fn build_hybrid_indexed_splits( async fn execute_full_text_route( table: &Table, route: &HybridSearchRoute, -) -> crate::Result { +) -> crate::Result { let mut builder = table.new_full_text_search_builder(); builder .with_text_column(&route.field_name) @@ -913,21 +905,21 @@ async fn execute_full_text_route( ) .with_limit(route.limit); let result = builder.execute_scored().await?; - Ok(SearchResult::new(result.row_ids, result.scores)) + Ok(ScoredRowIds::new(result.row_ids, result.scores)) } #[cfg(not(feature = "fulltext"))] async fn execute_full_text_route( _table: &Table, _route: &HybridSearchRoute, -) -> crate::Result { +) -> crate::Result { Err(crate::Error::ConfigInvalid { message: "Full-text hybrid routes require the fulltext feature".to_string(), }) } struct WeightedRouteResult { - result: SearchResult, + result: ScoredRowIds, weight: f32, } @@ -935,7 +927,7 @@ fn rank_results( ranker: HybridSearchRanker, route_results: &[WeightedRouteResult], limit: usize, -) -> SearchResult { +) -> ScoredRowIds { match ranker { HybridSearchRanker::Rrf => rrf(route_results, limit), HybridSearchRanker::WeightedScore => weighted_score(route_results, limit), @@ -943,7 +935,7 @@ fn rank_results( } } -fn rrf(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult { +fn rrf(route_results: &[WeightedRouteResult], limit: usize) -> ScoredRowIds { let mut scores = HashMap::new(); for route_result in route_results { for (rank, (row_id, _score)) in ranked_row_ids(&route_result.result).iter().enumerate() { @@ -954,7 +946,7 @@ fn rrf(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult { top_k(scores, limit) } -fn mrr(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult { +fn mrr(route_results: &[WeightedRouteResult], limit: usize) -> ScoredRowIds { let mut scores = HashMap::new(); for route_result in route_results { for (rank, (row_id, _score)) in ranked_row_ids(&route_result.result).iter().enumerate() { @@ -965,7 +957,7 @@ fn mrr(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult { top_k(scores, limit) } -fn weighted_score(route_results: &[WeightedRouteResult], limit: usize) -> SearchResult { +fn weighted_score(route_results: &[WeightedRouteResult], limit: usize) -> ScoredRowIds { let mut scores = HashMap::new(); for route_result in route_results { let ranked = ranked_row_ids(&route_result.result); @@ -992,7 +984,7 @@ fn weighted_score(route_results: &[WeightedRouteResult], limit: usize) -> Search top_k(scores, limit) } -fn ranked_row_ids(result: &SearchResult) -> Vec<(u64, f32)> { +fn ranked_row_ids(result: &ScoredRowIds) -> Vec<(u64, f32)> { let mut best_scores = HashMap::new(); for (&row_id, &score) in result.row_ids.iter().zip(&result.scores) { best_scores @@ -1022,9 +1014,9 @@ fn add_score(scores: &mut HashMap, row_id: u64, score: f32) { .or_insert(score); } -fn top_k(scores: HashMap, limit: usize) -> SearchResult { +fn top_k(scores: HashMap, limit: usize) -> ScoredRowIds { if scores.is_empty() || limit == 0 { - return SearchResult::empty(); + return ScoredRowIds::empty(); } let mut entries: Vec<_> = scores.into_iter().collect(); @@ -1037,7 +1029,7 @@ fn top_k(scores: HashMap, limit: usize) -> SearchResult { entries.truncate(limit); let (row_ids, scores): (Vec<_>, Vec<_>) = entries.into_iter().unzip(); - SearchResult::new(row_ids, scores) + ScoredRowIds::new(row_ids, scores) } #[cfg(test)] @@ -1046,7 +1038,7 @@ mod tests { fn route_result(row_ids: Vec, scores: Vec, weight: f32) -> WeightedRouteResult { WeightedRouteResult { - result: SearchResult::new(row_ids, scores), + result: ScoredRowIds::new(row_ids, scores), weight, } } @@ -1549,6 +1541,54 @@ mod pk_hybrid_tests { } } + #[tokio::test] + async fn pk_hybrid_weighted_score_uses_vector_scores_and_reads_in_fused_order() { + let table = build_hybrid_table( + "memory:/pk_hybrid_weighted_vector", + &[100, 101, 102], + &[ + [8.0, 0.0, 0.0, 0.0], + [10.0, 0.0, 0.0, 0.0], + [9.0, 0.0, 0.0, 0.0], + ], + &["alpha", "beta", "gamma"], + &[], + ) + .await; + let mut builder = table.new_hybrid_search_builder(); + builder + .add_vector_route( + VECTOR_COLUMN, + vec![10.0, 0.0, 0.0, 0.0], + 3, + 2.0, + HashMap::new(), + ) + .unwrap() + .with_limit(3) + .with_weighted_score_ranker(); + + let batches: Vec = builder + .execute_read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(column_i32(&batches, "id"), vec![101, 102, 100]); + // L2 distances [0, 1, 4] become scores [1, .5, .2]. Min-max normalization + // with weight 2 yields [2, .75, 0]. Raw distances or a second distance-to- + // score conversion would change both the order and these scores. + let scores = column_f32(&batches, SEARCH_SCORE_COLUMN); + assert_eq!(scores.len(), 3); + for (score, expected) in scores.into_iter().zip([2.0, 0.75, 0.0]) { + assert!( + (score - expected).abs() < 1e-6, + "expected {expected}, got {score}" + ); + } + } + // (c) A mixed PK/global route set must fail loud on execute_read. #[tokio::test] async fn mixed_pk_and_global_routes_fail_loud() { @@ -1859,5 +1899,13 @@ mod pk_hybrid_tests { Some(&"1".to_string()), "read-latest hybrid must pin the resolved latest snapshot id" ); + let route = builder + .pk_vector_route(&pinned, &builder.routes[0]) + .await + .unwrap(); + assert_eq!(route.snapshot_id, 1); + assert_eq!(route.positions.len(), 2); + assert_eq!(route.source_splits.len(), 1); + assert_eq!(route.source_splits[0].snapshot_id(), route.snapshot_id); } } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index bd98ba8ee..a909052e7 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -19,6 +19,7 @@ pub(crate) mod aggregator; mod audit_log_table; +mod batch_vector_search_builder; pub(crate) mod bin_pack; mod bitmap_global_index_format; mod bitmap_global_index_reader; @@ -39,6 +40,8 @@ mod data_evolution_reader; pub mod data_evolution_writer; mod data_file_reader; mod data_file_writer; +mod de_vector_read; +mod de_vector_scan; mod dedicated_format_file_writer; mod format_partition; mod format_read_builder; @@ -77,7 +80,9 @@ mod pk_vector_data_file_reader; mod pk_vector_indexed_split_read; mod pk_vector_orchestrator; mod pk_vector_position_read; +mod pk_vector_read; mod pk_vector_scan; +mod pk_vector_search_params; mod postpone_bucket_plan; mod postpone_file_writer; mod postpone_fixed_bucket_router; @@ -105,19 +110,27 @@ mod table_update; pub(crate) mod table_write; mod tag_manager; pub(crate) mod time_travel; +mod vector_read; +mod vector_scan; mod vector_search_builder; +mod vector_search_common; +pub(crate) mod vector_search_result; +#[cfg(test)] +mod vector_search_test_utils; mod vindex_index_build_builder; mod write_builder; use crate::Result; use arrow_array::RecordBatch; pub use audit_log_table::AuditLogTable; +pub use batch_vector_search_builder::BatchVectorSearchBuilder; pub use blob_resolver::{BlobReader, BlobStream}; pub use branch_manager::BranchManager; pub use commit_message::CommitMessage; pub use consumer_manager::ConsumerManager; pub use cow_writer::{CopyOnWriteMergeWriter, FileInfo}; pub use data_evolution_writer::{DataEvolutionDeleteWriter, DataEvolutionWriter}; +pub use de_vector_scan::PreparedVectorSearchFilter; pub use format_partition::{ format_partition_value, parse_format_partition_value, FormatTablePartitionPaths, }; @@ -161,9 +174,9 @@ pub use table_scan::TableScan; pub use table_update::TableUpdate; pub use table_write::TableWrite; pub use tag_manager::TagManager; -pub use vector_search_builder::{ - BatchVectorSearchBuilder, PreparedVectorSearchFilter, VectorSearchBuilder, -}; +pub use vector_read::{BatchVectorRead, VectorRead}; +pub use vector_scan::{VectorScan, VectorScanPlan}; +pub use vector_search_builder::VectorSearchBuilder; pub use vindex_index_build_builder::VindexIndexBuildBuilder; pub use write_builder::WriteBuilder; diff --git a/crates/paimon/src/table/pk_full_text_read.rs b/crates/paimon/src/table/pk_full_text_read.rs index 2eadc2e85..18376cf57 100644 --- a/crates/paimon/src/table/pk_full_text_read.rs +++ b/crates/paimon/src/table/pk_full_text_read.rs @@ -34,7 +34,7 @@ use crate::table::pk_full_text_bucket_search::search_bucket; use crate::table::pk_full_text_scan::{PrimaryKeyFullTextScanPlan, PrimaryKeyFullTextSearchSplit}; use crate::table::pk_vector_indexed_split_read::{PkVectorIndexedSplit, PkVectorIndexedSplitRead}; use crate::table::source::DataSplitBuilder; -use crate::table::vector_search_builder::{ +use crate::table::vector_search_common::{ collect_ranked_rows, reorder_and_strip_position, RankedRow, }; use crate::table::{ArrowRecordBatchStream, RowRange}; diff --git a/crates/paimon/src/table/pk_search_position.rs b/crates/paimon/src/table/pk_search_position.rs index 5cf5f3e2a..2d8c36fa4 100644 --- a/crates/paimon/src/table/pk_search_position.rs +++ b/crates/paimon/src/table/pk_search_position.rs @@ -111,16 +111,15 @@ impl PrimaryKeySearchPosition { ) } - pub(crate) fn from_vector_candidate( - candidate: &crate::table::pk_vector_orchestrator::PkVectorCandidate, - metric: crate::vindex::pkvector::metric::VectorSearchMetric, + pub(crate) fn from_vector_position( + position: &crate::vector_search::PrimaryKeySearchPosition, ) -> crate::Result { Self::new( - candidate.partition.clone(), - candidate.bucket, - candidate.data_file_name.clone(), - candidate.row_position, - metric.distance_to_score(candidate.distance), + position.partition.clone(), + position.bucket, + position.data_file_name.clone(), + position.row_position, + position.score, ) } @@ -169,8 +168,6 @@ impl Hash for PrimaryKeySearchPosition { #[cfg(test)] mod tests { use super::*; - use crate::table::pk_vector_orchestrator::PkVectorCandidate; - use crate::vindex::pkvector::metric::VectorSearchMetric; use std::collections::HashSet; fn pos(row_position: i64, score: f32) -> crate::Result { @@ -213,27 +210,31 @@ mod tests { assert_eq!(set.len(), 2); } - fn vector_candidate(distance: f32) -> PkVectorCandidate { - PkVectorCandidate { - split_index: 0, + fn vector_position(score: f32) -> crate::vector_search::PrimaryKeySearchPosition { + crate::vector_search::PrimaryKeySearchPosition { partition: BinaryRow::new(0), bucket: 0, data_file_name: "f".to_string(), row_position: 0, - distance, + score, } } #[test] - fn from_vector_candidate_applies_distance_to_score() { - // L2: score = 1/(1+distance); distance 1.0 -> score 0.5 (score != distance). - let candidate = vector_candidate(1.0); - let position = - PrimaryKeySearchPosition::from_vector_candidate(&candidate, VectorSearchMetric::L2) - .unwrap(); + fn from_vector_position_preserves_score_and_validates_position() { + // SearchResult already converted distance to score; fusion must preserve it. + let mut hit = vector_position(0.5); + let position = PrimaryKeySearchPosition::from_vector_position(&hit).unwrap(); assert_eq!(position.score(), 0.5); assert_eq!(position.row_position(), 0); assert_eq!(position.data_file_name(), "f"); + for score in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + hit.score = score; + assert!(PrimaryKeySearchPosition::from_vector_position(&hit).is_err()); + } + hit.score = 0.5; + hit.row_position = -1; + assert!(PrimaryKeySearchPosition::from_vector_position(&hit).is_err()); } #[cfg(feature = "fulltext")] diff --git a/crates/paimon/src/table/pk_vector_data_file_reader.rs b/crates/paimon/src/table/pk_vector_data_file_reader.rs index dc42f1728..99bd319ca 100644 --- a/crates/paimon/src/table/pk_vector_data_file_reader.rs +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -226,7 +226,7 @@ impl DataFilePkVectorReaderFactory { /// Extract one batch's vector column into `out`, one entry per row (NULL row = /// `None`). The column must be a `FixedSizeList`/`List` of `Float32`; every /// non-null row's child slice must have exactly `dimension` elements. Mirrors -/// the layout handling in `vector_search_builder`. +/// the layout handling in `de_vector_read`. pub(crate) fn append_batch_vectors( batch: &arrow_array::RecordBatch, field_name: &str, diff --git a/crates/paimon/src/table/pk_vector_indexed_split_read.rs b/crates/paimon/src/table/pk_vector_indexed_split_read.rs index d234ddcac..4feeacb2f 100644 --- a/crates/paimon/src/table/pk_vector_indexed_split_read.rs +++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs @@ -51,6 +51,7 @@ fn data_invalid(message: impl Into) -> crate::Error { /// /// Deliberately NOT reusing `DataSplit.row_ranges`, whose ranges mean stable/global /// row ids on the append/data-evolution path. Not serialized. +#[derive(Debug, Clone)] pub(crate) struct PkVectorIndexedSplit { pub split: DataSplit, pub row_ranges: Vec, diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index ca38b08fa..0d8c68c38 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -111,6 +111,7 @@ pub(crate) fn validate_row_position( /// One bucket's search input. Rust equivalent of Java /// `BucketVectorSearchSplit`. Constructed from a snapshot/manifest plan by /// `PkVectorScan`. +#[derive(Clone)] pub(crate) struct PkVectorSearchSplit { /// The bucket's combined data split (>= 1 data file); source of the /// partition/bucket/bucket_path/snapshot, the per-file `DataFileMeta`, and the @@ -579,6 +580,9 @@ impl PkVectorOrchestrator { } }); + // Erase the borrowing map closure before awaiting Send search futures. + let per_bucket = per_bucket.collect::>().into_iter(); + // Drive the per-bucket futures. `concurrency == 1` uses a strictly // sequential loop so buckets are searched in split order; larger values fan // them out through `drain_indexed_jobs`. Either way each bucket's per-query lists diff --git a/crates/paimon/src/table/pk_vector_read.rs b/crates/paimon/src/table/pk_vector_read.rs new file mode 100644 index 000000000..5869e9427 --- /dev/null +++ b/crates/paimon/src/table/pk_vector_read.rs @@ -0,0 +1,1026 @@ +// 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. + +//! Executes primary-key vector searches and materializes hits by bucket/file position. + +use crate::arrow::format::FilePredicates; +use crate::arrow::residual::{evaluate_predicates_mask, widen_scan_fields}; +use crate::lumina::reader::LuminaVectorGlobalIndexReader; +use crate::lumina::LuminaIndexMeta; +use crate::spec::{CoreOptions, DataField, Predicate}; +use crate::table::bucket_filter::split_partition_and_data_predicates; +use crate::table::data_file_reader::DataFileReader; +use crate::table::pk_vector_data_file_reader::{ + append_batch_vectors, DataFilePkVectorReaderFactory, +}; +use crate::table::pk_vector_indexed_split_read::{expand_ranges, PkVectorIndexedSplitRead}; +use crate::table::pk_vector_orchestrator::{ + as_split_exact_file_search, build_indexed_splits, merge_candidates, OrchestratorSearchResult, + PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, +}; +use crate::table::pk_vector_position_read::{PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN}; +use crate::table::pk_vector_scan::PkVectorScanPlan; +use crate::table::pk_vector_search_params::PkVectorSearchParams; +use crate::table::source::DataSplit; +use crate::table::vector_read::Read; +use crate::table::vector_search_common::{ + collect_ranked_rows, current_tokio_runtime_handle, log_vindex_range_io_stats, + reorder_and_strip_position, vindex_concurrency_limits, RankedRow, VectorIndexBackend, +}; +use crate::table::{ArrowRecordBatchStream, RowRange, Table}; +use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; +use crate::vindex::pkvector::ann::{AnnSegmentSource, PkVectorAnnSearcher, VindexAnnSearcher}; +use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture}; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::pkvector::{FileRowSelection, FileRowSelections}; +use crate::vindex::range_reader::{RangeReadLimiter, VindexFileReader}; +use crate::vindex::reader::VindexVectorGlobalIndexReader; +use arrow_array::{Array, Int64Array, RecordBatch}; +use futures::{stream, TryStreamExt}; +use roaring::RoaringTreemap; +use std::collections::{HashMap, HashSet}; +use std::io::Cursor; +use std::sync::Arc; + +pub(super) struct PkVectorRead { + table: Table, + options: HashMap, + filter: Option, + vector_column: String, + queries: Vec>, + limit: usize, + params: PkVectorSearchParams, +} + +impl PkVectorRead { + /// Construct a reader from parameters resolved and validated by the builder. + pub(super) fn new( + table: &Table, + options: &HashMap, + filter: Option<&Predicate>, + vector_column: &str, + queries: &[&[f32]], + limit: usize, + params: PkVectorSearchParams, + ) -> Self { + Self { + table: table.clone(), + options: options.clone(), + filter: filter.cloned(), + vector_column: vector_column.to_string(), + queries: queries.iter().map(|query| query.to_vec()).collect(), + limit, + params, + } + } +} + +impl Read for PkVectorRead { + type Plan = PkVectorScanPlan; + + async fn read(&self, plan: PkVectorScanPlan) -> crate::Result> { + let core = CoreOptions::new(self.table.schema().options()); + let queries: Vec<&[f32]> = self.queries.iter().map(Vec::as_slice).collect(); + let candidates = search_pk_candidates_batch_with_plan( + &self.table, + &self.options, + self.filter.as_ref(), + &core, + &self.vector_column, + &queries, + self.limit, + &plan, + &self.params, + ) + .await?; + let table = Arc::new(self.table.clone()); + candidates + .into_iter() + .map(|candidates| { + SearchResult::from_primary_key( + table.clone(), + plan.snapshot_id, + candidates, + &plan.splits, + self.params.metric, + ) + }) + .collect() + } +} + +/// Materialize one best-first candidate list into an Arrow stream, best-first, +/// with a `__paimon_search_score` column and `_PKEY_VECTOR_POSITION` stripped. +/// An empty candidate list yields an empty stream (never skipped) so a batch +/// caller preserves per-query arity. `materialize_reader` must project the +/// output columns (predicate-free). Both the single-query and batch read paths +/// use this so their materialization is identical. +pub(super) async fn materialize_positions( + positions: &[crate::vector_search::PrimaryKeySearchPosition], + indexed_splits: &[crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit], + materialize_reader: &DataFileReader, +) -> crate::Result { + if positions.is_empty() { + return Ok(Box::pin(stream::empty())); + } + let rank_of = positions + .iter() + .enumerate() + .map(|(rank, p)| { + ( + ( + p.partition.to_serialized_bytes(), + p.bucket, + p.data_file_name.clone(), + p.row_position, + ), + rank, + ) + }) + .collect(); + + // Materialize every indexed split, retaining each batch and, per row, the + // (rank, batch_index, row_index) tuple so we can reorder to best-first. + // Top-K is small, so full in-memory collection is acceptable. + let mut batches: Vec = Vec::new(); + let mut ranked: Vec = Vec::new(); + for indexed in indexed_splits { + let partition_bytes = indexed.split.partition().to_serialized_bytes(); + let bucket = indexed.split.bucket(); + let file_name = indexed.split.data_files()[0].file_name.clone(); + let mut stream = PkVectorIndexedSplitRead::new(materialize_reader.clone()).read(indexed)?; + while let Some(batch) = stream.try_next().await? { + let batch_index = batches.len(); + collect_ranked_rows( + &batch, + batch_index, + &partition_bytes, + bucket, + &file_name, + &rank_of, + &mut ranked, + )?; + batches.push(batch); + } + } + + // Reorder to best-first and drop the position column. + let output = reorder_and_strip_position(&batches, ranked)?; + Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) +} + +/// Search an already-resolved plan across every query and return each query's raw +/// indexed and exact candidate lists, before any rerank or merge. +/// +/// Plan-dependent concurrency — the vindex segment count, batch-index parallelism +/// and the range-read bound — is derived here from the plan that is actually being +/// searched, so a narrowed plan can never be searched under limits computed for a +/// wider one. +/// Combine the two per-split row allow-lists a search can be handed: the physical +/// rows an engine-supplied plan restricts each file to, and the positions a residual +/// data predicate leaves behind. +/// +/// The two sides read a file's ABSENCE differently, and the merge has to respect +/// both readings: +/// +/// * The plan lists only what the engine's split narrowed, so an absent file is +/// unrestricted -- Java's `rowRangesByFile.get(file) == null`. +/// * The residual is exhaustive over the files a search can read from +/// (`residual_positions_by_file` registers every active file, empty when nothing +/// passed), so once a residual exists its silence about a file means "no rows". +/// +/// So: with no residual, a file the plan omits stays absent and unrestricted. With a +/// residual, a file it omits is excluded even if the plan restricted it, and a file +/// both describe keeps the intersection. Absent from BOTH is unrestricted, which is +/// what lets the ANN backend search unfiltered. +/// +/// The plan's ranges stay ranges. Expanding them into positions would be work sized +/// by row counts that arrived on the wire; where an intersection is genuinely needed +/// the residual positions — bounded by the rows its own read returned — are filtered +/// BY the ranges instead. When the residual was evaluated over those same ranges the +/// intersection cannot remove anything, and is kept as the invariant that says so. +fn intersect_row_allow_lists( + physical: Option<&[HashMap>]>, + residual: Option>>, + split_count: usize, +) -> crate::Result>> { + if let Some(maps) = physical { + if maps.len() != split_count { + return Err(crate::Error::DataInvalid { + message: format!( + "plan carries {} physical row allow-lists for {split_count} splits", + maps.len() + ), + source: None, + }); + } + } + if let Some(maps) = residual.as_ref() { + if maps.len() != split_count { + return Err(crate::Error::DataInvalid { + message: format!( + "residual carries {} row allow-lists for {split_count} splits", + maps.len() + ), + source: None, + }); + } + } + match (physical, residual) { + (None, None) => Ok(None), + (None, Some(residual)) => Ok(Some( + residual + .into_iter() + .map(|per_file| { + per_file + .into_iter() + .map(|(file, positions)| (file, FileRowSelection::Positions(positions))) + .collect() + }) + .collect(), + )), + (Some(physical), None) => Ok(Some( + physical + .iter() + .map(|per_file| { + per_file + .iter() + .map(|(file, ranges)| { + (file.clone(), FileRowSelection::Ranges(ranges.clone())) + }) + .collect() + }) + .collect(), + )), + (Some(physical), Some(residual)) => { + Ok(Some( + physical + .iter() + .zip(residual) + .map(|(physical, mut residual)| { + let mut merged: FileRowSelections = HashMap::new(); + for (file, ranges) in physical { + let range_selection = FileRowSelection::Ranges(ranges.clone()); + let selection = match residual.remove(file.as_str()) { + // Both restrict: keep the positions the ranges also + // allow. Filtering the positions (bounded by the read) + // by the ranges never expands the ranges. + Some(positions) => FileRowSelection::Positions( + positions + .iter() + .filter(|position| range_selection.contains(*position)) + .collect(), + ), + // The residual is exhaustive over the files the search + // can read from -- `residual_positions_by_file` + // registers every active file, empty when nothing + // passed. Its silence about a file therefore means "no + // rows", NOT "unrestricted", and must stay fail-closed + // here even though the plan has something to say. + None => FileRowSelection::Positions(RoaringTreemap::new()), + }; + merged.insert(file.clone(), selection); + } + // Whatever the residual restricted and the plan did not. + merged.extend(residual.into_iter().map(|(file, positions)| { + (file, FileRowSelection::Positions(positions)) + })); + merged + }) + .collect(), + )) + } + } +} + +#[allow(clippy::too_many_arguments)] +async fn search_pk_raw_candidates_batch_with_plan( + table: &Table, + query_options: &HashMap, + filter: Option<&Predicate>, + core: &CoreOptions<'_>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, + plan: &PkVectorScanPlan, + params: &PkVectorSearchParams, +) -> crate::Result> { + // An empty plan has nothing to search. Returned before the backend is resolved + // so a table with no searchable data never errors on an unrecognized index type. + if plan.splits.is_empty() { + return Ok(queries + .iter() + .map(|_| OrchestratorSearchResult { + indexed: Vec::new(), + exact: Vec::new(), + }) + .collect()); + } + + let metric = params.metric; + let concurrency = params.concurrency; + let index_type = params.index_type.clone(); + let vector_field = params.vector_field.clone(); + let skip_exact_fallback = params.skip_exact_fallback; + let indexed_limit = params.indexed_limit; + + // Resolve the vector index backend from the single configured index type. + // Java enforces one index type per PK table and Rust filters segments to it, + // so one backend serves every segment. Computed after the empty-plan return so + // an empty table never errors on an unrecognized type. + let backend = VectorIndexBackend::from_index_type(&index_type).ok_or_else(|| { + crate::Error::DataInvalid { + message: format!("unsupported PK vector index backend/type: '{index_type}'"), + source: None, + } + })?; + let (batch_index_parallelism, range_read_concurrency) = match backend { + VectorIndexBackend::Vindex => vindex_concurrency_limits( + core, + plan.splits + .iter() + .map(|split| split.ann_segments.len()) + .sum(), + concurrency, + )?, + VectorIndexBackend::Lumina => (1, 0), + }; + + // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` + // but projecting only the vector column with no predicates. + let reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + vec![vector_field.clone()], + Vec::new(), + ); + + // Real ANN scorer + loader. Each segment source is opened lazily inside its + // bucket leaf and dropped after scoring. Lumina keeps its buffered-byte path; + // vindex remains range-backed and reads only metadata and probed lists. + let options = { + let mut o = table.schema().options().clone(); + o.extend(query_options.clone()); + o + }; + let search_options = options.clone(); + let field_name = pk_col.to_string(); + + let loader_io = table.file_io().clone(); + let loader_range_read_limiter = match backend { + VectorIndexBackend::Vindex => Some(RangeReadLimiter::new(range_read_concurrency)), + VectorIndexBackend::Lumina => None, + }; + let loader: crate::vindex::pkvector::ann::SourceSegmentLoader = Box::new( + move |segment: &BucketAnnSegment| { + let io = loader_io.clone(); + let range_read_limiter = loader_range_read_limiter.clone(); + let path = segment.path.clone(); + let file_size = segment.file_size; + Box::pin(async move { + let input = io.new_input(&path)?; + match backend { + VectorIndexBackend::Lumina => input + .read() + .await + .map(AnnSegmentSource::Buffered) + .map_err(|error| crate::Error::DataInvalid { + message: format!("failed to read ANN index file '{path}': {error}"), + source: None, + }), + VectorIndexBackend::Vindex => { + let file_reader = + input + .reader() + .await + .map_err(|error| crate::Error::DataInvalid { + message: format!( + "failed to open ANN index file '{path}' for range reads: {error}" + ), + source: None, + })?; + Ok(AnnSegmentSource::Vindex( + VindexFileReader::new_with_limiter( + Arc::new(file_reader), + current_tokio_runtime_handle()?, + range_read_limiter.expect("Vindex range-read limiter"), + file_size, + path, + ), + )) + } + } + }) + }, + ); + + let scorer: crate::vindex::pkvector::ann::SourceBatchScorer = Box::new( + move |segment: &BucketAnnSegment, source: AnnSegmentSource, searches: &[VectorSearch]| { + let io_meta = GlobalIndexIOMeta::new( + segment.path.clone(), + segment.file_size, + segment.index_meta.clone(), + ); + match (backend, source) { + (VectorIndexBackend::Lumina, AnnSegmentSource::Buffered(data)) => { + let lumina_metric = + LuminaIndexMeta::deserialize(&segment.index_meta)?.metric()?; + verify_segment_metric(metric, VectorSearchMetric::from_lumina(lumina_metric))?; + let mut reader = LuminaVectorGlobalIndexReader::new(io_meta, options.clone()); + reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data))) + } + (VectorIndexBackend::Vindex, AnnSegmentSource::Vindex(source)) => { + let range_io_stats = source.range_io_stats(); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()) + .with_batch_index_parallelism(batch_index_parallelism); + let results = reader.visit_batch_vector_search_validated( + searches, + |_| Ok(source), + |metadata| { + verify_segment_metric( + metric, + VectorSearchMetric::from_vindex(metadata.metric), + ) + }, + )?; + if let Some(stats) = range_io_stats { + log_vindex_range_io_stats(&segment.path, searches.len(), &stats); + } + Ok(results) + } + (VectorIndexBackend::Lumina, AnnSegmentSource::Vindex(_)) + | (VectorIndexBackend::Vindex, AnnSegmentSource::Buffered(_)) => { + Err(crate::Error::DataInvalid { + message: format!( + "ANN segment '{}' was loaded with the wrong backend source", + segment.path + ), + source: None, + }) + } + } + }, + ); + let ann_searcher: Arc = Arc::new(VindexAnnSearcher::new_with_source( + field_name, scorer, loader, + )); + + // Residual (post-recall) filtering: for each candidate file, re-read its + // physical rows and keep the positions whose rows satisfy the filter. The + // per-split allow-list is threaded into the bucket search so the residual folds + // into recall (best-first order and Top-K are preserved). Built only when the + // filter has data (non-partition) conjuncts; a partition-only filter (or no + // filter) leaves `None`, which leaves the search unfiltered — partition + // pruning is already handled in planning. The residual depends only on the + // filter and the plan, not the query vector, so it is computed once here and + // shared across every query in the batch. The residual reader projects only + // the predicate columns and carries no pushdown; `residual_positions_by_file` + // recovers each surviving row's file-local physical position from its ordinal + // in the unfiltered scan (no `_ROW_ID`, no `first_row_id`). A file the + // allow-list leaves empty is skipped by the bucket search without opening an + // exact reader. + let residual_by_split: Option>> = match filter { + Some(filter) => { + // The whole filter is pushed into scan planning (`PkVectorScan`), where + // partition-only conjuncts already prune partitions/files. Re-applying + // them as a per-row residual would be redundant, so keep only the data + // conjuncts here — a partition-only filter then needs no residual at + // all. Mixed partition/data conjuncts stay whole in `data_predicates` + // and evaluate against the materialized partition column (partition + // columns are physically present in primary-key data files), so there + // is no missing-column case to reject. + let (_partition_predicate, data_predicates) = split_partition_and_data_predicates( + filter.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ); + if data_predicates.is_empty() { + None + } else { + let file_predicates = FilePredicates { + predicates: data_predicates, + row_filter_factory: None, + file_fields: table.schema().fields().to_vec(), + }; + let residual_read_type = widen_scan_fields(&[], Some(&file_predicates)); + let residual_reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + residual_read_type, + Vec::new(), + ); + let mut per_split = Vec::with_capacity(plan.splits.len()); + for (index, split) in plan.splits.iter().enumerate() { + // The plan's selection for this split, so the residual is + // evaluated over the rows an engine-supplied split allows rather + // than over the whole file. + let allowed_rows = plan + .physical_row_ranges_by_split + .as_ref() + .and_then(|per_split| per_split.get(index)); + per_split.push( + residual_positions_by_file( + &residual_reader, + &split.data_split, + &split.active_files, + &file_predicates, + allowed_rows, + ) + .await?, + ); + } + Some(per_split) + } + } + None => None, + }; + // Fold the plan's own positional restriction into the same allow-list. A plan + // built from engine-supplied bucket splits carries the physical positions each + // file is limited to; a plan read from the index manifest carries none. Both + // sides list what is permitted, so combining them is an intersection. + let row_selections_by_split = intersect_row_allow_lists( + plan.physical_row_ranges_by_split.as_deref(), + residual_by_split, + plan.splits.len(), + )?; + + // Build the exact-fallback search on demand: the kernel calls this only for a + // file it actually searches (uncovered by ANN, residual-allowed, and only when + // the search mode is not FAST). Everything the future needs is cloned/owned up + // front so it borrows neither the split nor the file across the await. The + // search streams the file's vector column one Arrow batch at a time into + // per-query bounded heaps (all queries share one stream). + let reader_for_factory = reader.clone(); + let vector_field_for_factory = vector_field.clone(); + // The plan's own per-file selection, so an exact fallback reads only the rows an + // engine-supplied split allows. `is_excluded` still rejects on top of it, but it + // cannot un-read a row. + let physical_for_factory = plan.physical_row_ranges_by_split.clone(); + let factory = as_split_exact_file_search( + move |split_index: usize, + split: &PkVectorSearchSplit, + file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + let reader = reader_for_factory.clone(); + let vector_field = vector_field_for_factory.clone(); + let data_split = split.data_split.clone(); + let active = BucketActiveFile { + file_name: file.file_name.clone(), + row_count: file.row_count, + }; + let owned_queries: Vec> = queries.iter().map(|q| q.to_vec()).collect(); + let allowed_rows = physical_for_factory.as_ref().and_then(|per_split| { + per_split + .get(split_index) + .and_then(|per_file| per_file.get(&active.file_name)) + .cloned() + }); + Box::pin(async move { + let factory = DataFilePkVectorReaderFactory::new(reader, data_split, vector_field)?; + let query_refs: Vec<&[f32]> = owned_queries.iter().map(|q| q.as_slice()).collect(); + factory + .search_file( + &active, + &query_refs, + metric, + exact_limit, + is_excluded, + allowed_rows.as_deref(), + ) + .await + }) + }, + ); + + // Resolve the refine factor from the query options first, then fall back to the + // table options; a positive factor over-fetches indexed (approximate) + // candidates so the exact rerank below has a wider pool to reorder. Factor 0 + // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank + // path. The two option maps are kept distinct (query options passed separately + // from table options) so a broad query key cannot be overridden by a more + // specific table key: query options take precedence as a whole. `search_options` + // above is the merged view used only to drive the ANN read. + + let searches: Vec = PkVectorOrchestrator::new(reader) + .search_candidates_batch( + &plan.splits, + queries, + metric, + limit, + indexed_limit, + Some(ann_searcher), + &factory, + &search_options, + skip_exact_fallback, + row_selections_by_split.as_deref(), + concurrency, + ) + .await?; + + Ok(searches) +} + +/// Search an already-resolved plan and return one merged, best-first candidate list +/// per query: the raw layer above, followed by the optional exact rerank of the +/// approximate candidates and the merge with the exact-fallback candidates. +#[allow(clippy::too_many_arguments)] +async fn search_pk_candidates_batch_with_plan( + table: &Table, + query_options: &HashMap, + filter: Option<&Predicate>, + core: &CoreOptions<'_>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, + plan: &PkVectorScanPlan, + params: &PkVectorSearchParams, +) -> crate::Result>> { + let searches = search_pk_raw_candidates_batch_with_plan( + table, + query_options, + filter, + core, + pk_col, + queries, + limit, + plan, + params, + ) + .await?; + + let metric = params.metric; + let refine_factor = params.refine_factor; + let vector_field = params.vector_field.clone(); + + // Per query: exact rerank of the approximate candidates when a refine factor is + // set (exact-fallback candidates are already exact and are not reranked), then + // merge the (possibly reranked) indexed list with the exact list into one + // best-first list bounded to the caller's limit. With no refine factor the + // rerank is a plain merge, byte-identical to the no-rerank path. Each query + // reranks its OWN indexed candidates. + let mut per_query_candidates = Vec::with_capacity(searches.len()); + for (query_index, search) in searches.into_iter().enumerate() { + let query_vector = queries[query_index]; + let indexed = if refine_factor > 0 && !search.indexed.is_empty() { + // Vector-only reader (project just the vector field); the position read + // appends _PKEY_VECTOR_POSITION itself and injects _ROW_ID internally. + let rerank_reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + vec![vector_field.clone()], + Vec::new(), + ); + rerank_indexed_positional( + &rerank_reader, + search.indexed, + &plan.splits, + query_vector, + metric, + limit, + &vector_field, + ) + .await? + } else { + search.indexed + }; + per_query_candidates.push(merge_candidates(indexed, search.exact, limit)); + } + + Ok(per_query_candidates) +} + +/// Compute, per data file in `split`, the set of file-LOCAL physical row +/// positions whose rows satisfy the residual predicate. Mirrors the +/// row-collecting half of Java `PrimaryKeyVectorRead`'s `executeFilter`: the +/// predicate is NOT pushed down (a pushed filter would drop rows before their +/// position could be recovered). Instead `reader` projects only the residual +/// columns and carries no pushdown predicate, the residual is evaluated here at the +/// Arrow level, and each surviving row's file-local 0-based position is recovered +/// from the selection the read was limited to. This needs no `_ROW_ID` and no +/// `first_row_id` — real primary-key tables never write one. +/// +/// `allowed_rows` is the plan's per-file physical selection, keyed by data-file +/// name, with the plan's three states: a file it does not list is unrestricted and +/// the whole file is scanned; an empty range list excludes the file, which is +/// registered empty without a read; a non-empty list is scanned over exactly those +/// ranges, because an engine-supplied bucket split can restrict a huge file to a +/// handful of ranges and reading all of it to discard the rest would defeat the +/// split. +/// +/// Every *active* data file in the split gets an entry in the RESULT, possibly +/// empty, and that exhaustiveness is load-bearing. The search kernel reads a file's +/// absence from its selections as "unrestricted", so an active file missing here +/// would reach the search with no predicate applied at all -- the residual would be +/// silently dropped for it. (The merge below reads a residual's silence about a +/// file the PLAN listed as exclusion, so only a file both omit falls through, which +/// is exactly the case this exhaustiveness rules out.) Non-active files (e.g. +/// level-0 files the bucket search excludes) are skipped entirely: they are never +/// searched, so re-reading them would be wasted IO. +/// +/// `reader` must be predicate-free and project the residual columns; +/// `residual.file_fields` are the fields the residual leaf indices point into +/// (resolved by name against each emitted batch). +async fn residual_positions_by_file( + reader: &DataFileReader, + split: &DataSplit, + active_files: &[BucketActiveFile], + residual: &FilePredicates, + allowed_rows: Option<&HashMap>>, +) -> crate::Result> { + let scan_fields = reader.read_type().to_vec(); + let active_names: HashSet<&str> = active_files.iter().map(|f| f.file_name.as_str()).collect(); + let mut out: HashMap = HashMap::new(); + for file_meta in split.data_files() { + // Only files the bucket search actually recalls from need residual + // positions; skip everything else to avoid a wasted read. + if !active_names.contains(file_meta.file_name.as_str()) { + continue; + } + // A file the plan lists an EMPTY range list for permits nothing; registering + // it empty says so and costs no read. A file the plan does not list at all + // is unrestricted, so the residual is evaluated over the whole file. + let selection = match allowed_rows.and_then(|by_file| by_file.get(&file_meta.file_name)) { + Some(ranges) if ranges.is_empty() => { + out.entry(file_meta.file_name.clone()).or_default(); + continue; + } + Some(ranges) => Some(ranges.clone()), + None => None, + }; + let data_fields = reader.derive_data_fields(file_meta).await?; + let mut stream = match selection.clone() { + Some(ranges) => reader.read_single_file_stream_local_ranges( + split, + file_meta.clone(), + data_fields, + None, + ranges, + )?, + None => { + reader.read_single_file_stream(split, file_meta.clone(), data_fields, None, None)? + } + }; + // Register the file up front so a file whose rows all fail the residual + // still appears in the map (empty set). + let positions = out.entry(file_meta.file_name.clone()).or_default(); + // Rows arrive in ascending physical order, and the read emitted exactly what + // was selected (no pushdown predicate, no deletion vector), so walking the + // selection in step with the rows recovers each row's file-local position. + let mut selected: Box + Send> = match &selection { + Some(ranges) => Box::new( + ranges + .clone() + .into_iter() + .flat_map(|range| (range.from() as u64)..=(range.to() as u64)), + ), + None => Box::new(0..file_meta.row_count.max(0) as u64), + }; + while let Some(batch) = stream.try_next().await? { + let num_rows = batch.num_rows(); + let mask = evaluate_predicates_mask( + &batch, + &residual.predicates, + &residual.file_fields, + &scan_fields, + )?; + for row_index in 0..num_rows { + let position = selected.next().ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "residual scan of '{}' emitted more rows than the selection allows", + file_meta.file_name + ), + source: None, + })?; + let keep = match &mask { + // NULL follows the same NULL -> false convention the Arrow filter + // kernel applies, so a null mask slot drops the row. + Some(mask) => mask.is_valid(row_index) && mask.value(row_index), + // No predicate contributed a mask (identity) -> keep every row. + None => true, + }; + if keep { + positions.insert(position); + } + } + } + if selected.next().is_some() { + return Err(crate::Error::DataInvalid { + message: format!( + "residual scan of '{}' emitted fewer rows than the selection allows", + file_meta.file_name + ), + source: None, + }); + } + } + Ok(out) +} + +fn verify_segment_metric( + configured: VectorSearchMetric, + segment_metric: VectorSearchMetric, +) -> crate::Result<()> { + if segment_metric != configured { + return Err(crate::Error::DataInvalid { + message: format!( + "ANN segment metric {} does not match configured metric {}", + segment_metric.as_str(), + configured.as_str() + ), + source: None, + }); + } + Ok(()) +} + +/// Rerank approximate (indexed) candidates by rereading ONLY their candidate +/// positions and recomputing the exact distance, then keep the best `limit`. +/// +/// Unlike a whole-column preload, this reuses [`PkVectorPositionRead`] to read +/// just the selected physical rows of each hit file (positions -> row ranges -> +/// local ranges), so a rerank over a large ANN-covered file touches only the +/// candidate rows. Mirrors Java's IndexedSplit rerank. +/// +/// Each returned row is matched back to its candidate by the +/// `_PKEY_VECTOR_POSITION` column VALUE (never batch order). The recomputed +/// distance is written into the ORIGINAL candidate so `split_index` / +/// partition / bucket survive (`build_indexed_splits` does not carry +/// `split_index`). A DV loaded exactly as [`PkVectorIndexedSplitRead::read`] +/// does drops deleted positions, so a candidate at a deleted position returns no +/// row and trips the leftover guard — a deleted candidate reaching rerank is a +/// real inconsistency (the search path already DV-filters), so fail loud. +#[allow(clippy::too_many_arguments)] +async fn rerank_indexed_positional( + rerank_reader: &DataFileReader, + indexed: Vec, + plan_splits: &[PkVectorSearchSplit], + query_vector: &[f32], + metric: VectorSearchMetric, + limit: usize, + vector_field: &DataField, +) -> crate::Result> { + // Original per-position candidates keyed by (split_index, file, position); + // the recomputed distance is written back into these so split_index and + // partition/bucket survive (build_indexed_splits does not carry split_index). + let mut by_key: HashMap<(usize, String, i64), PkVectorCandidate> = HashMap::new(); + for c in &indexed { + if by_key + .insert( + (c.split_index, c.data_file_name.clone(), c.row_position), + c.clone(), + ) + .is_some() + { + return Err(crate::Error::DataInvalid { + message: "duplicate primary-key vector candidate for reranking".to_string(), + source: None, + }); + } + } + + // Rebuild the split_index lookup by (partition bytes, bucket, file): the + // indexed split exposes partition/bucket/file but not split_index. + let mut split_index_of: HashMap<(Vec, i32, String), usize> = HashMap::new(); + for (i, s) in plan_splits.iter().enumerate() { + let p = s.data_split.partition().to_serialized_bytes(); + let b = s.data_split.bucket(); + for f in s.data_split.data_files() { + split_index_of.insert((p.clone(), b, f.file_name.clone()), i); + } + } + + // Every candidate must reference a (partition, bucket, file) that the plan + // actually carries. Checking up front — before build_indexed_splits, which + // indexes plan_splits by split_index — turns an absent file into a fail-loud + // error rather than an out-of-range panic, and keeps the per-split lookup + // below a self-consistent backstop. + for c in &indexed { + let key = ( + c.partition.to_serialized_bytes(), + c.bucket, + c.data_file_name.clone(), + ); + if !split_index_of.contains_key(&key) { + return Err(crate::Error::DataInvalid { + message: format!("rerank split for {} not found in plan", c.data_file_name), + source: None, + }); + } + } + + // Group the candidates into per-file indexed splits (position ranges + file + // meta), reusing the exact grouping/validation the materialization path uses. + let indexed_splits = build_indexed_splits(indexed, plan_splits, metric)?; + + let dimension = query_vector.len(); + let mut reranked: Vec = Vec::new(); + for split in indexed_splits { + let data_split = split.split.clone(); + let file_meta = data_split.data_files()[0].clone(); + let file_name = file_meta.file_name.clone(); + let partition_bytes = data_split.partition().to_serialized_bytes(); + let bucket = data_split.bucket(); + let split_index = *split_index_of + .get(&(partition_bytes, bucket, file_name.clone())) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("rerank split for {file_name} not found in plan"), + source: None, + })?; + + // DV loaded exactly as PkVectorIndexedSplitRead::read does; skipping it + // would score deleted rows. + let dv_factory = rerank_reader.build_split_dv_factory(&data_split).await?; + let dv = DataFileReader::deletion_vector_for_file(dv_factory.as_ref(), &file_name); + let data_fields = rerank_reader.derive_data_fields(&file_meta).await?; + + // Positions from the split's row_ranges (ascending); read only those. + let positions = expand_ranges(&split.row_ranges, file_meta.row_count)?; + let mut stream = PkVectorPositionRead::new(rerank_reader).read( + &data_split, + file_meta, + data_fields, + dv, + positions, + None, // no scores; rerank recomputes distance + )?; + + while let Some(batch) = stream.try_next().await? { + let pos_idx = batch + .schema() + .index_of(PKEY_VECTOR_POSITION_COLUMN) + .map_err(|_| crate::Error::DataInvalid { + message: format!("rerank batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), + source: None, + })?; + let pos_col = batch + .column(pos_idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is not Int64"), + source: None, + })?; + let mut vectors: Vec>> = Vec::new(); + append_batch_vectors(&batch, vector_field.name(), dimension, &mut vectors)?; + for (row, vector) in vectors.iter().enumerate() { + let position = pos_col.value(row); + let mut candidate = by_key + .remove(&(split_index, file_name.clone(), position)) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("rerank read unexpected position {file_name}@{position}"), + source: None, + })?; + let vector = vector.as_ref().ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "primary-key vector candidate {file_name}@{position} contains a null vector" + ), + source: None, + })?; + candidate.distance = metric.compute_distance(query_vector, vector); + reranked.push(candidate); + } + } + } + + if !by_key.is_empty() { + return Err(crate::Error::DataInvalid { + message: format!( + "failed to read {} primary-key vector candidate(s) for reranking", + by_key.len() + ), + source: None, + }); + } + + Ok(merge_candidates(reranked, Vec::new(), limit)) +} + +#[cfg(test)] +mod tests; + +#[cfg(test)] +mod residual_positions_tests; diff --git a/crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs b/crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs new file mode 100644 index 000000000..7dc65762b --- /dev/null +++ b/crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs @@ -0,0 +1,536 @@ +// 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::*; +use crate::arrow::build_target_arrow_schema; +use crate::arrow::format::FilePredicates; +use crate::io::FileIOBuilder; +use crate::spec::stats::BinaryTableStats; +use crate::spec::{ + BigIntType, BinaryRow, DataField, DataFileMeta, DataType, Datum, IntType, PredicateBuilder, + ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, +}; +use crate::table::data_file_reader::DataFileReader; +use crate::table::merge_row_ranges; +use crate::table::schema_manager::SchemaManager; +use crate::table::source::{DataSplit, DataSplitBuilder}; +use crate::table::vector_search_common::take_only_result; +use arrow_array::{Int32Array, RecordBatch}; +use bytes::Bytes; +use paimon_mosaic_core::spec::COMPRESSION_NONE; +use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions}; +use std::io; +use std::sync::Arc; + +struct MemOutputFile { + data: Vec, +} + +impl OutputFile for MemOutputFile { + fn write(&mut self, data: &[u8]) -> io::Result<()> { + self.data.extend_from_slice(data); + Ok(()) + } + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + fn pos(&self) -> u64 { + self.data.len() as u64 + } +} + +fn id_field() -> DataField { + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())) +} + +fn row_id_field() -> DataField { + DataField::new( + ROW_ID_FIELD_ID, + ROW_ID_FIELD_NAME.to_string(), + DataType::BigInt(BigIntType::new()), + ) +} + +fn id_batch(ids: Vec) -> RecordBatch { + let schema = build_target_arrow_schema(&[id_field()]).unwrap(); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(ids))]).unwrap() +} + +fn write_mosaic(batch: &RecordBatch) -> Bytes { + let mut writer = MosaicWriter::new( + MemOutputFile { data: Vec::new() }, + batch.schema().as_ref(), + WriterOptions { + compression: COMPRESSION_NONE, + num_buckets: 2, + row_group_max_size: u64::MAX, + ..Default::default() + }, + ) + .unwrap(); + writer.write_batch(batch).unwrap(); + writer.close().unwrap(); + Bytes::from(writer.output().data.to_vec()) +} + +fn data_file( + file_name: &str, + file_size: i64, + row_count: i64, + first_row_id: Option, +) -> DataFileMeta { + DataFileMeta { + file_name: file_name.to_string(), + file_size, + row_count, + min_key: Vec::new(), + max_key: Vec::new(), + key_stats: BinaryTableStats::empty(), + value_stats: BinaryTableStats::empty(), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 1, + level: 0, + extra_files: Vec::new(), + creation_time: None, + delete_row_count: None, + embedded_index: None, + file_source: None, + value_stats_cols: None, + external_path: None, + first_row_id, + write_cols: None, + column_max_sequence_numbers: None, + } +} + +/// Build a predicate-free reader (read_type = `id` + `_ROW_ID`) over a split +/// containing `files` (each `(name, ids, first_row_id)`), written as Mosaic +/// data files in the same bucket. The returned active-file list covers every +/// file (all files active). +async fn build_reader_and_split( + table_path: &str, + files: &[(&str, Vec, i64)], +) -> (DataFileReader, DataSplit, Vec) { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let bucket_path = format!("{table_path}/bucket-0"); + let mut metas = Vec::new(); + let mut active_files = Vec::new(); + for (name, ids, first_row_id) in files { + let data = write_mosaic(&id_batch(ids.clone())); + file_io + .new_output(&format!("{bucket_path}/{name}")) + .unwrap() + .write(data.clone()) + .await + .unwrap(); + metas.push(data_file( + name, + data.len() as i64, + ids.len() as i64, + Some(*first_row_id), + )); + active_files.push(BucketActiveFile { + file_name: name.to_string(), + row_count: ids.len() as i64, + }); + } + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(metas) + .build() + .unwrap(); + let reader = DataFileReader::new( + file_io.clone(), + SchemaManager::new(file_io, table_path.to_string()), + 1, + vec![id_field()], + vec![id_field(), row_id_field()], + Vec::new(), + ); + (reader, split, active_files) +} + +/// `id > threshold`, with `file_fields` = `[id]` so the leaf index resolves. +fn residual_id_gt(threshold: i32) -> FilePredicates { + let pred = PredicateBuilder::new(&[id_field()]) + .greater_than("id", Datum::Int(threshold)) + .unwrap(); + FilePredicates { + predicates: vec![pred], + row_filter_factory: None, + file_fields: vec![id_field()], + } +} + +fn sorted(t: &roaring::RoaringTreemap) -> Vec { + t.iter().collect() +} + +#[tokio::test] +async fn test_residual_selects_matching_positions() { + // ids [1,2,3,4,5] at first_row_id 0; id > 2 -> ids 3,4,5 -> positions 2,3,4. + let (reader, split, active) = build_reader_and_split( + "memory:/rpf_basic", + &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], + ) + .await; + let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), None) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]); +} + +#[tokio::test] +async fn test_residual_only_evaluates_the_rows_the_plan_allows() { + // ids [1,2,3,4,5]; the plan allows positions 3-4 only. `id > 2` matches 2,3,4 + // over the whole file, so a result of 3,4 is the plan's restriction taking + // effect *before* evaluation: position 2 is never seen. + // + // This also cannot pass under a full read. The scan walks the selection in + // step with the emitted rows, so a read that emitted all five would run the + // selection dry and fail loudly rather than return a filtered answer. + let (reader, split, active) = build_reader_and_split( + "memory:/rpf_plan_ranges", + &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], + ) + .await; + let allowed = HashMap::from([("part-0.mosaic".to_string(), vec![RowRange::new(3, 4)])]); + let map = + residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), Some(&allowed)) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]); +} + +#[tokio::test] +async fn test_residual_does_not_read_a_file_the_plan_excludes() { + // An EMPTY range list is how a plan says "no rows of this file": it is + // registered empty and never opened. Absence means the opposite -- the plan + // narrowed nothing there -- so the residual reads the whole file. + let (reader, split, active) = build_reader_and_split( + "memory:/rpf_plan_excludes", + &[("part-0.mosaic", vec![1, 2, 3], 0)], + ) + .await; + + let excluded = HashMap::from([("part-0.mosaic".to_string(), Vec::new())]); + let map = residual_positions_by_file( + &reader, + &split, + &active, + &residual_id_gt(0), + Some(&excluded), + ) + .await + .unwrap(); + assert!(map.contains_key("part-0.mosaic")); + assert!(sorted(&map["part-0.mosaic"]).is_empty()); + + let unrestricted = HashMap::new(); + let map = residual_positions_by_file( + &reader, + &split, + &active, + &residual_id_gt(0), + Some(&unrestricted), + ) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); +} + +#[tokio::test] +async fn test_residual_matches_none_yields_empty_entry() { + // id > 100 matches nothing; the file still gets a (present, empty) entry. + let (reader, split, active) = + build_reader_and_split("memory:/rpf_none", &[("part-0.mosaic", vec![1, 2, 3], 0)]).await; + let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(100), None) + .await + .unwrap(); + assert!(map.contains_key("part-0.mosaic")); + assert!(map["part-0.mosaic"].is_empty()); +} + +#[tokio::test] +async fn test_residual_matches_all_yields_full_set() { + let (reader, split, active) = + build_reader_and_split("memory:/rpf_all", &[("part-0.mosaic", vec![1, 2, 3], 0)]).await; + let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(0), None) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); +} + +#[tokio::test] +async fn test_residual_positions_are_file_local_across_files() { + // Two files with distinct first_row_id; positions must be 0-based within + // each file, not global. id > 3 keeps ids 4,5 in both -> positions {3,4}. + let (reader, split, active) = build_reader_and_split( + "memory:/rpf_multi", + &[ + ("part-0.mosaic", vec![1, 2, 3, 4, 5], 0), + ("part-1.mosaic", vec![1, 2, 3, 4, 5], 100), + ], + ) + .await; + let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(3), None) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]); + assert_eq!(sorted(&map["part-1.mosaic"]), vec![3, 4]); +} + +#[tokio::test] +async fn test_non_active_files_are_skipped() { + // Two files in the split, but only `part-0.mosaic` is active. The bucket + // search never recalls from `part-1.mosaic` (level-0 / non-active), so it + // must not appear in the residual map — and even though it lacks a + // `first_row_id`, the query still succeeds because non-active files are + // skipped before the guard. + let (reader, split, mut active) = build_reader_and_split( + "memory:/rpf_nonactive", + &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], + ) + .await; + // Append a non-active file (missing first_row_id) directly to the split's + // data files, but leave it out of the active list. + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let bucket_path = "memory:/rpf_nonactive/bucket-0"; + let data = write_mosaic(&id_batch(vec![9, 9, 9])); + file_io + .new_output(&format!("{bucket_path}/part-1.mosaic")) + .unwrap() + .write(data.clone()) + .await + .unwrap(); + let mut metas = split.data_files().to_vec(); + metas.push(data_file("part-1.mosaic", data.len() as i64, 3, None)); + // `active` already lists only part-0.mosaic; keep it that way. + let _ = &mut active; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path.to_string()) + .with_total_buckets(1) + .with_data_files(metas) + .build() + .unwrap(); + let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), None) + .await + .unwrap(); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]); + assert!( + !map.contains_key("part-1.mosaic"), + "non-active file must be skipped" + ); +} + +#[tokio::test] +async fn test_missing_first_row_id_recovers_local_positions() { + // Real primary-key data files carry no `first_row_id`. Positions are + // recovered from each row's ordinal in the scan, so the residual still + // works: ids [1,2,3] with id > 0 -> all match -> local positions [0,1,2]. + let (reader, split, active) = build_reader_and_split_no_first_row_id().await; + let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(0), None) + .await + .expect("missing first_row_id must not fail the residual read"); + assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); +} + +async fn build_reader_and_split_no_first_row_id( +) -> (DataFileReader, DataSplit, Vec) { + let table_path = "memory:/rpf_nofrid"; + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let bucket_path = format!("{table_path}/bucket-0"); + let data = write_mosaic(&id_batch(vec![1, 2, 3])); + file_io + .new_output(&format!("{bucket_path}/part-0.mosaic")) + .unwrap() + .write(data.clone()) + .await + .unwrap(); + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![data_file("part-0.mosaic", data.len() as i64, 3, None)]) + .build() + .unwrap(); + let reader = DataFileReader::new( + file_io.clone(), + SchemaManager::new(file_io, table_path.to_string()), + 1, + vec![id_field()], + vec![id_field(), row_id_field()], + Vec::new(), + ); + // The lone file is active and carries no first_row_id, exercising the + // ordinal-based position recovery. + let active = vec![BucketActiveFile { + file_name: "part-0.mosaic".to_string(), + row_count: 3, + }]; + (reader, split, active) +} + +// ---- combining the plan's positional restriction with the residual ---- + +fn allow_list(entries: &[(&str, &[u64])]) -> HashMap { + entries + .iter() + .map(|(file, positions)| ((*file).to_string(), positions.iter().copied().collect())) + .collect() +} + +/// The plan side carries ranges, so its fixtures are built from the positions +/// each file allows and coalesced the way the planner normalizes them. +fn range_allow_list(entries: &[(&str, &[u64])]) -> HashMap> { + entries + .iter() + .map(|(file, positions)| { + let ranges = positions + .iter() + .map(|p| RowRange::new(*p as i64, *p as i64)) + .collect(); + ((*file).to_string(), merge_row_ranges(ranges)) + }) + .collect() +} + +/// The positions a merged selection allows, expanded for readable assertions. +/// Test-only: the production path never expands a range. +fn listed(map: &FileRowSelections, file: &str) -> Vec { + match map.get(file) { + None => Vec::new(), + Some(FileRowSelection::Positions(positions)) => positions.iter().collect(), + Some(FileRowSelection::Ranges(ranges)) => ranges + .iter() + .flat_map(|range| (range.from() as u64)..=(range.to() as u64)) + .collect(), + } +} + +#[test] +fn no_restriction_on_either_side_stays_unrestricted() { + assert!(intersect_row_allow_lists(None, None, 1).unwrap().is_none()); +} + +#[test] +fn one_side_alone_passes_through() { + let physical = vec![range_allow_list(&[("d0", &[1, 2])])]; + let only_physical = intersect_row_allow_lists(Some(&physical), None, 1) + .unwrap() + .expect("a plan restriction survives on its own"); + assert_eq!(listed(&only_physical[0], "d0"), vec![1, 2]); + // Still intervals. Expanding them here is the unbounded step the plan side + // must never take, and the positions above cannot tell the two apart. + assert!( + matches!(only_physical[0]["d0"], FileRowSelection::Ranges(_)), + "the plan's ranges must reach the search as ranges" + ); + + let residual = vec![allow_list(&[("d0", &[3])])]; + let only_residual = intersect_row_allow_lists(None, Some(residual), 1) + .unwrap() + .expect("a residual survives on its own"); + assert_eq!(listed(&only_residual[0], "d0"), vec![3]); +} + +#[test] +fn both_sides_intersect_and_the_residual_stays_fail_closed() { + // `d0`: both restrict it, so only the shared positions survive. `d1`: the + // residual says nothing about it. The residual registers EVERY file the + // search can read from, so its silence is "no rows" -- the plan's ranges + // must not resurrect the file, and neither may its absence make it + // unrestricted. + let physical = vec![range_allow_list(&[("d0", &[1, 2, 3]), ("d1", &[0, 1])])]; + let residual = vec![allow_list(&[("d0", &[2, 3, 4])])]; + let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) + .unwrap() + .expect("both sides restrict"); + assert_eq!(listed(&combined[0], "d0"), vec![2, 3]); + assert!( + combined[0]["d1"].is_excluded(), + "a file the residual omits must stay excluded" + ); +} + +#[test] +fn a_file_neither_side_restricts_stays_absent() { + // Absence is how "every row" is spelled. A merged map must not invent an + // entry for a file no one narrowed, or the ANN backend takes the filtered + // path for a query that filters nothing. + let physical = vec![range_allow_list(&[("d0", &[1])])]; + let combined = intersect_row_allow_lists(Some(&physical), None, 1) + .unwrap() + .expect("the plan restricts d0"); + assert!(!combined[0].contains_key("d1")); + + let residual = vec![allow_list(&[("d0", &[1])])]; + let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) + .unwrap() + .expect("both restrict d0"); + assert!(!combined[0].contains_key("d1")); + assert!(!combined[0].contains_key("d2")); +} + +#[test] +fn a_plan_that_restricts_nothing_produces_an_empty_selection_map() { + // The no-pre-filter split: the plan carries a map with no entries at all, + // and that must survive the merge as an empty map (which the ANN layer reads + // as "nothing to mask"), not become a per-file all-permitting mask. + let physical = vec![HashMap::new()]; + let combined = intersect_row_allow_lists(Some(&physical), None, 1) + .unwrap() + .expect("a split-driven plan is always Some"); + assert!(combined[0].is_empty()); +} + +/// The batch terminals here are handed exactly one query, so a result vector of any +/// other length means the batch ran the wrong number of searches. The +/// `debug_assert_eq!` this replaced was compiled out of release builds, where an +/// empty vector panicked on `remove(0)` and a longer one silently returned another +/// query's result. +#[test] +fn take_only_result_rejects_bad_batch_arity() { + assert_eq!(take_only_result(vec![7], "test").unwrap(), 7); + assert!(take_only_result::(Vec::new(), "test").is_err()); + assert!(take_only_result(vec![1, 2], "test").is_err()); +} + +#[test] +fn rejects_allow_lists_that_do_not_cover_every_split() { + let physical = vec![range_allow_list(&[("d0", &[1])])]; + let error = intersect_row_allow_lists(Some(&physical), None, 2) + .map(|_| ()) + .expect_err("an allow-list per split is what makes the index meaningful"); + assert!(error.to_string().contains("for 2 splits"), "{error}"); + + let residual = vec![allow_list(&[("d0", &[1])])]; + let error = intersect_row_allow_lists(Some(&physical), Some(residual), 2) + .map(|_| ()) + .expect_err("the residual must cover every split too"); + assert!(error.to_string().contains("for 2 splits"), "{error}"); +} diff --git a/crates/paimon/src/table/pk_vector_read/tests.rs b/crates/paimon/src/table/pk_vector_read/tests.rs new file mode 100644 index 000000000..59df08f93 --- /dev/null +++ b/crates/paimon/src/table/pk_vector_read/tests.rs @@ -0,0 +1,1256 @@ +// 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::*; +use crate::catalog::Identifier; +use crate::io::{FileIO, FileIOBuilder}; +use crate::lumina::LuminaIndexMeta; +use crate::spec::stats::BinaryTableStats; +use crate::spec::{ + BinaryRow, DataField, DataFileMeta, DataType, Datum, FloatType, IntType, Predicate, + PredicateBuilder, Schema, TableSchema, +}; +use crate::table::data_file_reader::DataFileReader; +use crate::table::pk_vector_orchestrator::{PkVectorCandidate, PkVectorSearchSplit}; +use crate::table::source::DataSplitBuilder; +use crate::table::vector_scan::Scan; +use crate::table::vector_search_test_utils::{build_vindex_segment_bytes, pk_vector_table}; +use crate::table::{Table, TableCommit, TableWrite}; +use crate::vindex::pkvector::bucket::BucketAnnSegment; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::IVF_FLAT_IDENTIFIER; +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; +use arrow_array::{ArrayRef, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use paimon_vindex_core::index::VectorIndexReader as VIndexReader; +use std::collections::HashMap; +use std::io::Cursor; +use std::sync::{Arc, Mutex, Once}; + +const VECTOR_SEARCH_LOG_TARGET: &str = "paimon::vector_search"; + +static VECTOR_SEARCH_TEST_LOGGER: VectorSearchTestLogger = VectorSearchTestLogger; + +static VECTOR_SEARCH_TEST_LOGS: Mutex> = Mutex::new(Vec::new()); + +struct VectorSearchTestLogger; + +impl log::Log for VectorSearchTestLogger { + fn enabled(&self, metadata: &log::Metadata<'_>) -> bool { + metadata.target() == VECTOR_SEARCH_LOG_TARGET && metadata.level() <= log::Level::Debug + } + + fn log(&self, record: &log::Record<'_>) { + if self.enabled(record.metadata()) { + VECTOR_SEARCH_TEST_LOGS + .lock() + .unwrap() + .push(record.args().to_string()); + } + } + + fn flush(&self) {} +} + +fn reset_vector_search_test_logs() { + static INIT: Once = Once::new(); + INIT.call_once(|| log::set_logger(&VECTOR_SEARCH_TEST_LOGGER).unwrap()); + log::set_max_level(log::LevelFilter::Debug); + VECTOR_SEARCH_TEST_LOGS.lock().unwrap().clear(); +} + +fn pk_data_file(name: &str, row_count: i64, first_row_id: Option) -> DataFileMeta { + DataFileMeta { + file_name: name.to_string(), + file_size: 1, + row_count, + min_key: Vec::new(), + max_key: Vec::new(), + key_stats: BinaryTableStats::empty(), + value_stats: BinaryTableStats::empty(), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 1, + level: 0, + extra_files: Vec::new(), + creation_time: None, + delete_row_count: None, + embedded_index: None, + file_source: None, + value_stats_cols: None, + external_path: None, + first_row_id, + write_cols: None, + column_max_sequence_numbers: None, + } +} + +fn pk_search_split(bucket: i32, files: Vec) -> PkVectorSearchSplit { + PkVectorSearchSplit { + data_split: DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(bucket) + .with_bucket_path(format!("memory:/t/bucket-{bucket}")) + .with_total_buckets(1) + .with_data_files(files) + .build() + .unwrap(), + ann_segments: Vec::new(), + active_files: Vec::new(), + } +} + +fn pk_candidate( + split_index: usize, + bucket: i32, + file: &str, + pos: i64, + distance: f32, +) -> PkVectorCandidate { + PkVectorCandidate { + split_index, + partition: BinaryRow::new(0), + bucket, + data_file_name: file.to_string(), + row_position: pos, + distance, + } +} + +// Candidate with a fixed empty (arity-0) partition and bucket 0, keyed only by +// (split_index, file, position) — the dimensions the rerank core groups on. +fn cand_at(split_index: usize, file: &str, pos: i64, dist: f32) -> PkVectorCandidate { + pk_candidate(split_index, 0, file, pos, dist) +} + +/// The single data-file name every rerank fixture writes. +const RERANK_FILE: &str = "part-0.parquet"; + +/// Serialize a Paimon deletion-vector blob covering `deleted_rows` and write it +/// at `path`, returning the matching `DeletionFile`. Byte layout mirrors the +/// position-read tests: `[length][magic][roaring bitmap][0]`. +async fn write_deletion_blob( + file_io: &FileIO, + path: &str, + deleted_rows: &[u32], +) -> crate::table::source::DeletionFile { + use roaring::RoaringBitmap; + + const MAGIC_NUMBER: i32 = 1581511376; + let mut bitmap = RoaringBitmap::new(); + for row in deleted_rows { + bitmap.insert(*row); + } + let mut bitmap_bytes = Vec::new(); + bitmap.serialize_into(&mut bitmap_bytes).unwrap(); + let bitmap_length = 4 + bitmap_bytes.len() as i32; + let mut blob = Vec::new(); + blob.extend_from_slice(&bitmap_length.to_be_bytes()); + blob.extend_from_slice(&MAGIC_NUMBER.to_be_bytes()); + blob.extend_from_slice(&bitmap_bytes); + blob.extend_from_slice(&0i32.to_be_bytes()); + file_io + .new_output(path) + .unwrap() + .write(bytes::Bytes::from(blob)) + .await + .unwrap(); + crate::table::source::DeletionFile::new( + path.to_string(), + 0, + bitmap_length as i64, + Some(deleted_rows.len() as i64), + ) +} + +/// Write a single-file vector data file (`FixedSizeList` of width +/// `dim`) holding `rows` (a `None` entry is a NULL vector row) as Parquet, and +/// return a vector-only `DataFileReader`, the enclosing `PkVectorSearchSplit`, +/// and the vector `DataField`. When `deleted_rows` is non-empty a deletion +/// vector covering those physical positions is attached to the split, so the +/// position read drops them exactly as `PkVectorIndexedSplitRead::read` does. +/// +/// This is the position-only analogue of the old `ArrayReader`: rerank now +/// re-reads real stored rows through `PkVectorPositionRead`, so the fixtures +/// exercise that path rather than an in-memory preloaded column. +async fn vector_rerank_fixture( + table_path: &str, + dim: u32, + rows: &[Option>], + deleted_rows: &[u32], +) -> (DataFileReader, PkVectorSearchSplit, DataField) { + use crate::arrow::build_target_arrow_schema; + use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; + use crate::spec::VectorType; + use crate::table::schema_manager::SchemaManager; + + let vector_type = VectorType::try_new(true, dim, DataType::Float(FloatType::new())).unwrap(); + let vector_field = DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)); + let read_fields = vec![vector_field.clone()]; + let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); + + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), dim as i32).with_field( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + ); + for row in rows { + match row { + Some(values) => { + for v in values { + builder.values().append_value(*v); + } + builder.append(true); + } + None => { + for _ in 0..dim { + builder.values().append_value(0.0); + } + builder.append(false); + } + } + } + let vec_array = builder.finish(); + let batch = + arrow_array::RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]).unwrap(); + + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let bucket_path = format!("{table_path}/bucket-0"); + let output = file_io + .new_output(&format!("{bucket_path}/{RERANK_FILE}")) + .unwrap(); + let mut writer: Box = Box::new( + ParquetFormatWriter::new( + &output, + arrow_schema.clone(), + "zstd", + 1, + None, + &HashMap::new(), + ) + .await + .unwrap(), + ); + writer.write(&batch).await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; + + let schema_id = 1; + let file_meta = pk_data_file(RERANK_FILE, rows.len() as i64, Some(0)); + let file_meta = DataFileMeta { + file_size: file_size as i64, + schema_id, + ..file_meta + }; + + let mut split_builder = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![file_meta]); + if !deleted_rows.is_empty() { + let df = + write_deletion_blob(&file_io, &format!("{table_path}/index/dv-0"), deleted_rows).await; + split_builder = split_builder.with_data_deletion_files(vec![Some(df)]); + } + let data_split = split_builder.build().unwrap(); + let split = PkVectorSearchSplit { + data_split, + ann_segments: Vec::new(), + active_files: Vec::new(), + }; + + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let reader = DataFileReader::new( + file_io, + schema_manager, + schema_id, + read_fields.clone(), + read_fields, + Vec::new(), + ); + (reader, split, vector_field) +} + +fn pk_split_with_lumina_segment(path: &str, metric: &str) -> PkVectorSearchSplit { + let mut split = pk_search_split(0, vec![pk_data_file("file-a", 3, Some(0))]); + let source_meta = crate::spec::PrimaryKeyIndexSourceMeta::new( + 1, + vec![crate::spec::PrimaryKeyIndexSourceFile::new("file-a".to_string(), 3).unwrap()], + ) + .unwrap(); + let mut segment = BucketAnnSegment::for_test(source_meta); + segment.path = path.to_string(); + // Lumina stores its metric in the serialized index metadata blob, not in + // the segment file bytes. `deserialize` requires both keys present. + let meta = crate::lumina::LuminaIndexMeta::new(HashMap::from([ + ("index.dimension".to_string(), "2".to_string()), + ("distance.metric".to_string(), metric.to_string()), + ])); + segment.index_meta = meta.serialize().unwrap(); + split.ann_segments = vec![segment]; + split +} + +/// One Java `DataOutput#writeUTF` value (u16-BE length + modified UTF-8), used +/// to assemble the `PrimaryKeyIndexSourceMeta` frame below. +fn java_write_utf(s: &str) -> Vec { + let mut body = Vec::new(); + for c in s.encode_utf16() { + if (0x0001..=0x007F).contains(&c) { + body.push(c as u8); + } else if c > 0x07FF { + body.push(0xE0 | (c >> 12) as u8); + body.push(0x80 | ((c >> 6) & 0x3F) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } else { + body.push(0xC0 | (c >> 6) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } + } + let mut out = (body.len() as u16).to_be_bytes().to_vec(); + out.extend_from_slice(&body); + out +} + +/// The Java `PrimaryKeyIndexSourceMeta` frame: `i32-BE version=1`, `i32-BE +/// data_level`, `i32-BE count`, then per source file a `writeUTF` name and an +/// `i64-BE` row count. +fn pk_source_meta_bytes(data_level: i32, files: &[(&str, i64)]) -> Vec { + let mut out = Vec::new(); + out.extend_from_slice(&1i32.to_be_bytes()); + out.extend_from_slice(&data_level.to_be_bytes()); + out.extend_from_slice(&(files.len() as i32).to_be_bytes()); + for (name, rows) in files { + out.extend_from_slice(&java_write_utf(name)); + out.extend_from_slice(&rows.to_be_bytes()); + } + out +} + +/// Build a committed primary-key vector table (memory FS) over `vectors` +/// (dimension 2): write a real data file via the write path, promote its meta +/// to a compacted, non-level-0 file (the PK index-source precondition), then +/// build + commit a real vindex IVF-flat ANN segment naming that file. Single +/// bucket, `nlist = 1`, so the ANN search is exact. Returns the opened table, +/// ready for vector search. +async fn build_committed_pk_vector_table(vectors: &[[f32; 2]]) -> Table { + use crate::spec::{GlobalIndexMeta, IndexFileMeta, VectorType}; + use crate::table::CommitMessage; + use bytes::Bytes; + use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; + use paimon_vindex_core::io::PosWriter; + + const DIM: usize = 2; + let table_path = "memory:/pk_vector_route_test"; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Vector( + VectorType::try_new(true, DIM as u32, DataType::Float(FloatType::new())).unwrap(), + ), + ) + .primary_key(["id"]) + .option("bucket", "1") + .option("deletion-vectors.enabled", "true") + .option("pk-vector.index.columns", "embedding") + .option("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER) + .option("fields.embedding.pk-vector.distance.metric", "l2") + .build() + .unwrap(); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "pk_vector_route_test"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + for dir in ["snapshot", "manifest", "index"] { + file_io + .mkdirs(&format!("{table_path}/{dir}")) + .await + .unwrap(); + } + + // id + FixedSizeList batch matching the table's target schema. + let ids: Vec = (0..vectors.len() as i32).collect(); + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vec_builder = FixedSizeListBuilder::new(Float32Builder::new(), DIM as i32) + .with_field(element_field.clone()); + for v in vectors { + for &x in v { + vec_builder.values().append_value(x); + } + vec_builder.append(true); + } + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList(element_field, DIM as i32), + true, + ), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(vec_builder.finish()) as ArrayRef, + ], + ) + .unwrap(); + + // Real data-file meta via the write path (these messages are not committed + // as-is; the meta is promoted below and committed with the index). + let mut writer = TableWrite::new(&table, "route-test".to_string()).unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + let base = &messages[0]; + let base_meta = base.new_files[0].clone(); + let bucket = base.bucket; + let partition = base.partition.clone(); + let data_file_name = base_meta.file_name.clone(); + let row_count = base_meta.row_count; + + // PK index-source precondition: compacted, non-level-0, first_row_id pinned. + let indexed_meta = DataFileMeta { + level: 1, + file_source: Some(1), + first_row_id: Some(0), + ..base_meta + }; + + // Real vindex IVF-flat segment (nlist=1 -> exact) over the vectors. + let n = vectors.len(); + let flat: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); + let seg_ids: Vec = (0..n as i64).collect(); + let native = HashMap::from([ + ("index.type".to_string(), "ivf_flat".to_string()), + ("dimension".to_string(), DIM.to_string()), + ("nlist".to_string(), "1".to_string()), + ("metric".to_string(), "l2".to_string()), + ]); + let config = VectorIndexConfig::from_options(&native).unwrap(); + let training = VectorIndexTrainer::train(config, &flat, n).unwrap(); + let mut ann_writer = VectorIndexWriter::new(training); + ann_writer.add_vectors(&seg_ids, &flat, n).unwrap(); + let mut seg_bytes = Vec::new(); + { + let mut out = PosWriter::new(&mut seg_bytes); + ann_writer.write(&mut out).unwrap(); + } + let index_file_name = "vector-ivf-flat-route.index".to_string(); + let index_file_size = seg_bytes.len() as u64; + file_io + .new_output(&format!("{table_path}/index/{index_file_name}")) + .unwrap() + .write(Bytes::from(seg_bytes)) + .await + .unwrap(); + + let vector_field_id = schema + .fields() + .iter() + .find(|f| f.name() == "embedding") + .unwrap() + .id(); + let index_file = IndexFileMeta { + index_type: IVF_FLAT_IDENTIFIER.to_string(), + file_name: index_file_name, + file_size: i64::try_from(index_file_size).unwrap(), + row_count, + deletion_vectors_ranges: None, + external_path: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: 0, + row_range_end: row_count - 1, + index_field_id: vector_field_id, + extra_field_ids: None, + source_meta: Some(pk_source_meta_bytes(1, &[(&data_file_name, row_count)])), + index_meta: None, + }), + }; + + let mut message = CommitMessage::new(partition, bucket, vec![indexed_meta]); + message.new_index_files = vec![index_file]; + TableCommit::new(table.clone(), "route-test".to_string()) + .commit(vec![message]) + .await + .unwrap(); + table +} + +#[tokio::test] +async fn rerank_aligns_recomputed_distance_by_position_column() { + use crate::arrow::build_target_arrow_schema; + use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; + use crate::spec::VectorType; + use crate::table::schema_manager::SchemaManager; + + // A vector data file with 4 physical rows: positions 0,1,3 hold vectors + // and position 2 (a NON-candidate) holds a NULL vector. Candidates sit at + // non-contiguous positions {1, 3}. The ANN-reported distances are + // deliberately reversed relative to the true stored vectors; after rerank + // each candidate must carry compute_distance(query, vec_at_its_position), + // proving alignment is by the _PKEY_VECTOR_POSITION column value, not batch + // order. Position 2's NULL is never read (it is not a candidate), so it + // cannot trip the null-vector guard. + let vector_type = VectorType::try_new(true, 2, DataType::Float(FloatType::new())).unwrap(); + let vector_field = DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)); + let read_fields = vec![vector_field.clone()]; + let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); + + // pos0=[7,0], pos1=[1,0], pos2=NULL, pos3=[4,0]. + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new( + ArrowField::new("element", ArrowDataType::Float32, true), + )); + for row in [ + Some([7.0f32, 0.0]), + Some([1.0, 0.0]), + None, + Some([4.0, 0.0]), + ] { + match row { + Some([a, b]) => { + builder.values().append_value(a); + builder.values().append_value(b); + builder.append(true); + } + None => { + builder.values().append_value(0.0); + builder.values().append_value(0.0); + builder.append(false); + } + } + } + let vec_array = builder.finish(); + let batch = + arrow_array::RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]).unwrap(); + + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/rerank_positional"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_name = "part-0.parquet"; + let output = file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap(); + let mut writer: Box = Box::new( + ParquetFormatWriter::new( + &output, + arrow_schema.clone(), + "zstd", + 1, + None, + &HashMap::new(), + ) + .await + .unwrap(), + ); + writer.write(&batch).await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; + + let schema_id = 1; + let file_meta = pk_data_file(file_name, 4, Some(0)); + let file_meta = DataFileMeta { + file_size: file_size as i64, + schema_id, + ..file_meta + }; + let data_split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![file_meta]) + .build() + .unwrap(); + let split = PkVectorSearchSplit { + data_split, + ann_segments: Vec::new(), + active_files: Vec::new(), + }; + + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let reader = DataFileReader::new( + file_io, + schema_manager, + schema_id, + read_fields.clone(), + read_fields.clone(), + Vec::new(), + ); + + let query = vec![1.0f32, 0.0]; + // ANN-reported distances reversed vs. truth: pos1 reported worse (0.9) than + // pos3 (0.1), but the true L2 distances are pos1=0 and pos3=9. + let indexed = vec![cand_at(0, file_name, 1, 0.9), cand_at(0, file_name, 3, 0.1)]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .unwrap(); + + // Best-first after exact recompute: pos1 (d=0) then pos3 (d=9), each + // carrying the distance computed from its OWN position's stored vector. + assert_eq!(out.len(), 2); + assert_eq!(out[0].row_position, 1); + assert_eq!(out[0].distance, 0.0); + assert_eq!(out[1].row_position, 3); + assert_eq!(out[1].distance, 9.0); +} + +#[tokio::test] +async fn rerank_recomputes_distance_and_reorders() { + // pos0=[9,0], pos1=[1,0]; query=[1,0]. The ANN-reported distances are + // reversed relative to the truth (pos0 reported best at 0.1, pos1 worst at + // 0.9), so an implementation that trusted the ANN order would emit pos0 + // first. Exact L2 recompute yields pos0=64, pos1=0, so the output must + // reorder to pos1-then-pos0 with the recomputed distances. + let (reader, split, vector_field) = vector_rerank_fixture( + "memory:/rerank_reorder", + 2, + &[Some(vec![9.0, 0.0]), Some(vec![1.0, 0.0])], + &[], + ) + .await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 0, 0.1), + cand_at(0, RERANK_FILE, 1, 0.9), + ]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .unwrap(); + + assert_eq!(out.len(), 2); + assert_eq!(out[0].row_position, 1); + assert_eq!(out[0].distance, 0.0); + assert_eq!(out[1].row_position, 0); + assert_eq!(out[1].distance, 64.0); + // Order genuinely changed vs. the ANN-reported best-first (which was pos0). + assert!(out[0].distance < out[1].distance); +} + +#[tokio::test] +async fn rerank_is_independent_of_fast_mode_reranks_indexed() { + // The rerank core takes only the indexed (fast-path) candidates and always + // recomputes their true distance; there is no fast/exact switch that can + // skip it. The single candidate carries a bogus ANN distance (0.42) but its + // stored vector equals the query, so the recomputed L2 distance is exactly + // 0.0 — proving the indexed candidate WAS reranked rather than passed + // through with its ANN distance. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_indexed", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.42)]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .unwrap(); + + assert_eq!(out.len(), 1); + assert_eq!(out[0].row_position, 0); + assert_ne!(out[0].distance, 0.42); + assert_eq!(out[0].distance, 0.0); +} + +#[tokio::test] +async fn rerank_fails_loud_on_null_vector() { + // A NULL vector stored AT a candidate position must fail loud rather than + // silently scoring it: the candidate genuinely has no vector to rerank on. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_null", 2, &[None], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("null vector at a candidate position must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("null vector")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn rerank_fails_loud_on_leftover_candidate() { + // pos1 is deleted by the deletion vector, so the position read returns no + // row for it. The search path already DV-filters, so a deleted candidate + // reaching rerank is a real inconsistency: the leftover guard must fail + // loud rather than silently dropping the candidate. + let (reader, split, vector_field) = vector_rerank_fixture( + "memory:/rerank_leftover", + 2, + &[Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])], + &[1], + ) + .await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 0, 0.1), + cand_at(0, RERANK_FILE, 1, 0.9), + ]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .err() + .expect("a candidate returning no row must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("failed to read")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn rerank_fails_loud_on_dimension_mismatch() { + // Stored vectors are 3-dimensional but the query is 2-dimensional. The + // vector extraction validates each stored row against the query dimension + // and fails loud, so the recompute never runs against mismatched vectors. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_dim", 3, &[Some(vec![1.0, 0.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("dimension mismatch must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("dimension")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn rerank_fails_loud_on_duplicate_candidate_position() { + // Two candidates addressing the same (split_index, file, position) is a + // programming error upstream: the dedup guard fires before any read. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_dup", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 0, 0.1), + cand_at(0, RERANK_FILE, 0, 0.9), + ]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .err() + .expect("duplicate candidate position must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("duplicate")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn rerank_fails_loud_on_unexpected_position() { + // Every position the read surfaces must resolve to a candidate keyed by + // (split_index, file, position). Here the plan carries two splits for the + // SAME (partition, bucket, file), so `split_index_of` resolves the file to + // the LAST plan index (1). The single candidate is tagged with split_index + // 0, so its by_key entry is (0, file, 0) while the read looks up + // (1, file, 0). The lookup misses and the unexpected-position guard fires + // rather than silently dropping the surfaced row. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_unexpected", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + // Two plan entries for the same file: split_index_of ends up mapping the + // file to plan index 1, not the candidate's split_index 0. + let dup = PkVectorSearchSplit { + data_split: split.data_split.clone(), + ann_segments: Vec::new(), + active_files: Vec::new(), + }; + let plan = vec![dup, split]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &plan, + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("a read position absent from the candidate map must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("unexpected position")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn rerank_fails_loud_on_file_not_in_plan() { + // A candidate references a (partition, bucket, file) that is absent from + // plan_splits. build_indexed_splits groups it into an indexed split, but the + // split_index_of lookup — built only from plan_splits — has no entry, so the + // kernel fails loud rather than reading an unplanned file. + let (reader, _split, vector_field) = + vector_rerank_fixture("memory:/rerank_noplan", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + // Empty plan: the candidate's file resolves in no plan split. + let err = rerank_indexed_positional( + &reader, + indexed, + &[], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("a candidate file absent from the plan must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("not found in plan")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn rerank_reads_only_candidate_positions_not_whole_column() { + // A 6-row file where every NON-candidate position (0, 2, 4, 5) holds a NULL + // vector "poison" and only the two candidate positions (1, 3) hold real + // vectors. The rerank read is told to fetch only positions {1, 3}; every + // row it surfaces is looked up in the candidate map, and any position not in + // the map trips the "unexpected position" guard (a surfaced NULL row would + // additionally trip the null-vector guard). So if the read had surfaced any + // of the poison rows, rerank would fail. It succeeds and returns exactly the + // two candidates at positions {1, 3}, which proves the position selection + // reaching the read contained only the candidate positions (not the whole + // column). + let rows = &[ + None, // pos0 poison (non-candidate) + Some(vec![1.0, 0.0]), // pos1 candidate + None, // pos2 poison (non-candidate) + Some(vec![3.0, 0.0]), // pos3 candidate + None, // pos4 poison (non-candidate) + None, // pos5 poison (non-candidate) + ]; + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_spy", 2, rows, &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 1, 0.9), + cand_at(0, RERANK_FILE, 3, 0.1), + ]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .unwrap_or_else(|e| { + panic!("only candidate positions are read, so the poison NULLs never decode: {e:?}") + }); + + assert_eq!(out.len(), 2, "exactly the candidate count of rows was read"); + let mut positions: Vec = out.iter().map(|c| c.row_position).collect(); + positions.sort_unstable(); + assert_eq!( + positions, + vec![1, 3], + "only candidate positions reached the read" + ); + // Recomputed distances confirm each surviving row is its own candidate's vector. + assert_eq!(out[0].row_position, 1); + assert_eq!(out[0].distance, 0.0); + assert_eq!(out[1].row_position, 3); + assert_eq!(out[1].distance, 4.0); +} + +#[test] +fn verify_segment_metric_accepts_matching_lumina_metric() { + // Lumina segment metadata says cosine; configured cosine => Ok. No segment + // file bytes are needed on the Lumina path. + let split = pk_split_with_lumina_segment("seg-lumina", "cosine"); + let segment = &split.ann_segments[0]; + let lumina_metric = LuminaIndexMeta::deserialize(&segment.index_meta) + .unwrap() + .metric() + .unwrap(); + verify_segment_metric( + VectorSearchMetric::Cosine, + VectorSearchMetric::from_lumina(lumina_metric), + ) + .expect("matching lumina metric must pass"); +} + +#[test] +fn verify_segment_metric_rejects_mismatched_lumina_metric() { + // Lumina segment metadata says l2; configured inner_product => fail loud, + // naming both metrics. + let split = pk_split_with_lumina_segment("seg-lumina", "l2"); + let segment = &split.ann_segments[0]; + let lumina_metric = LuminaIndexMeta::deserialize(&segment.index_meta) + .unwrap() + .metric() + .unwrap(); + let err = verify_segment_metric( + VectorSearchMetric::InnerProduct, + VectorSearchMetric::from_lumina(lumina_metric), + ) + .expect_err("mismatched lumina metric must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("does not match configured metric") + && message.contains("l2") + && message.contains("inner_product")), + "unexpected error: {err:?}" + ); +} + +#[test] +fn verify_segment_metric_accepts_matching_vindex_metric() { + // Real IVF segment trained with L2; configured metric L2 => Ok. + let bytes = bytes::Bytes::from(build_vindex_segment_bytes("l2")); + let reader = VIndexReader::open(Cursor::new(bytes)).unwrap(); + verify_segment_metric( + VectorSearchMetric::L2, + VectorSearchMetric::from_vindex(reader.metadata().metric), + ) + .expect("matching metric must pass"); +} + +#[test] +fn verify_segment_metric_rejects_mismatched_vindex_metric() { + // Real IVF segment trained with L2; configured metric Cosine => fail loud. + let bytes = bytes::Bytes::from(build_vindex_segment_bytes("l2")); + let reader = VIndexReader::open(Cursor::new(bytes)).unwrap(); + let err = verify_segment_metric( + VectorSearchMetric::Cosine, + VectorSearchMetric::from_vindex(reader.metadata().metric), + ) + .expect_err("mismatched metric must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("does not match configured metric") + && message.contains("l2") + && message.contains("cosine")), + "unexpected error: {err:?}" + ); +} + +// ---- Search results retain scored positions and the planned source context. ---- +#[tokio::test] +async fn reader_uses_planned_splits_after_index_manifest_is_removed() { + let table = build_committed_pk_vector_table(&[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]).await; + let options = HashMap::new(); + let query = [0.0, 1.0]; + let params = + PkVectorSearchParams::resolve(&table, &options, None, "embedding", &[&query], 2).unwrap(); + let scan = crate::table::pk_vector_scan::PkVectorScan::new( + &table, + params.vector_field.id(), + params.index_type.clone(), + None, + ); + let read = PkVectorRead::new(&table, &options, None, "embedding", &[&query], 2, params); + let plan = scan.plan().await.unwrap(); + let snapshot_id = plan.snapshot_id; + let manager = table.snapshot_manager(); + let snapshot = manager.get_snapshot(snapshot_id).await.unwrap(); + let manifest_path = manager.manifest_path(snapshot.index_manifest().unwrap()); + table.file_io().delete_file(&manifest_path).await.unwrap(); + assert!( + scan.plan().await.is_err(), + "a second plan must need the removed manifest" + ); + + // Reading must use the original per-bucket source context, without replanning. + let mut results = read.read(plan).await.unwrap(); + assert_eq!(results.len(), 1); + let result = results.pop().unwrap(); + assert_eq!(result.snapshot_id(), Some(snapshot_id)); + assert_eq!( + result + .positions() + .unwrap() + .iter() + .map(|c| c.row_position) + .collect::>(), + vec![1, 2] + ); + assert_eq!( + result + .positions() + .unwrap() + .iter() + .map(|c| c.score) + .collect::>(), + vec![1.0, 0.5] + ); + assert_eq!(result.indexed_splits().unwrap().len(), 1); + + let batches: Vec = result + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let ids: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!( + ids, + vec![1, 2], + "materialization must also use the retained plan" + ); +} + +#[tokio::test] +async fn execute_returns_scored_positions_and_publishes_diagnostics() { + reset_vector_search_test_logs(); + let _timing = crate::vindex::enable_vector_search_timing_for_test(); + // query [0,1]: squared-L2 distances pos1=0 < pos2=1 < pos0=2, so the + // strict-gap top-2 is [pos1, pos2] (best-first, not physical order). + let table = build_committed_pk_vector_table(&[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]).await; + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![0.0, 1.0]) + .with_limit(2) + .execute() + .await + .unwrap(); + let positions = result.positions().unwrap(); + + // Two nearest neighbours recalled, best-first, without materialization. + assert_eq!(positions.len(), 2, "top-2 positions expected"); + assert_eq!(positions[0].row_position, 1, "nearest is position 1"); + assert_eq!(positions[1].row_position, 2, "second nearest is position 2"); + assert_eq!( + positions.iter().map(|p| p.score).collect::>(), + vec![1.0, 0.5] + ); + + // Every position retains the source file and snapshot needed for a later read. + assert_eq!(result.snapshot_id(), Some(1), "first commit -> snapshot 1"); + let splits = result.indexed_splits().unwrap(); + assert_eq!(splits.len(), 1); + assert_eq!(splits[0].split.snapshot_id(), 1); + assert!( + positions.iter().all(|p| { + p.data_file_name == splits[0].split.data_files()[0].file_name + && p.bucket == splits[0].split.bucket() + && p.partition.to_serialized_bytes() + == splits[0].split.partition().to_serialized_bytes() + }), + "positions must refer to their retained source split" + ); + + let logs = VECTOR_SEARCH_TEST_LOGS.lock().unwrap(); + assert!( + logs.iter().any(|entry| { + entry.contains("event=paimon_vindex_reader") + && entry.contains("vector-ivf-flat-route.index") + }), + "PK vector search must publish vindex reader timing" + ); + assert!( + logs.iter().any(|entry| { + entry.contains("event=paimon_vector_range_io") + && entry.contains("vector-ivf-flat-route.index") + }), + "PK vector search must publish range-I/O timing" + ); +} + +/// A table with no snapshot yields empty positions and splits without inventing +/// a snapshot or changing the result's PK address space. +#[tokio::test] +async fn execute_empty_plan_yields_empty_pk_result() { + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 128]) + .with_limit(3) + .execute() + .await + .unwrap(); + assert!( + result.positions().unwrap().is_empty(), + "no data -> no positions" + ); + assert!( + result.indexed_splits().unwrap().is_empty(), + "no data -> no source splits" + ); + assert_eq!(result.snapshot_id(), None); + assert!(result.row_ids().is_err()); +} + +/// The vector residual is derived from the DATA conjuncts of the filter: +/// partition-only conjuncts are enforced by scan planning (`PkVectorScan` +/// pushes the whole filter through the normal scan) and must not enter the +/// per-row residual, so a partition-only filter yields no residual at all. +#[test] +fn residual_uses_only_data_conjuncts_of_the_filter() { + use crate::spec::VarCharType; + use crate::table::bucket_filter::split_partition_and_data_predicates; + + // Partitioned table: `dt` (partition key) + `id`. + let schema = Schema::builder() + .column("dt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .partition_keys(["dt"]) + .build() + .unwrap(); + let ts = TableSchema::new(0, &schema); + let fields = ts.fields(); + let partition_keys = ts.partition_keys(); + let pb = PredicateBuilder::new(fields); + + // Partition-only `dt = 'a'` -> no residual data predicate (residual skipped; + // the partition is enforced by planning alone). + let (_p, data) = split_partition_and_data_predicates( + pb.equal("dt", Datum::String("a".to_string())).unwrap(), + fields, + partition_keys, + ); + assert!( + data.is_empty(), + "partition-only filter must leave no residual data predicate" + ); + + // Data-only `id > 5` -> kept as the residual. + let (_p, data) = split_partition_and_data_predicates( + pb.greater_than("id", Datum::Int(5)).unwrap(), + fields, + partition_keys, + ); + assert_eq!(data.len(), 1, "data-only filter must remain the residual"); + + // `dt = 'a' AND id > 5` -> only the data conjunct enters the residual. + let (_p, data) = split_partition_and_data_predicates( + Predicate::and(vec![ + pb.equal("dt", Datum::String("a".to_string())).unwrap(), + pb.greater_than("id", Datum::Int(5)).unwrap(), + ]), + fields, + partition_keys, + ); + assert_eq!( + data.len(), + 1, + "AND(partition, data) residual must drop the partition conjunct" + ); + + // `dt = 'a' OR id > 5` is a single mixed conjunct: it is NOT partition-only, + // so it stays whole in the residual (evaluated against the materialized + // partition column), rather than being dropped or split. + let mixed = Predicate::or(vec![ + pb.equal("dt", Datum::String("a".to_string())).unwrap(), + pb.greater_than("id", Datum::Int(5)).unwrap(), + ]); + let (_p, data) = split_partition_and_data_predicates(mixed.clone(), fields, partition_keys); + assert_eq!( + data, + vec![mixed], + "a mixed partition/data conjunct must stay whole in the residual" + ); +} diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 39815ee5f..e9fae268e 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -34,6 +34,7 @@ use crate::table::partition_filter::PartitionFilter; use crate::table::pk_vector_bucket_split::BucketVectorSearchSplit; use crate::table::pk_vector_orchestrator::PkVectorSearchSplit; use crate::table::source::{merge_row_ranges, DataSplit, DataSplitBuilder, DeletionFile, RowRange}; +use crate::table::vector_scan::Scan; use crate::table::Table; use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; @@ -207,6 +208,7 @@ impl BucketAccumulator { } /// The per-bucket search splits produced by planning. +#[derive(Clone)] pub(crate) struct PkVectorScanPlan { // The snapshot the plan resolved during planning (pinned before the index // manifest is read). It is authoritative even when planning yields zero @@ -230,29 +232,75 @@ pub(crate) struct PkVectorScanPlan { pub physical_row_ranges_by_split: Option>>>, } -pub(crate) struct PkVectorScan<'a> { - table: &'a Table, +pub(crate) struct PkVectorScan { + table: Table, vector_field_id: i32, index_type: String, filter: Option, } -impl<'a> PkVectorScan<'a> { +impl PkVectorScan { pub(crate) fn new( - table: &'a Table, + table: &Table, vector_field_id: i32, index_type: String, filter: Option, ) -> Self { Self { - table, + table: table.clone(), vector_field_id, index_type, filter, } } - pub(crate) async fn plan(&self) -> crate::Result { + /// Build a plan from bucket splits an engine planned elsewhere, instead of from + /// this table's index manifest. + /// + /// The splits are the planning input and are taken as authoritative: their + /// payload files, their per-file row ranges, and the snapshot they pin are used + /// as given, and no index manifest is read. Only the partition conjuncts of this + /// scan's filter are re-applied, because a caller may narrow the query further + /// than the planner that produced the splits. + /// + /// Mirrors what Java's `PrimaryKeyVectorRead` does with a + /// `BucketVectorSearchSplit`: search the payloads the split names, over the rows + /// the split allows. + pub(crate) fn plan_for_bucket_vector_splits( + &self, + splits: Vec, + ) -> crate::Result { + // Partition conjuncts only. Data conjuncts stay a per-row residual applied + // during the search: pruning a whole bucket on them would drop rows that + // still match. + let partition_filter = self.filter.as_ref().and_then(|filter| { + let (partition_predicate, _data_predicates) = split_partition_and_data_predicates( + filter.clone(), + self.table.schema().fields(), + self.table.schema().partition_keys(), + ); + partition_predicate.map(|predicate| { + PartitionFilter::from_predicate(predicate, &self.table.schema().partition_fields()) + }) + }); + plan_from_bucket_splits( + &self.index_type, + self.vector_field_id, + partition_filter.as_ref(), + self.table.location().trim_end_matches('/'), + self.table + .schema() + .core_options() + .index_file_in_data_file_dir(), + splits, + ) + } +} + +impl Scan for PkVectorScan { + type Plan = PkVectorScanPlan; + + async fn plan(&self) -> crate::Result { let snapshot_manager = self.table.snapshot_manager(); // Data splits first, via the table's own scan resolution (which honors @@ -370,48 +418,6 @@ impl<'a> PkVectorScan<'a> { physical_row_ranges_by_split: None, }) } - - /// Build a plan from bucket splits an engine planned elsewhere, instead of from - /// this table's index manifest. - /// - /// The splits are the planning input and are taken as authoritative: their - /// payload files, their per-file row ranges, and the snapshot they pin are used - /// as given, and no index manifest is read. Only the partition conjuncts of this - /// scan's filter are re-applied, because a caller may narrow the query further - /// than the planner that produced the splits. - /// - /// Mirrors what Java's `PrimaryKeyVectorRead` does with a - /// `BucketVectorSearchSplit`: search the payloads the split names, over the rows - /// the split allows. - pub(crate) fn plan_for_bucket_vector_splits( - &self, - splits: Vec, - ) -> crate::Result { - // Partition conjuncts only. Data conjuncts stay a per-row residual applied - // during the search: pruning a whole bucket on them would drop rows that - // still match. - let partition_filter = self.filter.as_ref().and_then(|filter| { - let (partition_predicate, _data_predicates) = split_partition_and_data_predicates( - filter.clone(), - self.table.schema().fields(), - self.table.schema().partition_keys(), - ); - partition_predicate.map(|predicate| { - PartitionFilter::from_predicate(predicate, &self.table.schema().partition_fields()) - }) - }); - plan_from_bucket_splits( - &self.index_type, - self.vector_field_id, - partition_filter.as_ref(), - self.table.location().trim_end_matches('/'), - self.table - .schema() - .core_options() - .index_file_in_data_file_dir(), - splits, - ) - } } /// The `Table`-independent core of [`PkVectorScan::plan_for_bucket_vector_splits`], diff --git a/crates/paimon/src/table/pk_vector_search_params.rs b/crates/paimon/src/table/pk_vector_search_params.rs new file mode 100644 index 000000000..954c02851 --- /dev/null +++ b/crates/paimon/src/table/pk_vector_search_params.rs @@ -0,0 +1,218 @@ +// 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. + +//! Resolve and validate PK vector-search parameters before scan planning. + +use crate::lumina::{is_lumina_index_type, LuminaVectorIndexOptions}; +use crate::spec::{CoreOptions, DataField, DataType, GlobalIndexSearchMode, Predicate}; +use crate::table::bucket_filter::split_partition_and_data_predicates; +use crate::table::vector_search_common::{configured_refine_factor, indexed_search_limit}; +use crate::table::Table; +use crate::vindex::pkvector::exact::validate_query; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::VindexVectorIndexOptions; +use std::collections::HashMap; + +/// Query-level parameters for a primary-key vector search: everything resolvable +/// from the table schema, the options and the queries alone, independent of which +/// splits planning yields. Resolved before planning so a malformed query or option +/// fails loud even when the plan turns out empty. +pub(super) struct PkVectorSearchParams { + pub(super) metric: VectorSearchMetric, + /// Fan-out limit for bucket orchestration plus ANN and exact-file leaves (Java + /// `GLOBAL_INDEX_THREAD_NUM`); `1` reproduces strictly sequential execution. + pub(super) concurrency: usize, + pub(super) index_type: String, + pub(super) vector_field: DataField, + pub(super) skip_exact_fallback: bool, + pub(super) refine_factor: usize, + pub(super) indexed_limit: usize, +} + +impl PkVectorSearchParams { + /// Resolve the query-level parameters and reject a query the search cannot answer + /// correctly, before any planning or read happens. + pub(super) fn resolve( + table: &Table, + query_options: &HashMap, + filter: Option<&Predicate>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, + ) -> crate::Result { + let core = CoreOptions::new(table.schema().options()); + core.ensure_read_authorized()?; + // Residual pre-filter guard, mirroring Java `PrimaryKeyVectorScan`. A DATA + // predicate set via `with_filter` is applied post-recall by re-reading each + // candidate file's physical rows during search. That physical-position filtering + // only agrees with the bucket search when the table exposes physical rows + // directly: deletion vectors enabled and merge-on-read disabled. Under + // merge-on-read (or without deletion vectors) a read merges multiple key + // versions, so a scalar filter could retain a stale version whose live version + // does not match — a silent wrong-read. Reject such queries rather than answer + // them incorrectly. + // + // Guard on the DATA conjuncts, not the whole filter: partition-only conjuncts + // are enforced entirely by scan planning (partition pruning) and produce no + // per-row residual, so they need no physical-row read. This mirrors Java, where + // `BatchVectorSearchBuilderImpl.withFilter` splits at the builder level and + // leaves `this.filter == null` for a partition-only filter — the scan guard is + // then skipped. No data predicate (partition-only or no filter) → nothing to + // guard, so the search-only and read paths are unaffected. + let physical_row_read = + core.deletion_vectors_enabled() && !core.deletion_vectors_merge_on_read(); + let has_data_predicate = filter.is_some_and(|f| { + let (_partition, data) = split_partition_and_data_predicates( + f.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ); + !data.is_empty() + }); + if has_data_predicate && !physical_row_read { + return Err(crate::Error::DataInvalid { + message: + "primary-key vector pre-filter requires deletion vectors without merge-on-read" + .to_string(), + source: None, + }); + } + // `primary_key_vector_distance_metric` returns a validated name; re-parse into + // the enum for the numeric semantics. + let metric = VectorSearchMetric::parse(&core.primary_key_vector_distance_metric(pk_col)?)?; + // Fan-out limit for bucket orchestration plus ANN and exact-file leaves (Java + // `GLOBAL_INDEX_THREAD_NUM`); `1` reproduces strictly sequential execution. + let concurrency = core.global_index_thread_num()?; + let index_type = core.primary_key_vector_index_type(pk_col)?; + let vector_field = table + .schema() + .fields() + .iter() + .find(|f| f.name() == pk_col) + .cloned() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("PK-vector column '{pk_col}' not found in schema"), + source: None, + })?; + + let search_mode = core.vector_index_search_mode()?; + let skip_exact_fallback = search_mode == GlobalIndexSearchMode::Fast; + + // A non-positive limit is invalid regardless of the plan; reject it before + // planning so an empty plan cannot mask it with empty results. + if limit == 0 { + return Err(crate::Error::DataInvalid { + message: "vector search limit must be positive".to_string(), + source: None, + }); + } + + // Resolve the refine factor from the query options first, then fall back to + // the table options; a positive factor over-fetches indexed (approximate) + // candidates so the reader's exact rerank has a wider pool to reorder. Factor 0 + // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank + // path. The two option maps are kept distinct (query options passed + // separately from table options) so a broad query key cannot be overridden + // by a more specific table key: query options take precedence as a whole. + // Resolved before planning so an invalid factor (e.g. a non-numeric value) + // fails loud regardless of whether the table currently has searchable data. + let refine_factor = + configured_refine_factor(query_options, table.schema().options(), pk_col, &index_type)?; + let indexed_limit = indexed_search_limit(limit, refine_factor)?; + + // Validate every query against the vector column's dimension (and finiteness) + // before planning or any read, so a malformed query fails loud even when the + // plan turns out empty. VECTOR carries the dimension in its type; + // ARRAY gets the index dimension from the same vindex option resolver + // used by index reads. Both valid PK-vector column shapes must reject NaN/Inf + // up front, not only after a non-empty plan opens readers. + if let Some(dimension) = pk_vector_query_dimension( + table.schema().options(), + query_options, + &index_type, + &vector_field, + )? { + for query in queries { + validate_query(query, dimension)?; + } + } + + Ok(Self { + metric, + concurrency, + index_type, + vector_field, + skip_exact_fallback, + refine_factor, + indexed_limit, + }) + } +} + +fn pk_vector_query_dimension( + table_options: &HashMap, + query_options: &HashMap, + index_type: &str, + vector_field: &DataField, +) -> crate::Result> { + match vector_field.data_type() { + DataType::Vector(vector_type) + if matches!(vector_type.element_type(), DataType::Float(_)) => + { + Ok(Some(vector_type.length() as usize)) + } + DataType::Array(array_type) if matches!(array_type.element_type(), DataType::Float(_)) => { + // Resolve the dimension per the configured backend. An `ARRAY` + // column carries no dimension in its type, so it comes from options — + // but the option shape differs by backend. Lumina is not a vindex + // index type, so routing it through `VindexVectorIndexOptions` would + // reject it as unsupported before planning (even on an empty table). + if is_lumina_index_type(index_type) { + // Lumina reads `lumina.index.dimension` (default 128) from the + // merged table+query options, matching `resolve_lumina_options`. + let mut merged = table_options.clone(); + merged.extend(query_options.clone()); + let dimension = LuminaVectorIndexOptions::new(&merged)?.dimension; + Ok(Some(dimension as usize)) + } else { + let mut dimension_options = HashMap::new(); + for key in [ + "dimension".to_string(), + format!("{index_type}.dimension"), + format!("fields.{}.dimension", vector_field.name()), + ] { + if let Some(value) = query_options.get(&key) { + dimension_options.insert(key, value.clone()); + } + } + Ok(Some( + VindexVectorIndexOptions::new( + table_options, + &dimension_options, + index_type, + vector_field, + )? + .dimension(), + )) + } + } + _ => Ok(None), + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/pk_vector_search_params/tests.rs b/crates/paimon/src/table/pk_vector_search_params/tests.rs new file mode 100644 index 000000000..061c09d4b --- /dev/null +++ b/crates/paimon/src/table/pk_vector_search_params/tests.rs @@ -0,0 +1,41 @@ +// 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::*; +use crate::spec::{ArrayType, FloatType}; + +#[test] +fn vindex_array_dimension_accepts_diskann_search_options() { + let field = DataField::new( + 1, + "embedding".to_string(), + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + let query_options = HashMap::from([ + ("diskann.dimension".to_string(), "8".to_string()), + ("diskann.l_search".to_string(), "64".to_string()), + ( + "vindex.reader.memory-budget-bytes".to_string(), + "1048576".to_string(), + ), + ]); + + assert_eq!( + pk_vector_query_dimension(&HashMap::new(), &query_options, "diskann", &field).unwrap(), + Some(8) + ); +} diff --git a/crates/paimon/src/table/vector_read.rs b/crates/paimon/src/table/vector_read.rs new file mode 100644 index 000000000..a1c166c3d --- /dev/null +++ b/crates/paimon/src/table/vector_read.rs @@ -0,0 +1,117 @@ +// 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. + +//! Vector index execution contract and the shared scan -> plan -> read pipeline. + +use crate::spec::{CoreOptions, Predicate}; +use crate::table::de_vector_read::DeVectorRead; +use crate::table::de_vector_scan::PreparedVectorSearchFilter; +use crate::table::pk_vector_read::PkVectorRead; +use crate::table::pk_vector_search_params::PkVectorSearchParams; +use crate::table::vector_scan::{PlanContext, VectorScanPlan, VectorScanWork}; +use crate::table::vector_search_common::{take_only_result, targets_primary_key_column}; +use crate::table::Table; +use crate::vector_search::SearchResult; +use roaring::RoaringTreemap; +use std::collections::HashMap; +use std::future::Future; +use std::sync::Arc; + +/// Execute a resolved plan without resolving another snapshot or manifest. +/// +/// Both implementations return snapshot-scoped search results. The associated +/// plan type prevents passing a DE plan to a PK reader or vice versa. +pub(super) trait Read: Sync { + type Plan; + + fn read( + &self, + plan: Self::Plan, + ) -> impl Future>> + Send; +} + +/// Searches a common plan with one query, without replanning or materializing rows. +pub struct VectorRead { + pub(super) batch: BatchVectorRead, +} + +impl VectorRead { + pub async fn read(&self, plan: VectorScanPlan) -> crate::Result { + take_only_result(self.batch.read(plan).await?, "vector search") + } +} + +/// Searches a common plan with multiple queries, preserving input order and arity. +/// Owns its configuration so it can outlive the builder that created it. +pub struct BatchVectorRead { + context: PlanContext, + reader: VectorReadKind, +} + +enum VectorReadKind { + DataEvolution(DeVectorRead), + PrimaryKey(PkVectorRead), +} + +impl BatchVectorRead { + pub(super) fn new( + table: &Table, + column: &str, + queries: &[&[f32]], + limit: usize, + options: &HashMap, + filter: Option<&Predicate>, + include_row_ids: Option<&Arc>, + prepared: Option<&PreparedVectorSearchFilter>, + ) -> crate::Result { + let context = PlanContext::new(table, column, filter, include_row_ids, prepared)?; + let core = CoreOptions::new(table.schema().options()); + let reader = if targets_primary_key_column(&core, column) { + let pk_col = core.primary_key_vector_index_column()?; + let params = + PkVectorSearchParams::resolve(table, options, filter, &pk_col, queries, limit)?; + VectorReadKind::PrimaryKey(PkVectorRead::new( + table, options, filter, &pk_col, queries, limit, params, + )) + } else { + VectorReadKind::DataEvolution(DeVectorRead::new(column, queries, limit, options)?) + }; + Ok(Self { context, reader }) + } + + pub async fn read(&self, plan: VectorScanPlan) -> crate::Result> { + if self.context != plan.context { + return Err(crate::Error::DataInvalid { + message: "vector plan and reader must use the same table, column and pre-filter" + .to_string(), + source: None, + }); + } + match (&self.reader, plan.work) { + (VectorReadKind::DataEvolution(reader), VectorScanWork::DataEvolution(plan)) => { + reader.read(plan).await + } + (VectorReadKind::PrimaryKey(reader), VectorScanWork::PrimaryKey(plan)) => { + reader.read(plan).await + } + _ => Err(crate::Error::DataInvalid { + message: "vector plan and reader use different DE/PK address spaces".to_string(), + source: None, + }), + } + } +} diff --git a/crates/paimon/src/table/vector_scan.rs b/crates/paimon/src/table/vector_scan.rs new file mode 100644 index 000000000..a1190e004 --- /dev/null +++ b/crates/paimon/src/table/vector_scan.rs @@ -0,0 +1,201 @@ +// 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. + +//! Vector search planning contract, shared by DE and primary-key scans. + +use std::future::Future; +use std::sync::Arc; + +use roaring::RoaringTreemap; + +use crate::spec::{CoreOptions, Predicate}; +use crate::table::de_vector_scan::{DeVectorScan, DeVectorScanPlan, PreparedVectorSearchFilter}; +use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan}; +use crate::table::vector_search_common::targets_primary_key_column; +use crate::table::{find_field_id_by_name, BucketVectorSearchSplit, Table}; + +/// Resolve one snapshot and the work a vector reader will consume. +/// +/// The associated plan keeps global row IDs and physical bucket positions in +/// separate types. A reader cannot accidentally consume another route's plan. +pub(super) trait Scan: Sync { + type Plan; + + fn plan(&self) -> impl Future> + Send; +} + +/// Snapshot and search work shared by the single-query and batch readers. +/// Plans own their source context and remain usable after the scan is dropped. +#[derive(Clone)] +pub struct VectorScanPlan { + pub(super) context: PlanContext, + pub(super) work: VectorScanWork, +} + +#[derive(Clone)] +pub(super) enum VectorScanWork { + DataEvolution(DeVectorScanPlan), + PrimaryKey(PkVectorScanPlan), +} + +impl VectorScanPlan { + pub fn snapshot_id(&self) -> Option { + match &self.work { + VectorScanWork::DataEvolution(plan) => plan.table.travel_snapshot().map(|s| s.id()), + VectorScanWork::PrimaryKey(plan) => (plan.snapshot_id != 0).then_some(plan.snapshot_id), + } + } +} + +/// A plan must use the reader's table, column and pre-filter. Query vectors, +/// limits and index search options may differ when a plan is reused. +#[derive(Clone, PartialEq)] +pub(super) struct PlanContext { + location: String, + branch: String, + column: String, + filter: Option, + include_row_ids: Option>, +} + +impl PlanContext { + pub(super) fn new( + table: &Table, + column: &str, + filter: Option<&Predicate>, + include_row_ids: Option<&Arc>, + prepared: Option<&PreparedVectorSearchFilter>, + ) -> crate::Result { + let core = CoreOptions::new(table.schema().options()); + core.ensure_read_authorized()?; + if let Some(prepared) = prepared { + if table.location().trim_end_matches('/') + != prepared.table().location().trim_end_matches('/') + || table.branch() != prepared.table().branch() + { + return Err(crate::Error::DataInvalid { + message: "Prepared vector search filter belongs to a different table" + .to_string(), + source: None, + }); + } + CoreOptions::new(prepared.table().schema().options()).ensure_read_authorized()?; + } + if column.is_empty() { + return Err(crate::Error::ConfigInvalid { + message: "Vector column must be set via with_vector_column()".to_string(), + }); + } + if targets_primary_key_column(&core, column) + && (include_row_ids.is_some() || prepared.is_some()) + { + return Err(crate::Error::DataInvalid { + message: "global row-ID filters cannot be applied to primary-key file positions; use with_filter()".to_string(), source: None, + }); + } + Ok(Self { + location: table.location().trim_end_matches('/').to_string(), + branch: table.branch().to_string(), + column: column.to_string(), + filter: filter.cloned(), + include_row_ids: prepared + .map(PreparedVectorSearchFilter::include_row_ids) + .or(include_row_ids) + .cloned(), + }) + } +} + +/// Creates query-independent plans for DE or primary-key vector search. +pub struct VectorScan { + context: PlanContext, + scan: VectorScanKind, +} + +enum VectorScanKind { + DataEvolution(DeVectorScan), + PrimaryKey(PkVectorScan), +} + +impl VectorScan { + pub(super) fn new( + table: &Table, + column: &str, + filter: Option<&Predicate>, + include_row_ids: Option<&Arc>, + prepared: Option<&PreparedVectorSearchFilter>, + ) -> crate::Result { + let context = PlanContext::new(table, column, filter, include_row_ids, prepared)?; + let core = CoreOptions::new(table.schema().options()); + let scan = if targets_primary_key_column(&core, column) { + let column = core.primary_key_vector_index_column()?; + let field_id = + find_field_id_by_name(table.schema().fields(), &column).ok_or_else(|| { + crate::Error::DataInvalid { + message: format!("PK-vector column '{column}' not found in schema"), + source: None, + } + })?; + VectorScanKind::PrimaryKey(PkVectorScan::new( + table, + field_id, + core.primary_key_vector_index_type(&column)?, + filter.cloned(), + )) + } else { + VectorScanKind::DataEvolution(DeVectorScan::new( + table, + filter, + include_row_ids, + prepared, + )) + }; + Ok(Self { context, scan }) + } + + pub async fn plan(&self) -> crate::Result { + let work = match &self.scan { + VectorScanKind::DataEvolution(scan) => { + VectorScanWork::DataEvolution(scan.plan().await?) + } + VectorScanKind::PrimaryKey(scan) => VectorScanWork::PrimaryKey(scan.plan().await?), + }; + Ok(VectorScanPlan { + context: self.context.clone(), + work, + }) + } + + /// Adapt already-decoded Java PK bucket work into a common read plan. + /// The supplied snapshot, files and physical row ranges are authoritative; + /// this does not consult the table's snapshot or index manifests. + pub fn plan_from_bucket_splits( + &self, + splits: Vec, + ) -> crate::Result { + let VectorScanKind::PrimaryKey(scan) = &self.scan else { + return Err(crate::Error::DataInvalid { + message: "bucket splits require a primary-key vector scan".to_string(), + source: None, + }); + }; + Ok(VectorScanPlan { + context: self.context.clone(), + work: VectorScanWork::PrimaryKey(scan.plan_for_bucket_vector_splits(splits)?), + }) + } +} diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 9e2273968..f0c742f99 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -15,171 +15,14 @@ // specific language governing permissions and limitations // under the License. -use crate::arrow::format::FilePredicates; -use crate::arrow::residual::{evaluate_predicates_mask, widen_scan_fields}; -use crate::io::{FileIO, FileRead}; -use crate::lumina::reader::LuminaVectorGlobalIndexReader; -use crate::lumina::{ - is_lumina_index_type, LuminaIndexMeta, LuminaVectorIndexOptions, LuminaVectorMetric, -}; -use crate::spec::{ - row_id_data_field, CoreOptions, DataField, DataType, FileKind, GlobalIndexSearchMode, - IndexFileMeta, IndexManifest, IndexManifestEntry, Predicate, ROW_ID_FIELD_NAME, -}; -use crate::table::bucket_filter::split_partition_and_data_predicates; -use crate::table::data_file_reader::DataFileReader; -use crate::table::global_index_scanner::{ - deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, - unindexed_ranges_for_global_index_entries, RowRangeIndex, -}; -use crate::table::index_file_path::IndexFileLocation; -use crate::table::pk_vector_bucket_split::BucketVectorSearchSplit; -use crate::table::pk_vector_data_file_reader::{ - append_batch_vectors, DataFilePkVectorReaderFactory, -}; -use crate::table::pk_vector_indexed_split_read::{expand_ranges, PkVectorIndexedSplitRead}; -use crate::table::pk_vector_orchestrator::{ - as_split_exact_file_search, build_indexed_splits, merge_candidates, OrchestratorSearchResult, - PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, -}; -use crate::table::pk_vector_position_read::{ - PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN, -}; -use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan}; -use crate::table::read_builder::resolve_projected_fields; -use crate::table::row_id_predicate::intersect_sorted_ranges; -use crate::table::source::DataSplit; -use crate::table::{ - find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, Table, -}; -use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; -use crate::vindex::executor::{ - acquire_process_global_search_permit, drain_indexed_jobs, - ensure_global_index_executor_capacity, execute_global_index_with_guard, -}; -use crate::vindex::pkvector::ann::{AnnSegmentSource, PkVectorAnnSearcher, VindexAnnSearcher}; -use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture}; -use crate::vindex::pkvector::exact::validate_query; -use crate::vindex::pkvector::metric::VectorSearchMetric; -use crate::vindex::pkvector::{FileRowSelection, FileRowSelections}; -use crate::vindex::range_reader::{RangeIoStats, RangeReadLimiter, VindexFileReader}; -use crate::vindex::reader::VindexVectorGlobalIndexReader; -use crate::vindex::{is_vindex_index_type, vector_search_timing_enabled, VindexVectorIndexOptions}; -use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; -use arrow_select::interleave::interleave_record_batch; -use futures::{stream, TryStreamExt}; -use paimon_vindex_core::blas::sgemm_a_bt; -use paimon_vindex_core::diskann_io::DISKANN_HEADER_SIZE; -use paimon_vindex_core::distance::MetricType; -use paimon_vindex_core::index::VectorIndexReader as VIndexReader; -use paimon_vindex_core::io::SeekRead; -use roaring::RoaringTreemap; -use std::cmp::Ordering; -use std::collections::{HashMap, HashSet}; -use std::io::Cursor; -use std::sync::Arc; -use std::time::{Duration, Instant}; +//! Configures vector queries and dispatches to global-index or primary-key readers. -const RAW_SCORE_MATRIX_MIN_QUERY_COUNT: usize = 4; -const RAW_SCORE_MATRIX_TARGET_ELEMENTS: usize = 1 << 20; -const RAW_TOP_K_MIN_PARTITION_SIZE: usize = 1 << 12; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum VectorIndexBackend { - Lumina, - Vindex, -} - -impl VectorIndexBackend { - fn from_index_type(index_type: &str) -> Option { - 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", - } - } -} - -async fn execute_vindex_searches( - io_meta: GlobalIndexIOMeta, - options: HashMap, - vector_searches: Vec, - source: S, - file_name: String, - index_parallelism: usize, - guard: G, -) -> crate::Result>>> { - let panic_context = if vector_searches.len() > 1 { - "vindex global-index batch search task failed" - } else { - "vindex global-index search task failed" - }; - execute_global_index_with_guard(panic_context, guard, move || { - let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options) - .with_batch_index_parallelism(index_parallelism); - reader - .visit_batch_vector_search(&vector_searches, |_| Ok(source)) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to read vindex index file '{}': {}", file_name, e), - source: Some(Box::new(e)), - }) - }) - .await -} - -fn current_tokio_runtime_handle() -> crate::Result { - tokio::runtime::Handle::try_current().map_err(|error| crate::Error::UnexpectedError { - message: "Vector index range reader requires a Tokio runtime".to_string(), - source: Some(Box::new(error)), - }) -} - -fn vindex_index_parallelism(entry_count: usize, max_concurrency: usize) -> usize { - entry_count.min(max_concurrency).max(1) -} - -fn log_vindex_range_io_stats(file: &str, query_count: usize, stats: &RangeIoStats) { - let stats = stats.snapshot(); - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_range_io file={} nq={} logical_ranges={} requested_bytes={} file_read_calls={} returned_bytes={} read_ahead_hits={} io_wait_sum_ms={:.3} range_permit_wait_sum_ms={:.3} peak_in_flight_reads={} read_many_merged_ranges={} read_many_chunks={} read_many_chunk_size_sum={} read_many_chunk_size_min={} read_many_chunk_size_max={}", - file, - query_count, - stats.logical_ranges, - stats.requested_bytes, - stats.file_read_calls, - stats.returned_bytes, - stats.read_ahead_hits, - stats.io_wait_nanos as f64 / 1_000_000.0, - stats.range_permit_wait_nanos as f64 / 1_000_000.0, - stats.peak_in_flight_reads, - stats.read_many_merged_ranges, - stats.read_many_chunks, - stats.read_many_chunk_size_sum, - stats.read_many_chunk_size_min, - stats.read_many_chunk_size_max, - ); -} - -fn vindex_concurrency_limits( - core_options: &CoreOptions<'_>, - entry_count: usize, - max_concurrency: usize, -) -> crate::Result<(usize, usize)> { - Ok(( - vindex_index_parallelism(entry_count, max_concurrency), - core_options.global_index_vindex_read_thread_num()?, - )) -} +use crate::spec::{CoreOptions, Predicate}; +use crate::table::vector_read::{BatchVectorRead, VectorRead}; +use crate::table::vector_scan::VectorScan; +use crate::table::Table; +use crate::vector_search::SearchResult; +use std::collections::HashMap; pub struct VectorSearchBuilder<'a> { table: &'a Table, @@ -187,90 +30,7 @@ pub struct VectorSearchBuilder<'a> { query_vector: Option>, limit: Option, options: HashMap, - projection: Option>, - filter: Option, -} - -pub struct BatchVectorSearchBuilder<'a> { - table: &'a Table, - vector_column: Option, - query_vectors: Option>>, - limit: Option, - options: HashMap, - projection: Option>, filter: Option, - include_row_ids: Option>, - prepared_filter: Option, -} - -/// A scalar vector pre-filter resolved once against one pinned snapshot. -/// -/// Reusing this value avoids repeating the same scalar-index/table read for -/// every input batch of a lateral vector query. -#[derive(Debug, Clone)] -pub struct PreparedVectorSearchFilter { - table: Table, - include_row_ids: Arc, -} - -impl PreparedVectorSearchFilter { - pub fn table(&self) -> &Table { - &self.table - } - - pub fn include_row_ids(&self) -> &Arc { - &self.include_row_ids - } -} - -fn same_vector_search_table(left: &Table, right: &Table) -> bool { - left.location().trim_end_matches('/') == right.location().trim_end_matches('/') - && left.branch() == right.branch() -} - -/// Unwrap a single-query result from a batch entry point that must return exactly -/// one element per input query. -/// -/// The batch terminals below are handed one query, so their result vector holds -/// exactly one entry. A `debug_assert_eq!(len, 1)` followed by `remove(0)` checked -/// that only in debug builds, where a release build would instead panic on an index -/// out of bounds for an empty vector -- or SILENTLY return the first of several, -/// pairing the caller's single query with another query's result. A length that is -/// wrong means the batch ran the wrong number of searches, which is a programming -/// error in this crate rather than bad input, so it is reported as one. -fn take_only_result(results: Vec, operation: &str) -> crate::Result { - let mut results = results.into_iter(); - let result = results - .next() - .ok_or_else(|| crate::Error::UnexpectedError { - message: format!("{operation} returned no result for one query"), - source: None, - })?; - if results.next().is_some() { - return Err(crate::Error::UnexpectedError { - message: format!("{operation} returned more than one result for one query"), - source: None, - }); - } - Ok(result) -} - -/// The primary-key vector route's search output plus the source context a later -/// materialization (or a hybrid fusion across routes) needs. `candidates` are the -/// best-first hits; `splits` are the per-bucket source splits their `split_index` -/// refers into (the authority for re-associating a hit to its `DataFileMeta`); -/// `snapshot_id` is the single snapshot the plan resolved during planning -/// (authoritative even when the plan yields zero splits; `0` only for a table -/// with no snapshot at all); `metric` is the resolved distance metric used to -/// turn distances into scores. Produced by -/// [`VectorSearchBuilder::search_pk_route`]. -pub(crate) struct PkVectorRouteResult { - pub(crate) candidates: Vec, - // Read by the hybrid route consumer (fuses routes before materializing); the - // materialized read reaches candidates/splits/metric directly. - pub(crate) snapshot_id: i64, - pub(crate) splits: Vec, - pub(crate) metric: VectorSearchMetric, } impl<'a> VectorSearchBuilder<'a> { @@ -281,7 +41,6 @@ impl<'a> VectorSearchBuilder<'a> { query_vector: None, limit: None, options: HashMap::new(), - projection: None, filter: None, } } @@ -324,8311 +83,62 @@ impl<'a> VectorSearchBuilder<'a> { self } - /// Restrict the columns materialized by [`execute_read`](Self::execute_read) - /// to `cols` (plus the always-appended `__paimon_search_score`). Without this - /// call `execute_read` materializes every user table column. Only affects - /// `execute_read`; the search-only paths ignore it. - pub fn with_projection(&mut self, cols: &[&str]) -> &mut Self { - self.projection = Some(cols.iter().map(|c| c.to_string()).collect()); - self - } - - pub async fn execute(&self) -> crate::Result> { - self.execute_scored().await?.to_row_ranges() - } - - pub async fn execute_scored(&self) -> crate::Result { - // Fail closed: returns data-derived row ranges outside `TableScan`/`TableRead`. - let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; - let vector_column = - self.vector_column - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Vector column must be set via with_vector_column()".to_string(), - })?; - let query_vector = - self.query_vector - .as_ref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Query vector must be set via with_query_vector()".to_string(), - })?; - let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { - message: "Limit must be set via with_limit()".to_string(), - })?; - - // Primary-key vector search branch: mirrors Java `PrimaryKeyVectorRead`. - // Only taken when the table enables the PK-vector index AND this query - // targets a configured PK-vector column; otherwise fall through to the - // data-evolution (DE) global-index path below. - // - // Membership is resolved via the non-erroring columns accessor so a - // malformed PK-vector config (e.g. a blank list) cannot abort an unrelated - // DE query. A query that does target the PK-vector column fails loud here: - // the PK path produces physical positions, not global row ids, so scored - // search is unsupported and callers must use `execute_read` instead. - if core.primary_key_vector_index_enabled() { - let targets_pk_column = core - .primary_key_vector_index_columns() - .ok() - .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); - if targets_pk_column { - return Err(crate::Error::DataInvalid { - message: "primary-key vector search does not produce global row ids; use the materialized read (execute_read) instead".to_string(), - source: None, - }); - } - } - - let mut batch_builder = BatchVectorSearchBuilder::new(self.table); - batch_builder - .with_vector_column(vector_column) - .with_query_vectors(vec![query_vector.clone()]) - .with_limit(limit) - .with_options(self.options.clone()); - if let Some(filter) = &self.filter { - batch_builder.with_filter(filter.clone()); - } - let results = batch_builder.execute().await?; - - take_only_result(results, "vector search") - } - - /// Run the vector search and materialize the matching rows as Arrow batches, - /// ordered best-first. Supported for both primary-key vector indexes and - /// data-evolution (global-index) vector search; a query targeting a column - /// that is neither fails loud. Output columns are the projected user table - /// columns (all user columns by default, or those set via - /// [`with_projection`](Self::with_projection)) plus `__paimon_search_score`; - /// `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always hidden. - pub async fn execute_read(&self) -> crate::Result { - // Fail closed: returns data outside `TableScan`/`TableRead`. - let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; - let vector_column = - self.vector_column - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Vector column must be set via with_vector_column()".to_string(), - })?; - let query_vector = - self.query_vector - .as_ref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Query vector must be set via with_query_vector()".to_string(), - })?; - let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { - message: "Limit must be set via with_limit()".to_string(), - })?; - - // Only the primary-key vector path can materialize rows. The data-evolution - // (global-index) path returns data-derived row-ids, not table rows, so a - // read against it (or against a non-PK-vector column) fails loud. - if core.primary_key_vector_index_enabled() { - let targets_pk_column = core - .primary_key_vector_index_columns() - .ok() - .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); - if targets_pk_column { - let pk_col = core.primary_key_vector_index_column()?; - return self - .execute_primary_key_vector_read(&core, &pk_col, query_vector, limit) - .await; - } - } - - // Data-evolution (global-index) vector search: materialize rows from the - // scored global row-ids and attach the unified score column. A non-vector - // column or a set filter fails loud inside execute_scored below. - self.execute_de_vector_read(vector_column, query_vector, limit) - .await - } - - /// Run this search over bucket splits an engine planned elsewhere, and - /// materialize the hits. - /// - /// The unit of work is Java's `BucketVectorSearchSplit` byte form: a planner - /// running in Paimon Java enumerates one split per bucket -- a bucket is never - /// divided, because the ANN current-segment decision needs the bucket's whole - /// active file set -- and ships each to a worker that calls this. The splits - /// are the plan: their payload files, their per-file row ranges and the - /// snapshot they pin are used as given, and this table's index manifest is not - /// read. - /// - /// Everything after planning is the ordinary primary-key vector read, so - /// search, optional refine, local Top-K and materialization stay identical to - /// [`execute_read`](Self::execute_read): output is the projected user columns - /// plus `__paimon_search_score`, best-first. The Top-K is local to the supplied - /// splits; a caller distributing one call per bucket merges the per-bucket - /// results itself. - /// - /// Only a primary-key vector column can be read this way. The data-evolution - /// route plans through the global index rather than through bucket splits, so - /// it is rejected rather than silently answered from a different plan. - pub async fn execute_read_for_bucket_splits( - &self, - split_bytes: &[&[u8]], - ) -> crate::Result { - // Fail closed: returns data outside `TableScan`/`TableRead`. - let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; - let vector_column = - self.vector_column - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Vector column must be set via with_vector_column()".to_string(), - })?; - let query_vector = - self.query_vector - .as_ref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Query vector must be set via with_query_vector()".to_string(), - })?; - let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { - message: "Limit must be set via with_limit()".to_string(), - })?; - - let pk_col = if core.primary_key_vector_index_enabled() { - let targets_pk_column = core - .primary_key_vector_index_columns() - .ok() - .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); - if targets_pk_column { - core.primary_key_vector_index_column()? - } else { - return Err(bucket_split_route_error(vector_column)); - } - } else { - return Err(bucket_split_route_error(vector_column)); - }; - - // Decoding is the trust boundary: these bytes come from outside the - // process. Reject an empty request here rather than let it reach planning - // as "no splits", which cannot pin a snapshot. - if split_bytes.is_empty() { - return Err(crate::Error::DataInvalid { - message: "bucket-split read requires at least one split".to_string(), - source: None, - }); - } - let splits = split_bytes - .iter() - .map(|bytes| BucketVectorSearchSplit::deserialize(bytes)) - .collect::>>()?; - - // Resolve the query parameters (and reject a query the search cannot answer - // correctly) before planning, exactly as the manifest route does. - let params = resolve_pk_vector_search_params( - self.table, - &self.options, - self.filter.as_ref(), - &core, - &pk_col, - &[query_vector.as_slice()], - limit, - )?; - let plan = PkVectorScan::new( - self.table, - params.field_id, - params.index_type.clone(), - self.filter.clone(), - ) - .plan_for_bucket_vector_splits(splits)?; - - // Resolve the materialization read-type up front so an invalid projection - // fails loud even when the plan is empty and no rows will be read. - let read_type = self.resolve_materialize_read_type()?; - - let candidates = search_pk_candidates_batch_with_plan( - self.table, - &self.options, - self.filter.as_ref(), - &core, - &pk_col, - &[query_vector.as_slice()], - limit, - &plan, - ¶ms, - ) - .await?; - // One query in, so one candidate list out -- no more and no less. Checked - // rather than asserted, because this route is reached from the C ABI, where a - // debug-only assert leaves a release build indexing an empty vector or - // answering the caller's one query with another query's list. - let candidates = take_only_result(candidates, "bucket-split vector search")?; - - // A separate, predicate-free materialization reader projecting the user - // columns (the search reader projects only the vector column). - let materialize_reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - read_type, - Vec::new(), - ); - - Self::materialize_candidates(candidates, &plan.splits, params.metric, &materialize_reader) - .await - } - - /// Materialize the best-first data-evolution vector search hits into Arrow - /// rows. The global-index search returns global `_ROW_ID`s and their scores; a - /// subsequent row-range read materializes those rows, and each row's score is - /// joined back by `_ROW_ID`. Output columns are the projected user table - /// columns (all user columns by default) plus `__paimon_search_score`; `_ROW_ID` - /// is always hidden. A scalar filter is applied before vector Top-K by the - /// snapshot-pinned scored search below. - async fn execute_de_vector_read( - &self, - vector_column: &str, - query_vector: &[f32], - limit: usize, - ) -> crate::Result { - // Validate the target column exists and is a vector-bearing type before any - // work. The data-evolution search returns an empty result for an unknown - // field (its scored-path behavior), which would make a typo'd or scalar - // column look like a normal empty read here — violating `execute_read`'s - // fail-loud contract (a C/Doris caller would see EOF, not an input error). - // Reject it up front instead. - let field = self - .table - .schema() - .fields() - .iter() - .find(|f| f.name() == vector_column) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("vector search column '{vector_column}' does not exist"), - source: None, + /// Create a query-independent scan. Only the vector column must be configured. + pub fn new_scan(&self) -> crate::Result { + CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + let column = self + .vector_column + .as_deref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Vector column must be set via with_vector_column()".to_string(), })?; - // Require a FLOAT-element vector column: `ARRAY` or `VECTOR`, - // matching the element type the vector index/search operates on. An - // `ARRAY` (or any non-float element) is not a searchable vector column. - let is_float_vector = match field.data_type() { - DataType::Vector(t) => matches!(t.element_type(), DataType::Float(_)), - DataType::Array(t) => matches!(t.element_type(), DataType::Float(_)), - _ => false, - }; - if !is_float_vector { - return Err(crate::Error::DataInvalid { - message: format!( - "vector search column '{vector_column}' must be a FLOAT vector column \ - (ARRAY or VECTOR), got {:?}", - field.data_type() - ), - source: None, - }); - } - - // Resolve the projected user columns up front so an invalid projection - // fails loud even when the result is empty. - let mut read_type = self.resolve_materialize_read_type()?; - - let Some(snapshot) = crate::table::time_travel::resolve_snapshot(self.table).await? else { - return Ok(Box::pin(stream::empty())); - }; - let pinned_table = self.table.copy_with_resolved_snapshot(&snapshot).await?; - let mut search_builder = pinned_table.new_vector_search_builder(); - search_builder - .with_vector_column(vector_column) - .with_query_vector(query_vector.to_vec()) - .with_limit(limit) - .with_options(self.options.clone()); - if let Some(filter) = &self.filter { - search_builder.with_filter(filter.clone()); - } - let sr = search_builder.execute_scored().await?; - - if sr.is_empty() { - return Ok(Box::pin(stream::empty())); - } - - // rank = ordinal in the best-first scored result; score = the aligned score. - // Build ranges first (validates ids fit in i64::MAX) before constructing the map. - let ranges = sr.to_row_ranges()?; - let mut rank_score_of: HashMap = HashMap::new(); - for (rank, (&id, &score)) in sr.row_ids.iter().zip(sr.scores.iter()).enumerate() { - rank_score_of.insert(id as i64, (rank, score)); - } - - // Add _ROW_ID as the join key for score alignment; it is stripped before output. - if !read_type.iter().any(|f| f.name() == ROW_ID_FIELD_NAME) { - read_type.push(row_id_data_field()); - } - - let mut read_builder = pinned_table.new_read_builder(); - read_builder - .with_read_type(read_type) - .with_row_ranges(ranges); - let scan = read_builder.new_scan(); - let plan = scan.plan().await?; - let table_read = read_builder.new_read()?; - let mut stream = table_read.to_arrow(plan.splits())?; - - let mut batches: Vec = Vec::new(); - while let Some(batch) = stream.try_next().await? { - batches.push(batch); - } - let output = attach_scores_by_row_id(&batches, &rank_score_of, sr.len())?; - Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) - } - - /// Single-query wrapper over - /// [`plan_and_search_pk_candidates_batch`]: plan once, search the one query, - /// and return its candidate list. Output is byte-identical to the batch-of-one - /// path. - async fn plan_and_search_pk_candidates( - &self, - core: &CoreOptions<'_>, - pk_col: &str, - query_vector: &[f32], - limit: usize, - ) -> crate::Result<(Vec, PkVectorScanPlan, VectorSearchMetric)> { - let (candidates, plan, metric) = plan_and_search_pk_candidates_batch( - self.table, - &self.options, - self.filter.as_ref(), - core, - pk_col, - &[query_vector], - limit, - ) - .await?; - Ok(( - take_only_result(candidates, "planned vector search")?, - plan, - metric, - )) - } - - /// Plan + search the primary-key vector route and return the best-first - /// candidates together with the route source context needed to materialize - /// them later: the resolved snapshot id and the per-bucket source splits - /// (`split_index` is only meaningful against these originating splits), plus - /// the resolved distance metric. This is the hybrid-reachable entry point — - /// it runs exactly the plan/search core [`execute_read`](Self::execute_read) - /// uses but stops before materialization, so a caller can fuse these - /// candidates before materializing. The materialized read is layered on top - /// of it (see `execute_primary_key_vector_read`). The snapshot id is the one - /// the plan pinned during planning — authoritative even for an empty plan - /// (which also yields empty candidates and empty splits); it is `0` only for - /// a table with no snapshot at all. - pub(crate) async fn search_pk_route( - &self, - core: &CoreOptions<'_>, - pk_col: &str, - query_vector: &[f32], - limit: usize, - ) -> crate::Result { - let (candidates, plan, metric) = self - .plan_and_search_pk_candidates(core, pk_col, query_vector, limit) - .await?; - // Planning pins a single snapshot (`plan.snapshot_id`) even when it yields - // zero searchable splits, so report it unconditionally rather than deriving - // it from a split that may not exist. - Ok(PkVectorRouteResult { - candidates, - snapshot_id: plan.snapshot_id, - splits: plan.splits, - metric, - }) - } - - /// Materialize the best-first PK-vector search hits into Arrow rows. Mirrors - /// Java `PrimaryKeyVectorRead` feeding its result splits into an ordinary table - /// read: the search decides which rows, a subsequent read decides which - /// columns. - /// - /// Output columns are the projected user table columns (all user columns when - /// [`with_projection`](Self::with_projection) was not called) plus - /// `__paimon_search_score`; `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always - /// hidden. Rows are emitted best-first (the candidate order), which differs - /// from the file/position order the orchestrator materializes in. - async fn execute_primary_key_vector_read( - &self, - core: &CoreOptions<'_>, - pk_col: &str, - query_vector: &[f32], - limit: usize, - ) -> crate::Result { - let PkVectorRouteResult { - candidates, - splits, - metric, - .. - } = self - .search_pk_route(core, pk_col, query_vector, limit) - .await?; - - // Resolve the materialization read-type up front so an invalid projection - // (unknown column, or a reserved metadata / row-id name) fails loud - // unconditionally, even when the plan is empty and no rows will be read. - // Default (no `with_projection`) is every user table column. - let read_type = self.resolve_materialize_read_type()?; - - // A separate, predicate-free materialization reader projecting the user - // columns (the search reader projects only the vector column). Mirrors - // `table_read.rs::new_data_file_reader` with an empty predicate list. - let materialize_reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - read_type, - Vec::new(), - ); - - Self::materialize_candidates(candidates, &splits, metric, &materialize_reader).await - } - - /// Materialize one best-first candidate list into an Arrow stream, best-first, - /// with a `__paimon_search_score` column and `_PKEY_VECTOR_POSITION` stripped. - /// An empty candidate list yields an empty stream (never skipped) so a batch - /// caller preserves per-query arity. `materialize_reader` must project the - /// output columns (predicate-free). Both the single-query and batch read paths - /// use this so their materialization is identical. - async fn materialize_candidates( - candidates: Vec, - splits: &[PkVectorSearchSplit], - metric: VectorSearchMetric, - materialize_reader: &DataFileReader, - ) -> crate::Result { - if candidates.is_empty() { - return Ok(Box::pin(stream::empty())); - } - - // Rank each candidate by its best-first position, then reduce the physical - // materialization order back to best-first. The orchestrator emits rows in - // ascending (partition, bucket, file, position); the rank map keyed by - // (partition bytes, bucket, file, position) recovers the candidate order. - let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); - for (rank, c) in candidates.iter().enumerate() { - rank_of.insert( - ( - c.partition.to_serialized_bytes(), - c.bucket, - c.data_file_name.clone(), - c.row_position, - ), - rank, - ); - } - - let indexed_splits = build_indexed_splits(candidates, splits, metric)?; - - // Materialize every indexed split, retaining each batch and, per row, the - // (rank, batch_index, row_index) tuple so we can reorder to best-first. - // Top-K is small, so full in-memory collection is acceptable. - let mut batches: Vec = Vec::new(); - let mut ranked: Vec = Vec::new(); - for indexed in indexed_splits { - let partition_bytes = indexed.split.partition().to_serialized_bytes(); - let bucket = indexed.split.bucket(); - let file_name = indexed.split.data_files()[0].file_name.clone(); - let mut stream = - PkVectorIndexedSplitRead::new(materialize_reader.clone()).read(&indexed)?; - while let Some(batch) = stream.try_next().await? { - let batch_index = batches.len(); - collect_ranked_rows( - &batch, - batch_index, - &partition_bytes, - bucket, - &file_name, - &rank_of, - &mut ranked, - )?; - batches.push(batch); - } - } - - // Reorder to best-first and drop the position column. - let output = reorder_and_strip_position(&batches, ranked)?; - Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) - } - - /// Resolve the projected fields for the materialization read-type. Default - /// (no projection set) is all user table fields; otherwise the requested - /// names resolved via `resolve_projected_fields`. Rejects reserved metadata - /// names and `_ROW_ID` so a user cannot request a hidden column. - fn resolve_materialize_read_type(&self) -> crate::Result> { - let fields = match &self.projection { - None => self.table.schema().fields().to_vec(), - Some(names) => { - for name in names { - if is_reserved_read_column(name) { - return Err(crate::Error::DataInvalid { - message: format!( - "vector search read projection must not request reserved column '{name}'" - ), - source: None, - }); - } - } - resolve_projected_fields( - self.table.identifier().full_name(), - self.table.schema().fields(), - names, - true, - )? - } - }; - // The default projection returns every user column, so a user column - // whose name collides with an injected metadata column must be rejected - // on the resolved field list too — not only when explicitly requested. - ensure_no_reserved_read_columns(&fields)?; - Ok(fields) - } -} - -/// Names a read injects as metadata columns — `__paimon_search_score`, -/// `_PKEY_VECTOR_POSITION`, and `_ROW_ID` — that a materialized read type must -/// not reuse for a user column. -fn is_reserved_read_column(name: &str) -> bool { - name == PKEY_VECTOR_POSITION_COLUMN || name == SEARCH_SCORE_COLUMN || name == ROW_ID_FIELD_NAME -} - -/// Reject a materialized read type whose resolved fields contain a reserved -/// metadata column name. Applied to the RESOLVED field list so the default -/// (all user columns) projection is covered, not only an explicit one. -pub(crate) fn ensure_no_reserved_read_columns(fields: &[DataField]) -> crate::Result<()> { - for field in fields { - if is_reserved_read_column(field.name()) { - return Err(crate::Error::DataInvalid { - message: format!( - "search read must not include reserved column '{}'", - field.name() - ), - source: None, - }); - } - } - Ok(()) -} - -/// Batch PK-vector search core shared by the single and batch builders: plan ONE -/// per-bucket split set, lazy segment loader, ANN scorer, exact-fallback search -/// closure, and residual allow-list (all query-independent), then run -/// `search_candidates_batch` ONCE so N queries share the opened readers. Per -/// query, the approximate candidates are exact-reranked (when a refine factor is -/// set) and merged with the exact fallback into one best-first list bounded to -/// `limit`. Returns one candidate list per query (outer index aligned to -/// `queries`), together with the shared plan and resolved metric so the caller can -/// materialize each query's rows. An empty plan yields one empty candidate list -/// per query. -/// -/// The residual allow-list depends only on `filter` and the plan, NOT the query -/// vector, so it is computed once and the SAME slice is shared across all queries. -/// Rerank stays per-query (each query reranks its own indexed list). -#[allow(clippy::too_many_arguments)] -/// Query-level parameters for a primary-key vector search: everything resolvable -/// from the table schema, the options and the queries alone, independent of which -/// splits planning yields. Resolved before planning so a malformed query or option -/// fails loud even when the plan turns out empty. -struct PkVectorSearchParams { - metric: VectorSearchMetric, - /// Fan-out limit for bucket orchestration plus ANN and exact-file leaves (Java - /// `GLOBAL_INDEX_THREAD_NUM`); `1` reproduces strictly sequential execution. - concurrency: usize, - index_type: String, - field_id: i32, - vector_field: DataField, - skip_exact_fallback: bool, - refine_factor: usize, - indexed_limit: usize, -} - -/// A bucket split is a primary-key vector plan. The data-evolution route plans -/// through the global index instead, so answering it here would silently use a -/// different plan than the caller supplied. -fn bucket_split_route_error(vector_column: &str) -> crate::Error { - crate::Error::DataInvalid { - message: format!( - "bucket-split read requires a primary-key vector column, but '{vector_column}' is not one" - ), - source: None, - } -} - -/// Resolve the query-level parameters and reject a query the search cannot answer -/// correctly, before any planning or read happens. -fn resolve_pk_vector_search_params( - table: &Table, - query_options: &HashMap, - filter: Option<&Predicate>, - core: &CoreOptions<'_>, - pk_col: &str, - queries: &[&[f32]], - limit: usize, -) -> crate::Result { - // Residual pre-filter guard, mirroring Java `PrimaryKeyVectorScan`. A DATA - // predicate set via `with_filter` is applied post-recall by re-reading each - // candidate file's physical rows (see below). That physical-position filtering - // only agrees with the bucket search when the table exposes physical rows - // directly: deletion vectors enabled and merge-on-read disabled. Under - // merge-on-read (or without deletion vectors) a read merges multiple key - // versions, so a scalar filter could retain a stale version whose live version - // does not match — a silent wrong-read. Reject such queries rather than answer - // them incorrectly. - // - // Guard on the DATA conjuncts, not the whole filter: partition-only conjuncts - // are enforced entirely by scan planning (partition pruning) and produce no - // per-row residual, so they need no physical-row read. This mirrors Java, where - // `BatchVectorSearchBuilderImpl.withFilter` splits at the builder level and - // leaves `this.filter == null` for a partition-only filter — the scan guard is - // then skipped. No data predicate (partition-only or no filter) → nothing to - // guard, so the search-only and read paths are unaffected. - let physical_row_read = - core.deletion_vectors_enabled() && !core.deletion_vectors_merge_on_read(); - let has_data_predicate = filter.is_some_and(|f| { - let (_partition, data) = split_partition_and_data_predicates( - f.clone(), - table.schema().fields(), - table.schema().partition_keys(), - ); - !data.is_empty() - }); - if has_data_predicate && !physical_row_read { - return Err(crate::Error::DataInvalid { - message: - "primary-key vector pre-filter requires deletion vectors without merge-on-read" - .to_string(), - source: None, - }); - } - // `primary_key_vector_distance_metric` returns a validated name; re-parse into - // the enum for the numeric semantics. - let metric = VectorSearchMetric::parse(&core.primary_key_vector_distance_metric(pk_col)?)?; - // Fan-out limit for bucket orchestration plus ANN and exact-file leaves (Java - // `GLOBAL_INDEX_THREAD_NUM`); `1` reproduces strictly sequential execution. - let concurrency = core.global_index_thread_num()?; - let index_type = core.primary_key_vector_index_type(pk_col)?; - let field_id = find_field_id_by_name(table.schema().fields(), pk_col).ok_or_else(|| { - crate::Error::DataInvalid { - message: format!("PK-vector column '{pk_col}' not found in schema"), - source: None, - } - })?; - let vector_field = table - .schema() - .fields() - .iter() - .find(|f| f.name() == pk_col) - .cloned() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("PK-vector column '{pk_col}' not found in schema"), - source: None, - })?; - - let search_mode = core.vector_index_search_mode()?; - let skip_exact_fallback = search_mode == GlobalIndexSearchMode::Fast; - - // A non-positive limit is invalid regardless of the plan; reject it before - // planning so an empty plan cannot mask it with empty results. - if limit == 0 { - return Err(crate::Error::DataInvalid { - message: "vector search limit must be positive".to_string(), - source: None, - }); - } - - // Resolve the refine factor from the query options first, then fall back to - // the table options; a positive factor over-fetches indexed (approximate) - // candidates so the exact rerank below has a wider pool to reorder. Factor 0 - // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank - // path. The two option maps are kept distinct (query options passed - // separately from table options) so a broad query key cannot be overridden - // by a more specific table key: query options take precedence as a whole. - // Resolved before planning so an invalid factor (e.g. a non-numeric value) - // fails loud regardless of whether the table currently has searchable data. - let refine_factor = - configured_refine_factor(query_options, table.schema().options(), pk_col, &index_type)?; - let indexed_limit = indexed_search_limit(limit, refine_factor)?; - - // Validate every query against the vector column's dimension (and finiteness) - // before planning or any read, so a malformed query fails loud even when the - // plan turns out empty. VECTOR carries the dimension in its type; - // ARRAY gets the index dimension from the same vindex option resolver - // used by index reads. Both valid PK-vector column shapes must reject NaN/Inf - // up front, not only after a non-empty plan opens readers. - if let Some(dimension) = pk_vector_query_dimension( - table.schema().options(), - query_options, - &index_type, - &vector_field, - )? { - for query in queries { - validate_query(query, dimension)?; - } - } - - Ok(PkVectorSearchParams { - metric, - concurrency, - index_type, - field_id, - vector_field, - skip_exact_fallback, - refine_factor, - indexed_limit, - }) -} - -/// Search an already-resolved plan across every query and return each query's raw -/// indexed and exact candidate lists, before any rerank or merge. -/// -/// Plan-dependent concurrency — the vindex segment count, batch-index parallelism -/// and the range-read bound — is derived here from the plan that is actually being -/// searched, so a narrowed plan can never be searched under limits computed for a -/// wider one. -/// Combine the two per-split row allow-lists a search can be handed: the physical -/// rows an engine-supplied plan restricts each file to, and the positions a residual -/// data predicate leaves behind. -/// -/// The two sides read a file's ABSENCE differently, and the merge has to respect -/// both readings: -/// -/// * The plan lists only what the engine's split narrowed, so an absent file is -/// unrestricted -- Java's `rowRangesByFile.get(file) == null`. -/// * The residual is exhaustive over the files a search can read from -/// (`residual_positions_by_file` registers every active file, empty when nothing -/// passed), so once a residual exists its silence about a file means "no rows". -/// -/// So: with no residual, a file the plan omits stays absent and unrestricted. With a -/// residual, a file it omits is excluded even if the plan restricted it, and a file -/// both describe keeps the intersection. Absent from BOTH is unrestricted, which is -/// what lets the ANN backend search unfiltered. -/// -/// The plan's ranges stay ranges. Expanding them into positions would be work sized -/// by row counts that arrived on the wire; where an intersection is genuinely needed -/// the residual positions — bounded by the rows its own read returned — are filtered -/// BY the ranges instead. When the residual was evaluated over those same ranges the -/// intersection cannot remove anything, and is kept as the invariant that says so. -fn intersect_row_allow_lists( - physical: Option<&[HashMap>]>, - residual: Option>>, - split_count: usize, -) -> crate::Result>> { - if let Some(maps) = physical { - if maps.len() != split_count { - return Err(crate::Error::DataInvalid { - message: format!( - "plan carries {} physical row allow-lists for {split_count} splits", - maps.len() - ), - source: None, - }); - } - } - if let Some(maps) = residual.as_ref() { - if maps.len() != split_count { - return Err(crate::Error::DataInvalid { - message: format!( - "residual carries {} row allow-lists for {split_count} splits", - maps.len() - ), - source: None, - }); - } - } - match (physical, residual) { - (None, None) => Ok(None), - (None, Some(residual)) => Ok(Some( - residual - .into_iter() - .map(|per_file| { - per_file - .into_iter() - .map(|(file, positions)| (file, FileRowSelection::Positions(positions))) - .collect() - }) - .collect(), - )), - (Some(physical), None) => Ok(Some( - physical - .iter() - .map(|per_file| { - per_file - .iter() - .map(|(file, ranges)| { - (file.clone(), FileRowSelection::Ranges(ranges.clone())) - }) - .collect() - }) - .collect(), - )), - (Some(physical), Some(residual)) => { - Ok(Some( - physical - .iter() - .zip(residual) - .map(|(physical, mut residual)| { - let mut merged: FileRowSelections = HashMap::new(); - for (file, ranges) in physical { - let range_selection = FileRowSelection::Ranges(ranges.clone()); - let selection = match residual.remove(file.as_str()) { - // Both restrict: keep the positions the ranges also - // allow. Filtering the positions (bounded by the read) - // by the ranges never expands the ranges. - Some(positions) => FileRowSelection::Positions( - positions - .iter() - .filter(|position| range_selection.contains(*position)) - .collect(), - ), - // The residual is exhaustive over the files the search - // can read from -- `residual_positions_by_file` - // registers every active file, empty when nothing - // passed. Its silence about a file therefore means "no - // rows", NOT "unrestricted", and must stay fail-closed - // here even though the plan has something to say. - None => FileRowSelection::Positions(RoaringTreemap::new()), - }; - merged.insert(file.clone(), selection); - } - // Whatever the residual restricted and the plan did not. - merged.extend(residual.into_iter().map(|(file, positions)| { - (file, FileRowSelection::Positions(positions)) - })); - merged - }) - .collect(), - )) - } + VectorScan::new(self.table, column, self.filter.as_ref(), None, None) } -} - -#[allow(clippy::too_many_arguments)] -async fn search_pk_raw_candidates_batch_with_plan( - table: &Table, - query_options: &HashMap, - filter: Option<&Predicate>, - core: &CoreOptions<'_>, - pk_col: &str, - queries: &[&[f32]], - limit: usize, - plan: &PkVectorScanPlan, - params: &PkVectorSearchParams, -) -> crate::Result> { - // An empty plan has nothing to search. Returned before the backend is resolved - // so a table with no searchable data never errors on an unrecognized index type. - if plan.splits.is_empty() { - return Ok(queries - .iter() - .map(|_| OrchestratorSearchResult { - indexed: Vec::new(), - exact: Vec::new(), - }) - .collect()); - } - - let metric = params.metric; - let concurrency = params.concurrency; - let index_type = params.index_type.clone(); - let vector_field = params.vector_field.clone(); - let skip_exact_fallback = params.skip_exact_fallback; - let indexed_limit = params.indexed_limit; - - // Resolve the vector index backend from the single configured index type. - // Java enforces one index type per PK table and Rust filters segments to it, - // so one backend serves every segment. Computed after the empty-plan return so - // an empty table never errors on an unrecognized type. - let backend = VectorIndexBackend::from_index_type(&index_type).ok_or_else(|| { - crate::Error::DataInvalid { - message: format!("unsupported PK vector index backend/type: '{index_type}'"), - source: None, - } - })?; - let (batch_index_parallelism, range_read_concurrency) = match backend { - VectorIndexBackend::Vindex => vindex_concurrency_limits( - core, - plan.splits - .iter() - .map(|split| split.ann_segments.len()) - .sum(), - concurrency, - )?, - VectorIndexBackend::Lumina => (1, 0), - }; - - // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` - // but projecting only the vector column with no predicates. - let reader = DataFileReader::new( - table.file_io().clone(), - table.schema_manager().clone(), - table.schema().id(), - table.schema().fields().to_vec(), - vec![vector_field.clone()], - Vec::new(), - ); - - // Real ANN scorer + loader. Each segment source is opened lazily inside its - // bucket leaf and dropped after scoring. Lumina keeps its buffered-byte path; - // vindex remains range-backed and reads only metadata and probed lists. - let options = { - let mut o = table.schema().options().clone(); - o.extend(query_options.clone()); - o - }; - let search_options = options.clone(); - let field_name = pk_col.to_string(); - let loader_io = table.file_io().clone(); - let loader_range_read_limiter = match backend { - VectorIndexBackend::Vindex => Some(RangeReadLimiter::new(range_read_concurrency)), - VectorIndexBackend::Lumina => None, - }; - let loader: crate::vindex::pkvector::ann::SourceSegmentLoader = Box::new( - move |segment: &BucketAnnSegment| { - let io = loader_io.clone(); - let range_read_limiter = loader_range_read_limiter.clone(); - let path = segment.path.clone(); - let file_size = segment.file_size; - Box::pin(async move { - let input = io.new_input(&path)?; - match backend { - VectorIndexBackend::Lumina => input - .read() - .await - .map(AnnSegmentSource::Buffered) - .map_err(|error| crate::Error::DataInvalid { - message: format!("failed to read ANN index file '{path}': {error}"), - source: None, - }), - VectorIndexBackend::Vindex => { - let file_reader = - input - .reader() - .await - .map_err(|error| crate::Error::DataInvalid { - message: format!( - "failed to open ANN index file '{path}' for range reads: {error}" - ), - source: None, - })?; - Ok(AnnSegmentSource::Vindex( - VindexFileReader::new_with_limiter( - Arc::new(file_reader), - current_tokio_runtime_handle()?, - range_read_limiter.expect("Vindex range-read limiter"), - file_size, - path, - ), - )) - } - } - }) - }, - ); - - let scorer: crate::vindex::pkvector::ann::SourceBatchScorer = Box::new( - move |segment: &BucketAnnSegment, source: AnnSegmentSource, searches: &[VectorSearch]| { - let io_meta = GlobalIndexIOMeta::new( - segment.path.clone(), - segment.file_size, - segment.index_meta.clone(), - ); - match (backend, source) { - (VectorIndexBackend::Lumina, AnnSegmentSource::Buffered(data)) => { - let lumina_metric = - LuminaIndexMeta::deserialize(&segment.index_meta)?.metric()?; - verify_segment_metric(metric, VectorSearchMetric::from_lumina(lumina_metric))?; - let mut reader = LuminaVectorGlobalIndexReader::new(io_meta, options.clone()); - reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data))) - } - (VectorIndexBackend::Vindex, AnnSegmentSource::Vindex(source)) => { - let range_io_stats = source.range_io_stats(); - let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()) - .with_batch_index_parallelism(batch_index_parallelism); - let results = reader.visit_batch_vector_search_validated( - searches, - |_| Ok(source), - |metadata| { - verify_segment_metric( - metric, - VectorSearchMetric::from_vindex(metadata.metric), - ) - }, - )?; - if let Some(stats) = range_io_stats { - log_vindex_range_io_stats(&segment.path, searches.len(), &stats); - } - Ok(results) - } - (VectorIndexBackend::Lumina, AnnSegmentSource::Vindex(_)) - | (VectorIndexBackend::Vindex, AnnSegmentSource::Buffered(_)) => { - Err(crate::Error::DataInvalid { - message: format!( - "ANN segment '{}' was loaded with the wrong backend source", - segment.path - ), - source: None, - }) - } - } - }, - ); - let ann_searcher: Arc = Arc::new(VindexAnnSearcher::new_with_source( - field_name, scorer, loader, - )); - - // Residual (post-recall) filtering: for each candidate file, re-read its - // physical rows and keep the positions whose rows satisfy the filter. The - // per-split allow-list is threaded into the bucket search so the residual folds - // into recall (best-first order and Top-K are preserved). Built only when the - // filter has data (non-partition) conjuncts; a partition-only filter (or no - // filter) leaves `None`, which leaves the search unfiltered — partition - // pruning is already handled in planning. The residual depends only on the - // filter and the plan, not the query vector, so it is computed once here and - // shared across every query in the batch. The residual reader projects only - // the predicate columns and carries no pushdown; `residual_positions_by_file` - // recovers each surviving row's file-local physical position from its ordinal - // in the unfiltered scan (no `_ROW_ID`, no `first_row_id`). A file the - // allow-list leaves empty is skipped by the bucket search without opening an - // exact reader. - let residual_by_split: Option>> = match filter { - Some(filter) => { - // The whole filter is pushed into scan planning (`PkVectorScan`), where - // partition-only conjuncts already prune partitions/files. Re-applying - // them as a per-row residual would be redundant, so keep only the data - // conjuncts here — a partition-only filter then needs no residual at - // all. Mixed partition/data conjuncts stay whole in `data_predicates` - // and evaluate against the materialized partition column (partition - // columns are physically present in primary-key data files), so there - // is no missing-column case to reject. - let (_partition_predicate, data_predicates) = split_partition_and_data_predicates( - filter.clone(), - table.schema().fields(), - table.schema().partition_keys(), - ); - if data_predicates.is_empty() { - None - } else { - let file_predicates = FilePredicates { - predicates: data_predicates, - row_filter_factory: None, - file_fields: table.schema().fields().to_vec(), - }; - let residual_read_type = widen_scan_fields(&[], Some(&file_predicates)); - let residual_reader = DataFileReader::new( - table.file_io().clone(), - table.schema_manager().clone(), - table.schema().id(), - table.schema().fields().to_vec(), - residual_read_type, - Vec::new(), - ); - let mut per_split = Vec::with_capacity(plan.splits.len()); - for (index, split) in plan.splits.iter().enumerate() { - // The plan's selection for this split, so the residual is - // evaluated over the rows an engine-supplied split allows rather - // than over the whole file. - let allowed_rows = plan - .physical_row_ranges_by_split - .as_ref() - .and_then(|per_split| per_split.get(index)); - per_split.push( - residual_positions_by_file( - &residual_reader, - &split.data_split, - &split.active_files, - &file_predicates, - allowed_rows, - ) - .await?, - ); - } - Some(per_split) - } - } - None => None, - }; - // Fold the plan's own positional restriction into the same allow-list. A plan - // built from engine-supplied bucket splits carries the physical positions each - // file is limited to; a plan read from the index manifest carries none. Both - // sides list what is permitted, so combining them is an intersection. - let row_selections_by_split = intersect_row_allow_lists( - plan.physical_row_ranges_by_split.as_deref(), - residual_by_split, - plan.splits.len(), - )?; - - // Build the exact-fallback search on demand: the kernel calls this only for a - // file it actually searches (uncovered by ANN, residual-allowed, and only when - // the search mode is not FAST). Everything the future needs is cloned/owned up - // front so it borrows neither the split nor the file across the await. The - // search streams the file's vector column one Arrow batch at a time into - // per-query bounded heaps (all queries share one stream). - let reader_for_factory = reader.clone(); - let vector_field_for_factory = vector_field.clone(); - // The plan's own per-file selection, so an exact fallback reads only the rows an - // engine-supplied split allows. `is_excluded` still rejects on top of it, but it - // cannot un-read a row. - let physical_for_factory = plan.physical_row_ranges_by_split.clone(); - let factory = as_split_exact_file_search( - move |split_index: usize, - split: &PkVectorSearchSplit, - file: &BucketActiveFile, - queries: &[&[f32]], - metric: VectorSearchMetric, - exact_limit: usize, - is_excluded: &(dyn Fn(i64) -> bool + Sync)| - -> ExactFileSearchFuture<'_> { - let reader = reader_for_factory.clone(); - let vector_field = vector_field_for_factory.clone(); - let data_split = split.data_split.clone(); - let active = BucketActiveFile { - file_name: file.file_name.clone(), - row_count: file.row_count, - }; - let owned_queries: Vec> = queries.iter().map(|q| q.to_vec()).collect(); - let allowed_rows = physical_for_factory.as_ref().and_then(|per_split| { - per_split - .get(split_index) - .and_then(|per_file| per_file.get(&active.file_name)) - .cloned() - }); - Box::pin(async move { - let factory = DataFilePkVectorReaderFactory::new(reader, data_split, vector_field)?; - let query_refs: Vec<&[f32]> = owned_queries.iter().map(|q| q.as_slice()).collect(); - factory - .search_file( - &active, - &query_refs, - metric, - exact_limit, - is_excluded, - allowed_rows.as_deref(), - ) - .await - }) - }, - ); - - // Resolve the refine factor from the query options first, then fall back to the - // table options; a positive factor over-fetches indexed (approximate) - // candidates so the exact rerank below has a wider pool to reorder. Factor 0 - // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank - // path. The two option maps are kept distinct (query options passed separately - // from table options) so a broad query key cannot be overridden by a more - // specific table key: query options take precedence as a whole. `search_options` - // above is the merged view used only to drive the ANN read. - - let searches: Vec = PkVectorOrchestrator::new(reader) - .search_candidates_batch( - &plan.splits, - queries, - metric, - limit, - indexed_limit, - Some(ann_searcher), - &factory, - &search_options, - skip_exact_fallback, - row_selections_by_split.as_deref(), - concurrency, - ) - .await?; - - Ok(searches) -} - -/// Search an already-resolved plan and return one merged, best-first candidate list -/// per query: the raw layer above, followed by the optional exact rerank of the -/// approximate candidates and the merge with the exact-fallback candidates. -#[allow(clippy::too_many_arguments)] -async fn search_pk_candidates_batch_with_plan( - table: &Table, - query_options: &HashMap, - filter: Option<&Predicate>, - core: &CoreOptions<'_>, - pk_col: &str, - queries: &[&[f32]], - limit: usize, - plan: &PkVectorScanPlan, - params: &PkVectorSearchParams, -) -> crate::Result>> { - let searches = search_pk_raw_candidates_batch_with_plan( - table, - query_options, - filter, - core, - pk_col, - queries, - limit, - plan, - params, - ) - .await?; - - let metric = params.metric; - let refine_factor = params.refine_factor; - let vector_field = params.vector_field.clone(); - - // Per query: exact rerank of the approximate candidates when a refine factor is - // set (exact-fallback candidates are already exact and are not reranked), then - // merge the (possibly reranked) indexed list with the exact list into one - // best-first list bounded to the caller's limit. With no refine factor the - // rerank is a plain merge, byte-identical to the no-rerank path. Each query - // reranks its OWN indexed candidates. - let mut per_query_candidates = Vec::with_capacity(searches.len()); - for (query_index, search) in searches.into_iter().enumerate() { - let query_vector = queries[query_index]; - let indexed = if refine_factor > 0 && !search.indexed.is_empty() { - // Vector-only reader (project just the vector field); the position read - // appends _PKEY_VECTOR_POSITION itself and injects _ROW_ID internally. - let rerank_reader = DataFileReader::new( - table.file_io().clone(), - table.schema_manager().clone(), - table.schema().id(), - table.schema().fields().to_vec(), - vec![vector_field.clone()], - Vec::new(), - ); - rerank_indexed_positional( - &rerank_reader, - search.indexed, - &plan.splits, - query_vector, - metric, + /// Create an owned reader; query errors are reported before planning. + pub fn new_read(&self) -> crate::Result { + let (column, query, limit) = self.query()?; + Ok(VectorRead { + batch: BatchVectorRead::new( + self.table, + column, + &[query], limit, - &vector_field, - ) - .await? - } else { - search.indexed - }; - per_query_candidates.push(merge_candidates(indexed, search.exact, limit)); - } - - Ok(per_query_candidates) -} - -/// Plan the whole table and search it: resolve the query parameters, read the index -/// manifest into a plan, then search that plan. The plan and metric are returned -/// alongside the candidates because callers re-associate hits through the plan. -async fn plan_and_search_pk_candidates_batch( - table: &Table, - query_options: &HashMap, - filter: Option<&Predicate>, - core: &CoreOptions<'_>, - pk_col: &str, - queries: &[&[f32]], - limit: usize, -) -> crate::Result<( - Vec>, - PkVectorScanPlan, - VectorSearchMetric, -)> { - let params = resolve_pk_vector_search_params( - table, - query_options, - filter, - core, - pk_col, - queries, - limit, - )?; - let plan = PkVectorScan::new( - table, - params.field_id, - params.index_type.clone(), - filter.cloned(), - ) - .plan() - .await?; - let metric = params.metric; - let candidates = search_pk_candidates_batch_with_plan( - table, - query_options, - filter, - core, - pk_col, - queries, - limit, - &plan, - ¶ms, - ) - .await?; - Ok((candidates, plan, metric)) -} - -impl<'a> BatchVectorSearchBuilder<'a> { - pub(crate) fn new(table: &'a Table) -> Self { - Self { - table, - vector_column: None, - query_vectors: None, - limit: None, - options: HashMap::new(), - projection: None, - filter: None, - include_row_ids: None, - prepared_filter: None, - } - } - - pub fn with_vector_column(&mut self, name: &str) -> &mut Self { - self.vector_column = Some(name.to_string()); - self - } - - pub fn with_query_vectors(&mut self, vectors: Vec>) -> &mut Self { - self.query_vectors = Some(vectors); - self - } - - pub fn with_limit(&mut self, limit: usize) -> &mut Self { - self.limit = Some(limit); - self - } - - pub fn with_options(&mut self, options: HashMap) -> &mut Self { - self.options = options; - self - } - - /// Attach one scalar predicate shared by every query in the batch and applied - /// before vector Top-K. See [`VectorSearchBuilder::with_filter`] for the - /// primary-key and data-evolution execution semantics. - pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { - self.filter = Some(filter); - self.include_row_ids = None; - self.prepared_filter = None; - self - } - - /// Attach a prepared scalar pre-filter together with the exact table - /// snapshot against which its row-ID allow-list was evaluated. - pub fn with_prepared_filter( - &mut self, - prepared_filter: PreparedVectorSearchFilter, - ) -> &mut Self { - self.prepared_filter = Some(prepared_filter); - self.filter = None; - self.include_row_ids = None; - self - } - - /// Attach a caller-managed row-ID allow-list. - /// - /// This low-level API does not bind the allow-list to a table snapshot. - /// Prefer [`Self::with_prepared_filter`] for scalar pre-filters. - pub fn with_include_row_ids(&mut self, include_row_ids: RoaringTreemap) -> &mut Self { - self.include_row_ids = Some(Arc::new(include_row_ids)); - self.filter = None; - self.prepared_filter = None; - self + &self.options, + self.filter.as_ref(), + None, + None, + )?, + }) } - /// Restrict the columns materialized by [`execute_read`](Self::execute_read) to - /// `cols` (plus the always-appended `__paimon_search_score`). Without this call - /// `execute_read` materializes every user table column. Only affects - /// `execute_read`; `execute` ignores it. - pub fn with_projection(&mut self, cols: &[&str]) -> &mut Self { - self.projection = Some(cols.iter().map(|c| c.to_string()).collect()); - self + /// Search locally using the same Scan -> Plan -> Read API exposed to engines. + /// Use the result's `new_read_builder()` to materialize projected columns. + pub async fn execute(&self) -> crate::Result { + let read = self.new_read()?; + read.read(self.new_scan()?.plan().await?).await } - pub async fn execute(&self) -> crate::Result> { - let timing_enabled = vector_search_timing_enabled(); - let total_start = timing_enabled.then(Instant::now); - // The builder target is authoritative for current auth/type policy. - // A prepared filter only pins a snapshot and may carry older options. + fn query(&self) -> crate::Result<(&str, &[f32], usize)> { CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; - if let Some(prepared) = &self.prepared_filter { - if !same_vector_search_table(self.table, prepared.table()) { - return Err(crate::Error::DataInvalid { - message: format!( - "Prepared vector search filter belongs to a different table: builder target is '{}@{}', prepared filter target is '{}@{}'", - self.table.location(), - self.table.branch(), - prepared.table().location(), - prepared.table().branch(), - ), - source: None, - }); - } - } - // Check the pinned execution view as defense in depth before any fast - // path returns data-derived row ids/scores outside TableScan/TableRead. - let execution_table = self - .prepared_filter - .as_ref() - .map(PreparedVectorSearchFilter::table) - .unwrap_or(self.table); - let core = CoreOptions::new(execution_table.schema().options()); - core.ensure_read_authorized()?; - let vector_column = - self.vector_column - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Vector column must be set via with_vector_column()".to_string(), - })?; - if vector_column.is_empty() { - return Err(crate::Error::ConfigInvalid { + let column = self + .vector_column + .as_deref() + .ok_or_else(|| crate::Error::ConfigInvalid { message: "Vector column must be set via with_vector_column()".to_string(), - }); - } - - let query_vectors = - self.query_vectors - .as_ref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Query vectors must be set via with_query_vectors()".to_string(), - })?; - if query_vectors.is_empty() { - return Err(crate::Error::ConfigInvalid { - message: "Query vectors must be set via with_query_vectors()".to_string(), - }); - } - + })?; + let vector = self + .query_vector + .as_deref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Query vector must be set via with_query_vector()".to_string(), + })?; let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { message: "Limit must be set via with_limit()".to_string(), })?; - - // A primary-key vector table exposes no global row ids, so scored batch - // search is unsupported: `execute()` returns `SearchResult`s (global row - // ids). Fail loud and direct callers to the materialized batch - // `execute_read`, mirroring the single-query builder's PK guard. Membership - // is resolved via the non-erroring columns accessor so a malformed - // PK-vector config cannot abort an unrelated DE query. - if core.primary_key_vector_index_enabled() { - let targets_pk_column = core - .primary_key_vector_index_columns() - .ok() - .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); - if targets_pk_column { - return Err(crate::Error::DataInvalid { - message: "primary-key vector search does not produce global row ids; use the materialized read (execute_read) instead".to_string(), - source: None, - }); - } - } - - let mut vector_searches = query_vectors - .iter() - .map(|vector| { - VectorSearch::new(vector.clone(), limit, vector_column.to_string()) - .map(|search| search.with_options(self.options.clone())) - }) - .collect::>>()?; - - if self - .prepared_filter - .as_ref() - .is_some_and(|prepared| prepared.include_row_ids().is_empty()) - { - return Ok(vec![SearchResult::empty(); vector_searches.len()]); - } - - let snapshot_manager = execution_table.snapshot_manager(); - let setup = total_start.map_or(Duration::ZERO, |start| start.elapsed()); - - let snapshot_start = timing_enabled.then(Instant::now); - let snapshot = match crate::table::time_travel::resolve_snapshot(execution_table).await? { - Some(s) => s, - None => { - let snapshot = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); - let results = vec![SearchResult::empty(); vector_searches.len()]; - if let Some(total_start) = total_start { - let total = total_start.elapsed(); - let unattributed = total.saturating_sub(setup.saturating_add(snapshot)); - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_search_api nq={} index_entries=0 result_count=0 total_ms={:.3} setup_ms={:.3} snapshot_ms={:.3} manifest_ms=0.000 evaluate_ms=0.000 unattributed_ms={:.3}", - vector_searches.len(), - total.as_secs_f64() * 1000.0, - setup.as_secs_f64() * 1000.0, - snapshot.as_secs_f64() * 1000.0, - unattributed.as_secs_f64() * 1000.0, - ); - } - return Ok(results); - } - }; - let snapshot_elapsed = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); - let pinned_table = match &self.prepared_filter { - Some(prepared) => prepared.table().clone(), - None => { - execution_table - .copy_with_resolved_snapshot(&snapshot) - .await? - } - }; - - if let Some(prepared) = &self.prepared_filter { - for search in &mut vector_searches { - search.set_shared_include_row_ids(Arc::clone(prepared.include_row_ids())); - } - } else if let Some(include_row_ids) = &self.include_row_ids { - if include_row_ids.is_empty() { - return Ok(vec![SearchResult::empty(); vector_searches.len()]); - } - for search in &mut vector_searches { - search.set_shared_include_row_ids(Arc::clone(include_row_ids)); - } - } else if let Some(filter) = &self.filter { - let include_row_ids = matching_row_ids_for_filter(&pinned_table, filter).await?; - if include_row_ids.is_empty() { - return Ok(vec![SearchResult::empty(); vector_searches.len()]); - } - let include_row_ids = Arc::new(include_row_ids); - for search in &mut vector_searches { - search.set_shared_include_row_ids(Arc::clone(&include_row_ids)); - } - } - - let manifest_start = timing_enabled.then(Instant::now); - let index_entries = match snapshot.index_manifest() { - Some(index_manifest_name) => { - let manifest_path = snapshot_manager.manifest_path(index_manifest_name); - IndexManifest::read(execution_table.file_io(), &manifest_path).await? - } - None => Vec::new(), - }; - let manifest = manifest_start.map_or(Duration::ZERO, |start| start.elapsed()); - - let evaluate_start = timing_enabled.then(Instant::now); - let results = evaluate_batch_vector_search( - VectorSearchEvaluation { - table: Some(&pinned_table), - file_io: pinned_table.file_io(), - table_path: pinned_table.location(), - table_options: pinned_table.schema().options(), - schema_fields: pinned_table.schema().fields(), - next_row_id: snapshot.next_row_id(), - }, - &index_entries, - &vector_searches, - ) - .await?; - if let (Some(total_start), Some(evaluate_start)) = (total_start, evaluate_start) { - let total = total_start.elapsed(); - let evaluate = evaluate_start.elapsed(); - let children = setup - .saturating_add(snapshot_elapsed) - .saturating_add(manifest) - .saturating_add(evaluate); - let result_count = results - .iter() - .map(|result| result.row_ids.len()) - .sum::(); - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_search_api nq={} index_entries={} result_count={} total_ms={:.3} setup_ms={:.3} snapshot_ms={:.3} manifest_ms={:.3} evaluate_ms={:.3} unattributed_ms={:.3}", - vector_searches.len(), - index_entries.len(), - result_count, - total.as_secs_f64() * 1000.0, - setup.as_secs_f64() * 1000.0, - snapshot_elapsed.as_secs_f64() * 1000.0, - manifest.as_secs_f64() * 1000.0, - evaluate.as_secs_f64() * 1000.0, - total.saturating_sub(children).as_secs_f64() * 1000.0, - ); - } - Ok(results) + Ok((column, vector, limit)) } +} - /// Run a batch of vector searches and materialize each query's matching rows as - /// a best-first Arrow stream. Supported only for the primary-key vector path - /// (which alone can materialize physical rows). The returned `Vec` is aligned - /// strictly to the input query order and its length always equals the query - /// count — a query with no hits yields an empty stream, never a missing entry. - /// If ANY query errors (e.g. a malformed vector) the whole call fails loud with - /// no partial `Vec` of streams. Output columns are the projected user table - /// columns (all user columns by default, or those set via - /// [`with_projection`](Self::with_projection)) plus `__paimon_search_score`; - /// `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always hidden. - /// - /// A data-evolution (global-index) table fails loud: its batch search returns - /// scored global row-ids, not materialized rows, so callers use - /// [`execute`](Self::execute) instead. - pub async fn execute_read(&self) -> crate::Result> { - // Fail closed: returns data outside `TableScan`/`TableRead`. - let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; - let vector_column = - self.vector_column - .as_deref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Vector column must be set via with_vector_column()".to_string(), - })?; - let query_vectors = - self.query_vectors - .as_ref() - .ok_or_else(|| crate::Error::ConfigInvalid { - message: "Query vectors must be set via with_query_vectors()".to_string(), - })?; - if query_vectors.is_empty() { - return Err(crate::Error::ConfigInvalid { - message: "Query vectors must be set via with_query_vectors()".to_string(), - }); - } - let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { - message: "Limit must be set via with_limit()".to_string(), - })?; - - // Only the primary-key vector path can materialize rows. The data-evolution - // (global-index) path returns data-derived row-ids, not table rows, so a - // batch read against it (or a non-PK-vector column) fails loud, directing - // callers to `execute()`. - let targets_pk_column = core.primary_key_vector_index_enabled() - && core - .primary_key_vector_index_columns() - .ok() - .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); - if !targets_pk_column { - return Err(crate::Error::DataInvalid { - message: "batch vector read is only supported on the primary-key vector path; data-evolution batch search returns scored row ids, use execute() instead".to_string(), - source: None, - }); - } - - let pk_col = core.primary_key_vector_index_column()?; - let query_refs: Vec<&[f32]> = query_vectors.iter().map(|q| q.as_slice()).collect(); - - // Resolve the materialization read-type up front so an invalid projection - // (unknown column, or a reserved metadata / row-id name) fails loud - // unconditionally, before any read — a whole-call failure, not a partial - // Vec. - let read_type = self.resolve_materialize_read_type()?; - - // One shared plan / lazy segment loader / residual across all N queries; the - // per-query candidate lists come back in strict input order. Any query - // error (or a shared-plan error) propagates here, so no partial Vec is - // returned. - let (per_query_candidates, plan, metric) = plan_and_search_pk_candidates_batch( - self.table, - &self.options, - self.filter.as_ref(), - &core, - &pk_col, - &query_refs, - limit, - ) - .await?; - - let materialize_reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - read_type, - Vec::new(), - ); - - // Materialize each query's candidates into its own stream, preserving - // arity: an empty candidate list yields an empty stream. Build every stream - // before returning so a materialization error fails the whole call with no - // partial Vec. - let mut streams = Vec::with_capacity(per_query_candidates.len()); - for candidates in per_query_candidates { - streams.push( - VectorSearchBuilder::materialize_candidates( - candidates, - &plan.splits, - metric, - &materialize_reader, - ) - .await?, - ); - } - Ok(streams) - } - - /// Resolve the projected fields for the materialization read-type. Default - /// (no projection set) is all user table fields; otherwise the requested names - /// resolved via `resolve_projected_fields`. Rejects reserved metadata names and - /// `_ROW_ID` so a user cannot request a hidden column. Mirrors the single - /// builder's resolver. - fn resolve_materialize_read_type(&self) -> crate::Result> { - let fields = match &self.projection { - None => self.table.schema().fields().to_vec(), - Some(names) => { - for name in names { - if is_reserved_read_column(name) { - return Err(crate::Error::DataInvalid { - message: format!( - "vector search read projection must not request reserved column '{name}'" - ), - source: None, - }); - } - } - resolve_projected_fields( - self.table.identifier().full_name(), - self.table.schema().fields(), - names, - true, - )? - } - }; - // The default projection returns every user column, so a user column - // whose name collides with an injected metadata column must be rejected - // on the resolved field list too — not only when explicitly requested. - ensure_no_reserved_read_columns(&fields)?; - Ok(fields) - } -} - -#[derive(Clone, Copy)] -struct VectorSearchEvaluation<'a> { - table: Option<&'a Table>, - file_io: &'a FileIO, - table_path: &'a str, - table_options: &'a HashMap, - schema_fields: &'a [DataField], - next_row_id: Option, -} - -async fn matching_row_ids_for_filter( - table: &Table, - filter: &Predicate, -) -> crate::Result { - let mut read_builder = table.new_read_builder(); - read_builder - .with_projection(&[ROW_ID_FIELD_NAME])? - .with_filter(filter.clone()); - let plan = read_builder.new_scan().plan().await?; - let read = read_builder.new_read()?; - let mut stream = read.to_arrow(plan.splits())?; - let mut row_ids = RoaringTreemap::new(); - while let Some(batch) = stream.try_next().await? { - let index = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|_| crate::Error::DataInvalid { - message: format!( - "scalar vector pre-filter read is missing {ROW_ID_FIELD_NAME}" - ), - source: None, - })?; - let values = batch - .column(index) - .as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "scalar vector pre-filter {ROW_ID_FIELD_NAME} column is not Int64" - ), - source: None, - })?; - for row in 0..values.len() { - if values.is_null(row) { - return Err(crate::Error::DataInvalid { - message: format!( - "scalar vector pre-filter produced a null {ROW_ID_FIELD_NAME}" - ), - source: None, - }); - } - let row_id = values.value(row); - let row_id = u64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { - message: format!( - "scalar vector pre-filter produced a negative {ROW_ID_FIELD_NAME}: {row_id}" - ), - source: None, - })?; - row_ids.insert(row_id); - } - } - Ok(row_ids) -} - -impl Table { - /// Resolve a scalar predicate once and pin all later vector-search/read - /// stages to the same snapshot. - pub async fn prepare_vector_search_filter( - &self, - filter: Predicate, - ) -> crate::Result { - CoreOptions::new(self.schema().options()).ensure_read_authorized()?; - let Some(snapshot) = crate::table::time_travel::resolve_snapshot(self).await? else { - return Ok(PreparedVectorSearchFilter { - table: self.clone(), - include_row_ids: Arc::new(RoaringTreemap::new()), - }); - }; - let table = self.copy_with_resolved_snapshot(&snapshot).await?; - let include_row_ids = matching_row_ids_for_filter(&table, &filter).await?; - Ok(PreparedVectorSearchFilter { - table, - include_row_ids: Arc::new(include_row_ids), - }) - } -} - -#[derive(Default)] -struct IndexSearchTiming { - permit_wait: Duration, - file_reader_open: Duration, -} - -#[cfg(test)] -async fn evaluate_vector_search( - evaluation: VectorSearchEvaluation<'_>, - index_entries: &[IndexManifestEntry], - vector_search: &VectorSearch, -) -> crate::Result> { - let results = evaluate_batch_vector_search( - evaluation, - index_entries, - std::slice::from_ref(vector_search), - ) - .await?; - take_only_result(results, "vector search")?.to_row_ranges() -} - -async fn evaluate_batch_vector_search( - evaluation: VectorSearchEvaluation<'_>, - index_entries: &[IndexManifestEntry], - vector_searches: &[VectorSearch], -) -> crate::Result> { - let timing_enabled = vector_search_timing_enabled(); - let total_start = timing_enabled.then(Instant::now); - if vector_searches.is_empty() { - return Ok(Vec::new()); - } - - let table_path = evaluation.table_path.trim_end_matches('/'); - let core_options = CoreOptions::new(evaluation.table_options); - let search_mode = core_options.vector_index_search_mode()?; - let field_name = &vector_searches[0].field_name; - if vector_searches - .iter() - .any(|vector_search| vector_search.field_name != *field_name) - { - return Err(crate::Error::DataInvalid { - message: "Batch vector search requires all query vectors to use the same field" - .to_string(), - source: None, - }); - } - let search_options = vector_searches[0].options.clone(); - if vector_searches - .iter() - .any(|vector_search| vector_search.options != search_options) - { - return Err(crate::Error::DataInvalid { - message: "Batch vector search requires all query vectors to use the same options" - .to_string(), - source: None, - }); - } - - let field_id = match find_field_id_by_name(evaluation.schema_fields, field_name) { - Some(id) => id, - None => return Ok(vec![SearchResult::empty(); vector_searches.len()]), - }; - - let vector_entries: Vec<_> = index_entries - .iter() - .filter(|e| { - e.kind == FileKind::Add - && VectorIndexBackend::from_index_type(&e.index_file.index_type).is_some() - && e.index_file - .global_index_meta - .as_ref() - .is_some_and(|m| m.index_field_id == field_id) - }) - .collect(); - - if vector_entries.is_empty() && search_mode == GlobalIndexSearchMode::Fast { - return Ok(vec![SearchResult::empty(); vector_searches.len()]); - } - - let deletion_vector_start = timing_enabled.then(Instant::now); - let deleted_row_index = if core_options.data_evolution_enabled() { - match evaluation.table { - Some(table) => { - let ranges = - deleted_row_ranges_for_data_evolution_dvs(table, index_entries).await?; - (!ranges.is_empty()).then(|| RowRangeIndex::create(ranges)) - } - None => None, - } - } else { - None - }; - let deletion_vector = deletion_vector_start.map_or(Duration::ZERO, |start| start.elapsed()); - - let max_limit = vector_searches - .iter() - .map(|vector_search| vector_search.limit) - .max() - .unwrap_or(0); - let refine_factor = match vector_entries.first() { - Some(entry) => configured_refine_factor( - &search_options, - evaluation.table_options, - field_name, - &entry.index_file.index_type, - )?, - None => 0, - }; - let index_search_limit = indexed_search_limit(max_limit, refine_factor)?; - - let vector_entry_count = vector_entries.len(); - let vector_search_plans = if let Some(include_row_ids) = - shared_batch_include_row_ids(vector_searches) - { - let ranges = vector_entries - .iter() - .map(|entry| { - let meta = entry.index_file.global_index_meta.as_ref().ok_or_else(|| { - crate::Error::DataInvalid { - message: format!( - "Vector index '{}' is missing global index metadata", - entry.index_file.file_name - ), - source: None, - } - })?; - Ok((meta.row_range_start, meta.row_range_end)) - }) - .collect::>>()?; - vector_entries - .iter() - .copied() - .zip(localize_shared_include_row_ids( - include_row_ids.as_ref(), - &ranges, - )?) - .filter_map(|(entry, local_filter)| local_filter.map(|filter| (entry, Some(filter)))) - .collect::>() - } else { - vector_entries - .iter() - .copied() - .map(|entry| (entry, None)) - .collect::>() - }; - let mut permit_wait = Duration::ZERO; - let mut file_reader_open = Duration::ZERO; - let mut index_search = Duration::ZERO; - let mut merge = Duration::ZERO; - let mut refine = Duration::ZERO; - let mut raw_fallback = Duration::ZERO; - let mut merged = vec![SearchResult::empty(); vector_searches.len()]; - if !vector_entries.is_empty() { - let index_search_start = timing_enabled.then(Instant::now); - let concurrency = core_options.global_index_thread_num()?; - if concurrency > tokio::sync::Semaphore::MAX_PERMITS { - return Err(crate::Error::DataInvalid { - message: format!( - "Global index thread count must not exceed {}", - tokio::sync::Semaphore::MAX_PERMITS - ), - source: None, - }); - } - ensure_global_index_executor_capacity(concurrency); - let vindex_entry_count = vector_entries - .iter() - .filter(|entry| is_vindex_index_type(&entry.index_file.index_type)) - .count(); - let (batch_index_parallelism, range_read_limiter) = if vindex_entry_count == 0 { - (1, None) - } else { - let (index_parallelism, range_read_concurrency) = - vindex_concurrency_limits(&core_options, vindex_entry_count, concurrency)?; - ( - index_parallelism, - Some(RangeReadLimiter::new(range_read_concurrency)), - ) - }; - let futures: Vec<_> = vector_search_plans - .into_iter() - .map(|(entry, shared_local_filter)| { - let range_read_limiter = range_read_limiter.clone(); - 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 = IndexFileLocation::Global { table_path } - .resolve(&entry.index_file.file_name, entry.index_file.external_path.as_deref()); - let file_name = entry.index_file.file_name.clone(); - let file_size = entry.index_file.file_size as u64; - let index_meta_bytes = global_meta.index_meta.clone().unwrap_or_default(); - let row_range_start = global_meta.row_range_start; - let row_range_end = global_meta.row_range_end; - let index_limit = search_limit_with_deleted_rows( - index_search_limit, - row_range_start, - row_range_end, - deleted_row_index.as_ref(), - ) - .min(i32::MAX as usize); - let mut vector_searches = vector_searches.to_vec(); - for vector_search in &mut vector_searches { - vector_search.limit = index_limit; - } - let mut options = evaluation.table_options.clone(); - options.extend(search_options.clone()); - let input = evaluation.file_io.new_input(&path); - async move { - if let Some(local_filter) = shared_local_filter { - let local_filter = Arc::new(local_filter); - for vector_search in &mut vector_searches { - vector_search - .set_shared_include_row_ids(Arc::clone(&local_filter)); - } - } else { - for vector_search in &mut vector_searches { - if let Some(include_row_ids) = - vector_search.effective_include_row_ids() - { - vector_search.set_shared_include_row_ids(Arc::new( - localize_include_row_ids( - include_row_ids, - row_range_start, - row_range_end, - )?, - )); - } - } - } - if vector_searches.iter().all(|search| { - search - .effective_include_row_ids() - .is_some_and(|row_ids| row_ids.is_empty()) - }) { - return Ok(( - vec![SearchResult::empty(); vector_searches.len()], - IndexSearchTiming::default(), - )); - } - let permit_start = timing_enabled.then(Instant::now); - let permit = acquire_process_global_search_permit(concurrency).await?; - let permit_wait = - permit_start.map_or(Duration::ZERO, |start| start.elapsed()); - let input = input?; - let query_count = vector_searches.len(); - let mut file_reader_open = Duration::ZERO; - let mut full_file_read = None; - let io_meta = - GlobalIndexIOMeta::new(file_name.clone(), file_size, index_meta_bytes); - let results = match backend { - VectorIndexBackend::Lumina => { - let read_start = timing_enabled.then(Instant::now); - let data = input.read().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to read {} index file '{}': {}", - backend.error_name(), - file_name, - e - ), - source: None, - } - })?; - if let Some(start) = read_start { - full_file_read = Some((start.elapsed(), data.len())); - } - execute_global_index_with_guard( - "Lumina global-index batch search task failed", - permit, - move || { - let mut reader = - LuminaVectorGlobalIndexReader::new(io_meta, options); - reader.visit_batch_vector_search(&vector_searches, |_| { - Ok(Cursor::new(data)) - }) - }, - ) - .await? - } - VectorIndexBackend::Vindex => { - match tokio::runtime::Handle::try_current() { - Ok(runtime) => { - let file_reader_open_start = - timing_enabled.then(Instant::now); - let file_reader = input.reader().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to open vindex file '{}' for range reads: {}", - file_name, e - ), - source: None, - } - })?; - file_reader_open = file_reader_open_start - .map_or(Duration::ZERO, |start| start.elapsed()); - let source = VindexFileReader::new_with_limiter( - Arc::new(file_reader), - runtime, - range_read_limiter.expect("Vindex range-read limiter"), - file_size, - file_name.clone(), - ); - let range_io_stats = source.range_io_stats(); - let results = execute_vindex_searches( - io_meta, - options, - vector_searches, - source, - file_name.clone(), - batch_index_parallelism, - permit, - ) - .await?; - if let Some(stats) = range_io_stats { - log_vindex_range_io_stats( - &file_name, - query_count, - &stats, - ); - } - results - } - Err(_) if query_count > 1 => { - let read_start = timing_enabled.then(Instant::now); - let data = input.read().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to read vindex index file '{}': {}", - file_name, e - ), - source: None, - } - })?; - if let Some(start) = read_start { - full_file_read = Some((start.elapsed(), data.len())); - } - execute_vindex_searches( - io_meta, - options, - vector_searches, - Cursor::new(data), - file_name.clone(), - batch_index_parallelism, - permit, - ) - .await? - } - Err(error) => { - return Err(crate::Error::UnexpectedError { - message: - "Vector index range reader requires a Tokio runtime" - .to_string(), - source: Some(Box::new(error)), - }); - } - } - } - }; - if let Some((read, returned_bytes)) = full_file_read { - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_full_file_io backend={} file={} nq={} requested_bytes={} returned_bytes={} read_ms={:.3}", - backend.error_name(), - file_name, - query_count, - file_size, - returned_bytes, - read.as_secs_f64() * 1000.0, - ); - } - if results.len() != query_count { - return Err(crate::Error::DataInvalid { - message: format!( - "Batch vector search backend returned {} results for {} query vectors", - results.len(), - query_count - ), - source: None, - }); - } - - Ok::<_, crate::Error>(( - results - .into_iter() - .map(|result| match result { - Some(scored_map) => SearchResult::from_scored_map(scored_map) - .offset(row_range_start), - None => SearchResult::empty(), - }) - .collect::>(), - IndexSearchTiming { - permit_wait, - file_reader_open, - }, - )) - } - }) - .collect(); - - let results = drain_indexed_jobs(futures.into_iter(), concurrency).await?; - index_search = index_search_start.map_or(Duration::ZERO, |start| start.elapsed()); - let merge_start = timing_enabled.then(Instant::now); - for (per_entry, entry_timing) in &results { - permit_wait = permit_wait.saturating_add(entry_timing.permit_wait); - file_reader_open = file_reader_open.saturating_add(entry_timing.file_reader_open); - for (query_index, result) in per_entry.iter().enumerate() { - merged[query_index] = merged[query_index].or(result); - } - } - merge = merge_start.map_or(Duration::ZERO, |start| start.elapsed()); - } - - if refine_factor != 0 { - let refine_start = timing_enabled.then(Instant::now); - merged = maybe_rerank_indexed_batch_results( - evaluation, - index_entries, - field_id, - field_name, - vector_searches, - merged, - index_search_limit, - ) - .await?; - refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); - } - - if search_mode != GlobalIndexSearchMode::Fast { - let raw_fallback_start = timing_enabled.then(Instant::now); - let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail { - let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { - message: "Vector raw search in detail mode requires table context".to_string(), - source: None, - })?; - detail_data_ranges_for_table(table).await? - } else { - Vec::new() - }; - let field_ids = HashSet::from([field_id]); - let raw_ranges = unindexed_ranges_for_global_index_entries( - index_entries, - &field_ids, - search_mode, - evaluation.next_row_id, - &detail_ranges, - is_vector_global_index_file, - ); - if !raw_ranges.is_empty() { - let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { - message: "Vector raw search requires table context".to_string(), - source: None, - })?; - let metric_start = timing_enabled.then(Instant::now); - let metric = resolve_raw_vector_metric( - evaluation.file_io, - table_path, - evaluation.table_options, - index_entries, - field_id, - field_name, - ) - .await?; - let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); - let (raw_results, raw_timing) = - read_raw_batch_vector_search(table, vector_searches, &raw_ranges, metric).await?; - if let Some(raw_timing) = raw_timing { - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_raw_fallback nq={} row_ranges={} metric_resolve_ms={:.3} raw_plan_ms={:.3} split_count={} file_count={} raw_stream_wait_ms={:.3} raw_score_cpu_ms={:.3} arrow_batches={} arrow_rows={} total_raw_read_ms={:.3}", - vector_searches.len(), - raw_ranges.len(), - metric_resolve.as_secs_f64() * 1000.0, - raw_timing.plan.as_secs_f64() * 1000.0, - raw_timing.split_count, - raw_timing.file_count, - raw_timing.stream_wait.as_secs_f64() * 1000.0, - raw_timing.score_cpu.as_secs_f64() * 1000.0, - raw_timing.batch_count, - raw_timing.row_count, - raw_timing.total.as_secs_f64() * 1000.0, - ); - } - for (query_index, result) in raw_results.iter().enumerate() { - merged[query_index] = merged[query_index].or(result); - } - } - raw_fallback = raw_fallback_start.map_or(Duration::ZERO, |start| start.elapsed()); - } - - let finalize_start = timing_enabled.then(Instant::now); - let results = merged - .into_iter() - .zip(vector_searches) - .map(|(result, vector_search)| { - Ok(result - .without_deleted_row_ranges(deleted_row_index.as_ref())? - .top_k(vector_search.limit)) - }) - .collect::>>()?; - let finalize = finalize_start.map_or(Duration::ZERO, |start| start.elapsed()); - if let Some(total_start) = total_start { - let total = total_start.elapsed(); - let children = deletion_vector - .saturating_add(index_search) - .saturating_add(merge) - .saturating_add(refine) - .saturating_add(raw_fallback) - .saturating_add(finalize); - let result_count = results - .iter() - .map(|result| result.row_ids.len()) - .sum::(); - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_search_evaluate nq={} index_entries={} index_files={} result_count={} refine_factor={} total_ms={:.3} deletion_vector_ms={:.3} index_search_ms={:.3} global_permit_wait_sum_ms={:.3} file_reader_open_sum_ms={:.3} merge_ms={:.3} refine_ms={:.3} raw_fallback_ms={:.3} finalize_ms={:.3} unattributed_ms={:.3}", - vector_searches.len(), - index_entries.len(), - vector_entry_count, - result_count, - refine_factor, - total.as_secs_f64() * 1000.0, - deletion_vector.as_secs_f64() * 1000.0, - index_search.as_secs_f64() * 1000.0, - permit_wait.as_secs_f64() * 1000.0, - file_reader_open.as_secs_f64() * 1000.0, - merge.as_secs_f64() * 1000.0, - refine.as_secs_f64() * 1000.0, - raw_fallback.as_secs_f64() * 1000.0, - finalize.as_secs_f64() * 1000.0, - total.saturating_sub(children).as_secs_f64() * 1000.0, - ); - } - Ok(results) -} - -fn is_vector_global_index_file(index_file: &IndexFileMeta) -> bool { - VectorIndexBackend::from_index_type(&index_file.index_type).is_some() -} - -/// Compute, per data file in `split`, the set of file-LOCAL physical row -/// positions whose rows satisfy the residual predicate. Mirrors the -/// row-collecting half of Java `PrimaryKeyVectorRead`'s `executeFilter`: the -/// predicate is NOT pushed down (a pushed filter would drop rows before their -/// position could be recovered). Instead `reader` projects only the residual -/// columns and carries no pushdown predicate, the residual is evaluated here at the -/// Arrow level, and each surviving row's file-local 0-based position is recovered -/// from the selection the read was limited to. This needs no `_ROW_ID` and no -/// `first_row_id` — real primary-key tables never write one. -/// -/// `allowed_rows` is the plan's per-file physical selection, keyed by data-file -/// name, with the plan's three states: a file it does not list is unrestricted and -/// the whole file is scanned; an empty range list excludes the file, which is -/// registered empty without a read; a non-empty list is scanned over exactly those -/// ranges, because an engine-supplied bucket split can restrict a huge file to a -/// handful of ranges and reading all of it to discard the rest would defeat the -/// split. -/// -/// Every *active* data file in the split gets an entry in the RESULT, possibly -/// empty, and that exhaustiveness is load-bearing. The search kernel reads a file's -/// absence from its selections as "unrestricted", so an active file missing here -/// would reach the search with no predicate applied at all -- the residual would be -/// silently dropped for it. (The merge below reads a residual's silence about a -/// file the PLAN listed as exclusion, so only a file both omit falls through, which -/// is exactly the case this exhaustiveness rules out.) Non-active files (e.g. -/// level-0 files the bucket search excludes) are skipped entirely: they are never -/// searched, so re-reading them would be wasted IO. -/// -/// `reader` must be predicate-free and project the residual columns; -/// `residual.file_fields` are the fields the residual leaf indices point into -/// (resolved by name against each emitted batch). -async fn residual_positions_by_file( - reader: &DataFileReader, - split: &DataSplit, - active_files: &[BucketActiveFile], - residual: &FilePredicates, - allowed_rows: Option<&HashMap>>, -) -> crate::Result> { - let scan_fields = reader.read_type().to_vec(); - let active_names: HashSet<&str> = active_files.iter().map(|f| f.file_name.as_str()).collect(); - let mut out: HashMap = HashMap::new(); - for file_meta in split.data_files() { - // Only files the bucket search actually recalls from need residual - // positions; skip everything else to avoid a wasted read. - if !active_names.contains(file_meta.file_name.as_str()) { - continue; - } - // A file the plan lists an EMPTY range list for permits nothing; registering - // it empty says so and costs no read. A file the plan does not list at all - // is unrestricted, so the residual is evaluated over the whole file. - let selection = match allowed_rows.and_then(|by_file| by_file.get(&file_meta.file_name)) { - Some(ranges) if ranges.is_empty() => { - out.entry(file_meta.file_name.clone()).or_default(); - continue; - } - Some(ranges) => Some(ranges.clone()), - None => None, - }; - let data_fields = reader.derive_data_fields(file_meta).await?; - let mut stream = match selection.clone() { - Some(ranges) => reader.read_single_file_stream_local_ranges( - split, - file_meta.clone(), - data_fields, - None, - ranges, - )?, - None => { - reader.read_single_file_stream(split, file_meta.clone(), data_fields, None, None)? - } - }; - // Register the file up front so a file whose rows all fail the residual - // still appears in the map (empty set). - let positions = out.entry(file_meta.file_name.clone()).or_default(); - // Rows arrive in ascending physical order, and the read emitted exactly what - // was selected (no pushdown predicate, no deletion vector), so walking the - // selection in step with the rows recovers each row's file-local position. - let mut selected: Box + Send> = match &selection { - Some(ranges) => Box::new( - ranges - .clone() - .into_iter() - .flat_map(|range| (range.from() as u64)..=(range.to() as u64)), - ), - None => Box::new(0..file_meta.row_count.max(0) as u64), - }; - while let Some(batch) = stream.try_next().await? { - let num_rows = batch.num_rows(); - let mask = evaluate_predicates_mask( - &batch, - &residual.predicates, - &residual.file_fields, - &scan_fields, - )?; - for row_index in 0..num_rows { - let position = selected.next().ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "residual scan of '{}' emitted more rows than the selection allows", - file_meta.file_name - ), - source: None, - })?; - let keep = match &mask { - // NULL follows the same NULL -> false convention the Arrow filter - // kernel applies, so a null mask slot drops the row. - Some(mask) => mask.is_valid(row_index) && mask.value(row_index), - // No predicate contributed a mask (identity) -> keep every row. - None => true, - }; - if keep { - positions.insert(position); - } - } - } - if selected.next().is_some() { - return Err(crate::Error::DataInvalid { - message: format!( - "residual scan of '{}' emitted fewer rows than the selection allows", - file_meta.file_name - ), - source: None, - }); - } - } - Ok(out) -} - -fn verify_segment_metric( - configured: VectorSearchMetric, - segment_metric: VectorSearchMetric, -) -> crate::Result<()> { - if segment_metric != configured { - return Err(crate::Error::DataInvalid { - message: format!( - "ANN segment metric {} does not match configured metric {}", - segment_metric.as_str(), - configured.as_str() - ), - source: None, - }); - } - Ok(()) -} - -fn pk_vector_query_dimension( - table_options: &HashMap, - query_options: &HashMap, - index_type: &str, - vector_field: &DataField, -) -> crate::Result> { - match vector_field.data_type() { - DataType::Vector(vector_type) - if matches!(vector_type.element_type(), DataType::Float(_)) => - { - Ok(Some(vector_type.length() as usize)) - } - DataType::Array(array_type) if matches!(array_type.element_type(), DataType::Float(_)) => { - // Resolve the dimension per the configured backend. An `ARRAY` - // column carries no dimension in its type, so it comes from options — - // but the option shape differs by backend. Lumina is not a vindex - // index type, so routing it through `VindexVectorIndexOptions` would - // reject it as unsupported before planning (even on an empty table). - if is_lumina_index_type(index_type) { - // Lumina reads `lumina.index.dimension` (default 128) from the - // merged table+query options, matching `resolve_lumina_options`. - let mut merged = table_options.clone(); - merged.extend(query_options.clone()); - let dimension = LuminaVectorIndexOptions::new(&merged)?.dimension; - Ok(Some(dimension as usize)) - } else { - let mut dimension_options = HashMap::new(); - for key in [ - "dimension".to_string(), - format!("{index_type}.dimension"), - format!("fields.{}.dimension", vector_field.name()), - ] { - if let Some(value) = query_options.get(&key) { - dimension_options.insert(key, value.clone()); - } - } - Ok(Some( - VindexVectorIndexOptions::new( - table_options, - &dimension_options, - index_type, - vector_field, - )? - .dimension(), - )) - } - } - _ => Ok(None), - } -} - -/// Rerank approximate (indexed) candidates by rereading ONLY their candidate -/// positions and recomputing the exact distance, then keep the best `limit`. -/// -/// Unlike a whole-column preload, this reuses [`PkVectorPositionRead`] to read -/// just the selected physical rows of each hit file (positions -> row ranges -> -/// local ranges), so a rerank over a large ANN-covered file touches only the -/// candidate rows. Mirrors Java's IndexedSplit rerank. -/// -/// Each returned row is matched back to its candidate by the -/// `_PKEY_VECTOR_POSITION` column VALUE (never batch order). The recomputed -/// distance is written into the ORIGINAL candidate so `split_index` / -/// partition / bucket survive (`build_indexed_splits` does not carry -/// `split_index`). A DV loaded exactly as [`PkVectorIndexedSplitRead::read`] -/// does drops deleted positions, so a candidate at a deleted position returns no -/// row and trips the leftover guard — a deleted candidate reaching rerank is a -/// real inconsistency (the search path already DV-filters), so fail loud. -#[allow(clippy::too_many_arguments)] -async fn rerank_indexed_positional( - rerank_reader: &DataFileReader, - indexed: Vec, - plan_splits: &[PkVectorSearchSplit], - query_vector: &[f32], - metric: VectorSearchMetric, - limit: usize, - vector_field: &DataField, -) -> crate::Result> { - // Original per-position candidates keyed by (split_index, file, position); - // the recomputed distance is written back into these so split_index and - // partition/bucket survive (build_indexed_splits does not carry split_index). - let mut by_key: HashMap<(usize, String, i64), PkVectorCandidate> = HashMap::new(); - for c in &indexed { - if by_key - .insert( - (c.split_index, c.data_file_name.clone(), c.row_position), - c.clone(), - ) - .is_some() - { - return Err(crate::Error::DataInvalid { - message: "duplicate primary-key vector candidate for reranking".to_string(), - source: None, - }); - } - } - - // Rebuild the split_index lookup by (partition bytes, bucket, file): the - // indexed split exposes partition/bucket/file but not split_index. - let mut split_index_of: HashMap<(Vec, i32, String), usize> = HashMap::new(); - for (i, s) in plan_splits.iter().enumerate() { - let p = s.data_split.partition().to_serialized_bytes(); - let b = s.data_split.bucket(); - for f in s.data_split.data_files() { - split_index_of.insert((p.clone(), b, f.file_name.clone()), i); - } - } - - // Every candidate must reference a (partition, bucket, file) that the plan - // actually carries. Checking up front — before build_indexed_splits, which - // indexes plan_splits by split_index — turns an absent file into a fail-loud - // error rather than an out-of-range panic, and keeps the per-split lookup - // below a self-consistent backstop. - for c in &indexed { - let key = ( - c.partition.to_serialized_bytes(), - c.bucket, - c.data_file_name.clone(), - ); - if !split_index_of.contains_key(&key) { - return Err(crate::Error::DataInvalid { - message: format!("rerank split for {} not found in plan", c.data_file_name), - source: None, - }); - } - } - - // Group the candidates into per-file indexed splits (position ranges + file - // meta), reusing the exact grouping/validation the materialization path uses. - let indexed_splits = build_indexed_splits(indexed, plan_splits, metric)?; - - let dimension = query_vector.len(); - let mut reranked: Vec = Vec::new(); - for split in indexed_splits { - let data_split = split.split.clone(); - let file_meta = data_split.data_files()[0].clone(); - let file_name = file_meta.file_name.clone(); - let partition_bytes = data_split.partition().to_serialized_bytes(); - let bucket = data_split.bucket(); - let split_index = *split_index_of - .get(&(partition_bytes, bucket, file_name.clone())) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("rerank split for {file_name} not found in plan"), - source: None, - })?; - - // DV loaded exactly as PkVectorIndexedSplitRead::read does; skipping it - // would score deleted rows. - let dv_factory = rerank_reader.build_split_dv_factory(&data_split).await?; - let dv = DataFileReader::deletion_vector_for_file(dv_factory.as_ref(), &file_name); - let data_fields = rerank_reader.derive_data_fields(&file_meta).await?; - - // Positions from the split's row_ranges (ascending); read only those. - let positions = expand_ranges(&split.row_ranges, file_meta.row_count)?; - let mut stream = PkVectorPositionRead::new(rerank_reader).read( - &data_split, - file_meta, - data_fields, - dv, - positions, - None, // no scores; rerank recomputes distance - )?; - - while let Some(batch) = stream.try_next().await? { - let pos_idx = batch - .schema() - .index_of(PKEY_VECTOR_POSITION_COLUMN) - .map_err(|_| crate::Error::DataInvalid { - message: format!("rerank batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), - source: None, - })?; - let pos_col = batch - .column(pos_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is not Int64"), - source: None, - })?; - let mut vectors: Vec>> = Vec::new(); - append_batch_vectors(&batch, vector_field.name(), dimension, &mut vectors)?; - for (row, vector) in vectors.iter().enumerate() { - let position = pos_col.value(row); - let mut candidate = by_key - .remove(&(split_index, file_name.clone(), position)) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("rerank read unexpected position {file_name}@{position}"), - source: None, - })?; - let vector = vector.as_ref().ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "primary-key vector candidate {file_name}@{position} contains a null vector" - ), - source: None, - })?; - candidate.distance = metric.compute_distance(query_vector, vector); - reranked.push(candidate); - } - } - } - - if !by_key.is_empty() { - return Err(crate::Error::DataInvalid { - message: format!( - "failed to read {} primary-key vector candidate(s) for reranking", - by_key.len() - ), - source: None, - }); - } - - Ok(merge_candidates(reranked, Vec::new(), limit)) -} - -/// One materialized row tagged with its best-first `rank` and its `(batch_index, -/// row_index)` location in the retained materialization batches. -pub(crate) struct RankedRow { - rank: usize, - batch_index: usize, - row_index: usize, -} - -/// For each row in a materialized batch, look up its best-first rank via the -/// `(partition bytes, bucket, file, position)` key and record its location. The -/// `_PKEY_VECTOR_POSITION` column supplies the physical position; every row must -/// map to a candidate rank (the batch came from that candidate's file), so a miss -/// fails loud rather than silently dropping a row. -#[allow(clippy::too_many_arguments)] -pub(crate) fn collect_ranked_rows( - batch: &RecordBatch, - batch_index: usize, - partition_bytes: &[u8], - bucket: i32, - file_name: &str, - rank_of: &HashMap<(Vec, i32, String, i64), usize>, - out: &mut Vec, -) -> crate::Result<()> { - let position_idx = batch - .schema() - .index_of(PKEY_VECTOR_POSITION_COLUMN) - .map_err(|_| crate::Error::DataInvalid { - message: format!("materialized batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), - source: None, - })?; - let positions = batch - .column(position_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is not Int64"), - source: None, - })?; - for row_index in 0..batch.num_rows() { - let position = positions.value(row_index); - let key = ( - partition_bytes.to_vec(), - bucket, - file_name.to_string(), - position, - ); - let rank = *rank_of.get(&key).ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "materialized row (file {file_name}, position {position}) has no matching search candidate" - ), - source: None, - })?; - out.push(RankedRow { - rank, - batch_index, - row_index, - }); - } - Ok(()) -} - -/// Reorder the materialized rows into best-first order and drop the internal -/// `_PKEY_VECTOR_POSITION` column, yielding a single output batch (empty input -/// yields no batches). The projected user columns and `__paimon_search_score` are -/// retained. -pub(crate) fn reorder_and_strip_position( - batches: &[RecordBatch], - mut ranked: Vec, -) -> crate::Result> { - if ranked.is_empty() { - return Ok(Vec::new()); - } - ranked.sort_by_key(|r| r.rank); - let indices: Vec<(usize, usize)> = ranked - .iter() - .map(|r| (r.batch_index, r.row_index)) - .collect(); - let refs: Vec<&RecordBatch> = batches.iter().collect(); - let reordered = - interleave_record_batch(&refs, &indices).map_err(|e| crate::Error::DataInvalid { - message: format!("failed to reorder vector search read rows: {e}"), - source: None, - })?; - - // Drop the internal position column; keep every other column (projected user - // columns + __paimon_search_score) in order. - let position_idx = reordered - .schema() - .index_of(PKEY_VECTOR_POSITION_COLUMN) - .map_err(|_| crate::Error::DataInvalid { - message: format!("reordered batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), - source: None, - })?; - let keep: Vec = (0..reordered.num_columns()) - .filter(|i| *i != position_idx) - .collect(); - let projected = reordered - .project(&keep) - .map_err(|e| crate::Error::DataInvalid { - message: format!("failed to drop position column: {e}"), - source: None, - })?; - Ok(vec![projected]) -} - -/// Collect materialized DE rows, join each row's `(rank, score)` by its global -/// `_ROW_ID`, reorder to the search rank order, append the `__paimon_search_score` -/// column, and drop `_ROW_ID`. Every row must map to a search candidate and the -/// total materialized count must equal `expected_len`; a miss or count mismatch -/// fails loud rather than silently dropping or NaN-scoring a row. Empty input -/// yields no batches. -fn attach_scores_by_row_id( - batches: &[RecordBatch], - rank_score_of: &HashMap, - expected_len: usize, -) -> crate::Result> { - // (rank, batch_index, row_index, score) per materialized row. - let mut ranked: Vec<(usize, usize, usize, f32)> = Vec::new(); - for (batch_index, batch) in batches.iter().enumerate() { - let row_id_idx = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|_| crate::Error::DataInvalid { - message: format!("materialized batch missing {ROW_ID_FIELD_NAME} column"), - source: None, - })?; - let col = batch.column(row_id_idx); - let ids = - col.as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("{ROW_ID_FIELD_NAME} column is not Int64"), - source: None, - })?; - for row_index in 0..batch.num_rows() { - if ids.is_null(row_index) { - return Err(crate::Error::DataInvalid { - message: format!( - "materialized DE vector row has null {ROW_ID_FIELD_NAME}; cannot align score" - ), - source: None, - }); - } - let id = ids.value(row_index); - let (rank, score) = - *rank_score_of - .get(&id) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "materialized DE vector row (row id {id}) has no matching search candidate" - ), - source: None, - })?; - ranked.push((rank, batch_index, row_index, score)); - } - } - - if ranked.len() != expected_len { - return Err(crate::Error::DataInvalid { - message: format!( - "DE vector materialization produced {} rows but search returned {expected_len}", - ranked.len() - ), - source: None, - }); - } - if ranked.is_empty() { - return Ok(Vec::new()); - } - - ranked.sort_by_key(|r| r.0); - let indices: Vec<(usize, usize)> = ranked.iter().map(|r| (r.1, r.2)).collect(); - let refs: Vec<&RecordBatch> = batches.iter().collect(); - let reordered = - interleave_record_batch(&refs, &indices).map_err(|e| crate::Error::DataInvalid { - message: format!("failed to reorder DE vector search rows: {e}"), - source: None, - })?; - - // Drop _ROW_ID. - let row_id_idx = reordered - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|_| crate::Error::DataInvalid { - message: format!("reordered batch missing {ROW_ID_FIELD_NAME} column"), - source: None, - })?; - let keep: Vec = (0..reordered.num_columns()) - .filter(|i| *i != row_id_idx) - .collect(); - let stripped = reordered - .project(&keep) - .map_err(|e| crate::Error::DataInvalid { - message: format!("failed to drop {ROW_ID_FIELD_NAME} column: {e}"), - source: None, - })?; - - // Append the score column in rank order. - let scores: Vec = ranked.iter().map(|r| r.3).collect(); - let score_array: Arc = Arc::new(Float32Array::from(scores)); - let mut fields: Vec> = - stripped.schema().fields().iter().cloned().collect(); - fields.push(Arc::new(arrow_schema::Field::new( - SEARCH_SCORE_COLUMN, - arrow_schema::DataType::Float32, - false, - ))); - let out_schema = Arc::new(arrow_schema::Schema::new(fields)); - let mut columns = stripped.columns().to_vec(); - columns.push(score_array); - let out = RecordBatch::try_new(out_schema, columns).map_err(|e| crate::Error::DataInvalid { - message: format!("failed to append DE vector score column: {e}"), - source: None, - })?; - Ok(vec![out]) -} - -fn indexed_search_limit(limit: usize, refine_factor: usize) -> crate::Result { - if refine_factor == 0 { - return Ok(limit); - } - let search_limit = - limit - .checked_mul(refine_factor) - .ok_or_else(|| crate::Error::ConfigInvalid { - message: format!( - "Vector search limit overflow: limit={limit}, refine factor={refine_factor}" - ), - })?; - if search_limit > i32::MAX as usize { - return Err(crate::Error::ConfigInvalid { - message: format!( - "Vector search limit overflow: limit={limit}, refine factor={refine_factor}" - ), - }); - } - Ok(search_limit) -} - -async fn maybe_rerank_indexed_batch_results( - evaluation: VectorSearchEvaluation<'_>, - index_entries: &[IndexManifestEntry], - field_id: i32, - field_name: &str, - vector_searches: &[VectorSearch], - results: Vec, - index_search_limit: usize, -) -> crate::Result> { - let timing_enabled = vector_search_timing_enabled(); - let total_start = timing_enabled.then(Instant::now); - let mut candidate_searches = Vec::with_capacity(vector_searches.len()); - let mut candidate_results = Vec::with_capacity(vector_searches.len()); - let mut union_candidates = RoaringTreemap::new(); - let mut candidate_references = 0usize; - - for (result, vector_search) in results.into_iter().zip(vector_searches) { - let candidates = result.top_k(index_search_limit); - candidate_references = candidate_references.saturating_add(candidates.row_ids.len()); - let mut include_row_ids = RoaringTreemap::new(); - for &row_id in &candidates.row_ids { - include_row_ids.insert(row_id); - union_candidates.insert(row_id); - } - - let mut candidate_search = vector_search.clone(); - candidate_search.set_shared_include_row_ids(Arc::new(include_row_ids)); - candidate_searches.push(candidate_search); - candidate_results.push(candidates); - } - - if union_candidates.iter().next().is_none() { - return Ok(candidate_results); - } - - let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { - message: "Vector index rerank requires table context".to_string(), - source: None, - })?; - let unique_candidates = union_candidates.len(); - let raw_ranges = sorted_row_ids_to_row_ranges(union_candidates.iter())?; - let metric_start = timing_enabled.then(Instant::now); - let metric = resolve_raw_vector_metric( - evaluation.file_io, - evaluation.table_path.trim_end_matches('/'), - evaluation.table_options, - index_entries, - field_id, - field_name, - ) - .await?; - let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); - - let (results, raw_timing) = - read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric).await?; - if let (Some(total_start), Some(raw_timing)) = (total_start, raw_timing) { - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_refine nq={} candidate_references={} unique_candidates={} row_ranges={} metric_resolve_ms={:.3} raw_plan_ms={:.3} split_count={} file_count={} raw_stream_wait_ms={:.3} raw_score_cpu_ms={:.3} arrow_batches={} arrow_rows={} total_refine_ms={:.3}", - vector_searches.len(), - candidate_references, - unique_candidates, - raw_ranges.len(), - metric_resolve.as_secs_f64() * 1000.0, - raw_timing.plan.as_secs_f64() * 1000.0, - raw_timing.split_count, - raw_timing.file_count, - raw_timing.stream_wait.as_secs_f64() * 1000.0, - raw_timing.score_cpu.as_secs_f64() * 1000.0, - raw_timing.batch_count, - raw_timing.row_count, - total_start.elapsed().as_secs_f64() * 1000.0, - ); - } - Ok(results) -} - -fn sorted_row_ids_to_row_ranges( - row_ids: impl IntoIterator, -) -> crate::Result> { - let mut row_ids = row_ids.into_iter(); - let Some(first) = row_ids.next() else { - return Ok(Vec::new()); - }; - let mut start = row_id_to_i64_for_range(first)?; - let mut end = start; - let mut ranges = Vec::new(); - for row_id in row_ids { - let row_id = row_id_to_i64_for_range(row_id)?; - if end.checked_add(1) == Some(row_id) { - end = row_id; - } else { - ranges.push(RowRange::new(start, end)); - start = row_id; - end = row_id; - } - } - ranges.push(RowRange::new(start, end)); - Ok(ranges) -} - -fn row_id_to_i64_for_range(row_id: u64) -> crate::Result { - i64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { - message: format!( - "Vector search row id {row_id} exceeds i64::MAX and cannot be converted to RowRange" - ), - source: None, - }) -} - -fn shared_batch_include_row_ids(vector_searches: &[VectorSearch]) -> Option<&Arc> { - let first = vector_searches.first()?.shared_include_row_ids.as_ref()?; - vector_searches - .iter() - .skip(1) - .all(|search| { - search - .shared_include_row_ids - .as_ref() - .is_some_and(|include_row_ids| Arc::ptr_eq(first, include_row_ids)) - }) - .then_some(first) -} - -fn prune_raw_ranges_by_include_row_ids( - raw_ranges: &[RowRange], - vector_searches: &[VectorSearch], -) -> crate::Result> { - if vector_searches - .iter() - .any(|search| search.effective_include_row_ids().is_none()) - { - return Ok(raw_ranges.to_vec()); - } - - let include_ranges = - if let Some(include_row_ids) = shared_batch_include_row_ids(vector_searches) { - sorted_row_ids_to_row_ranges(include_row_ids.iter())? - } else { - let mut union = RoaringTreemap::new(); - for include_row_ids in vector_searches - .iter() - .filter_map(VectorSearch::effective_include_row_ids) - { - for row_id in include_row_ids.iter() { - union.insert(row_id); - } - } - sorted_row_ids_to_row_ranges(union.iter())? - }; - Ok(intersect_sorted_ranges(raw_ranges, &include_ranges)) -} - -fn localize_include_row_ids( - include_row_ids: &RoaringTreemap, - row_range_start: i64, - row_range_end: i64, -) -> crate::Result { - let start = u64::try_from(row_range_start).map_err(|_| crate::Error::DataInvalid { - message: format!("Negative vector index row range start: {row_range_start}"), - source: None, - })?; - let end = u64::try_from(row_range_end).map_err(|_| crate::Error::DataInvalid { - message: format!("Negative vector index row range end: {row_range_end}"), - source: None, - })?; - let mut localized = RoaringTreemap::new(); - for row_id in include_row_ids.iter() { - if row_id >= start && row_id <= end { - localized.insert(row_id - start); - } - } - Ok(localized) -} - -fn localize_shared_include_row_ids( - include_row_ids: &RoaringTreemap, - ranges: &[(i64, i64)], -) -> crate::Result>> { - let mut validated_ranges = Vec::with_capacity(ranges.len()); - for (index, &(start, end)) in ranges.iter().enumerate() { - if start < 0 || end < start { - return Err(crate::Error::DataInvalid { - message: format!("Invalid vector index row range [{start}, {end}]"), - source: None, - }); - } - validated_ranges.push((start as u64, end as u64, index)); - } - validated_ranges.sort_unstable_by_key(|(start, _, _)| *start); - - let mut localized = (0..ranges.len()) - .map(|_| RoaringTreemap::new()) - .collect::>(); - let mut active = Vec::::new(); - let mut next_range = 0usize; - for row_id in include_row_ids.iter() { - while next_range < validated_ranges.len() && validated_ranges[next_range].0 <= row_id { - active.push(next_range); - next_range += 1; - } - active.retain(|range_index| validated_ranges[*range_index].1 >= row_id); - for range_index in &active { - let (start, _, original_index) = validated_ranges[*range_index]; - localized[original_index].insert(row_id - start); - } - if next_range == validated_ranges.len() && active.is_empty() { - break; - } - } - - Ok(localized - .into_iter() - .map(|filter| (!filter.is_empty()).then_some(filter)) - .collect()) -} - -async fn detail_data_ranges_for_table(table: &Table) -> crate::Result> { - let plan = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan() - .await?; - let mut ranges = Vec::new(); - for split in plan.splits() { - for file in split.data_files() { - if let Some((from, to)) = file.row_id_range() { - ranges.push(RowRange::new(from, to)); - } - } - } - Ok(merge_row_ranges(ranges)) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum RawVectorMetric { - L2, - Cosine, - InnerProduct, -} - -impl RawVectorMetric { - fn parse(value: &str) -> crate::Result { - Self::parse_normalized(&normalize_metric(value)).ok_or_else(|| crate::Error::DataInvalid { - message: format!("Unknown vector search metric: {value}"), - source: None, - }) - } - - fn parse_normalized(value: &str) -> Option { - match value { - "l2" => Some(Self::L2), - "cosine" => Some(Self::Cosine), - "inner_product" => Some(Self::InnerProduct), - _ => None, - } - } - - fn from_lumina(metric: LuminaVectorMetric) -> Self { - match metric { - LuminaVectorMetric::L2 => Self::L2, - LuminaVectorMetric::Cosine => Self::Cosine, - LuminaVectorMetric::InnerProduct => Self::InnerProduct, - } - } - - fn from_vindex(metric: MetricType) -> Self { - match metric { - MetricType::L2 => Self::L2, - MetricType::Cosine => Self::Cosine, - MetricType::InnerProduct => Self::InnerProduct, - } - } -} - -fn normalize_metric(metric: &str) -> String { - metric.to_ascii_lowercase().replace('-', "_") -} - -fn indexed_type_prefixes(field_name: &str, index_type: &str) -> Vec { - let mut prefixes = Vec::new(); - add_refine_prefixes(&mut prefixes, &format!("fields.{field_name}."), index_type); - add_refine_prefixes(&mut prefixes, "", index_type); - prefixes -} - -fn add_refine_prefixes(prefixes: &mut Vec, base: &str, index_type: &str) { - if !index_type.is_empty() { - prefixes.push(format!("{base}{index_type}.")); - let normalized = normalize_metric(index_type); - if normalized != index_type { - prefixes.push(format!("{base}{normalized}.")); - } - if normalized.starts_with("ivf") { - prefixes.push(format!("{base}ivf.")); - } - } - prefixes.push(base.to_string()); -} - -fn configured_refine_factor( - search_options: &HashMap, - table_options: &HashMap, - field_name: &str, - index_type: &str, -) -> crate::Result { - if let Some(value) = - configured_refine_factor_from_options(search_options, field_name, index_type) - { - return parse_refine_factor(&value); - } - if let Some(value) = - configured_refine_factor_from_options(table_options, field_name, index_type) - { - return parse_refine_factor(&value); - } - Ok(0) -} - -fn configured_refine_factor_from_options( - options: &HashMap, - field_name: &str, - index_type: &str, -) -> Option { - for prefix in indexed_type_prefixes(field_name, index_type) { - for suffix in [ - "refine_factor", - "refine-factor", - "rerank_factor", - "rerank-factor", - ] { - if let Some(value) = options.get(&(prefix.clone() + suffix)) { - return Some(value.trim().to_string()); - } - } - } - None -} - -fn parse_refine_factor(value: &str) -> crate::Result { - let factor = value - .parse::() - .map_err(|_| crate::Error::ConfigInvalid { - message: format!("Invalid vector refine factor: {value}. Must be an integer."), - })?; - if factor == 0 { - return Err(crate::Error::ConfigInvalid { - message: format!("Vector refine factor must be positive, got: {value}"), - }); - } - Ok(factor) -} - -async fn resolve_raw_vector_metric( - file_io: &FileIO, - table_path: &str, - table_options: &HashMap, - index_entries: &[IndexManifestEntry], - field_id: i32, - field_name: &str, -) -> crate::Result { - for entry in index_entries { - if entry.kind != FileKind::Add { - continue; - } - let Some(global_meta) = entry.index_file.global_index_meta.as_ref() else { - continue; - }; - if global_meta.index_field_id != field_id { - continue; - } - let Some(backend) = VectorIndexBackend::from_index_type(&entry.index_file.index_type) - else { - continue; - }; - match backend { - VectorIndexBackend::Lumina => { - if let Some(index_meta) = global_meta.index_meta.as_ref() { - if !index_meta.is_empty() { - let metric = LuminaIndexMeta::deserialize(index_meta)?.metric()?; - return Ok(RawVectorMetric::from_lumina(metric)); - } - } - } - VectorIndexBackend::Vindex => { - if let Some(index_meta) = global_meta.index_meta.as_ref() { - if let Ok(options) = - serde_json::from_slice::>(index_meta) - { - if let Some(metric) = options.get("metric") { - if let Some(metric) = - RawVectorMetric::parse_normalized(&normalize_metric(metric)) - { - return Ok(metric); - } - } - } - } - let path = IndexFileLocation::Global { table_path }.resolve( - &entry.index_file.file_name, - entry.index_file.external_path.as_deref(), - ); - let input = file_io.new_input(&path)?; - let read_error = |e| crate::Error::DataInvalid { - message: format!( - "Failed to read vindex index file '{}' for raw search metric: {}", - entry.index_file.file_name, e - ), - source: Some(Box::new(e)), - }; - let header_size = if entry.index_file.file_size > 0 { - (entry.index_file.file_size as u64).min(DISKANN_HEADER_SIZE as u64) - } else { - input - .metadata() - .await - .map_err(&read_error)? - .size - .min(DISKANN_HEADER_SIZE as u64) - }; - let file_reader = input.reader().await.map_err(&read_error)?; - let bytes = file_reader.read(0..header_size).await.map_err(read_error)?; - let reader = VIndexReader::open(Cursor::new(bytes)).map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to open paimon-vindex-core reader for raw search metric: {}", - e - ), - source: Some(Box::new(e)), - } - })?; - return Ok(RawVectorMetric::from_vindex(reader.metadata().metric)); - } - } - } - - configured_raw_vector_metric(table_options, field_name) -} - -fn configured_raw_vector_metric( - options: &HashMap, - field_name: &str, -) -> crate::Result { - let direct_keys = [ - format!("fields.{field_name}.distance.metric"), - format!("fields.{field_name}.metric"), - "test.vector.metric".to_string(), - "lumina.distance.metric".to_string(), - "distance.metric".to_string(), - "metric".to_string(), - ]; - for key in direct_keys { - if let Some(value) = options.get(&key) { - return RawVectorMetric::parse(value); - } - } - - let mut inferred = None; - for (key, value) in options { - if !(key.ends_with(".distance.metric") || key.ends_with(".metric")) { - continue; - } - let normalized = normalize_metric(value); - let Some(metric) = RawVectorMetric::parse_normalized(&normalized) else { - continue; - }; - if let Some(existing) = inferred { - if existing != metric { - return Ok(RawVectorMetric::L2); - } - } else { - inferred = Some(metric); - } - } - Ok(inferred.unwrap_or(RawVectorMetric::L2)) -} - -#[derive(Default)] -struct RawVectorReadTiming { - plan: Duration, - stream_wait: Duration, - score_cpu: Duration, - total: Duration, - split_count: usize, - file_count: usize, - batch_count: usize, - row_count: usize, -} - -async fn read_raw_batch_vector_search( - table: &Table, - vector_searches: &[VectorSearch], - raw_ranges: &[RowRange], - metric: RawVectorMetric, -) -> crate::Result<(Vec, Option)> { - let timing_enabled = vector_search_timing_enabled(); - let total_start = timing_enabled.then(Instant::now); - if vector_searches.is_empty() { - return Ok((Vec::new(), None)); - } - if raw_ranges.is_empty() { - return Ok((vec![SearchResult::empty(); vector_searches.len()], None)); - } - let raw_ranges = prune_raw_ranges_by_include_row_ids(raw_ranges, vector_searches)?; - if raw_ranges.is_empty() { - return Ok((vec![SearchResult::empty(); vector_searches.len()], None)); - } - - let field_name = &vector_searches[0].field_name; - if vector_searches - .iter() - .any(|vector_search| vector_search.field_name != *field_name) - { - return Err(crate::Error::DataInvalid { - message: "Batch vector raw search requires all query vectors to use the same field" - .to_string(), - source: None, - }); - } - - let plan_start = timing_enabled.then(Instant::now); - let mut read_builder = table.new_read_builder(); - read_builder - .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])? - .with_row_ranges(raw_ranges); - let plan = read_builder.new_scan().plan().await?; - let plan_elapsed = plan_start.map_or(Duration::ZERO, |start| start.elapsed()); - let split_count = plan.splits().len(); - let file_count = plan - .splits() - .iter() - .map(|split| split.data_files().len()) - .sum(); - if plan.splits().is_empty() { - return Ok(( - vec![SearchResult::empty(); vector_searches.len()], - total_start.map(|start| RawVectorReadTiming { - plan: plan_elapsed, - total: start.elapsed(), - ..RawVectorReadTiming::default() - }), - )); - } - let read = read_builder.new_read()?; - let mut stream = read.to_arrow(plan.splits())?; - - let scoring_plan = RawScoringPlan::new(vector_searches, metric); - let mut top_k = vector_searches - .iter() - .map(|vector_search| RawScoreTopK::new(vector_search.limit)) - .collect::>(); - let mut timing = timing_enabled.then(|| RawVectorReadTiming { - plan: plan_elapsed, - split_count, - file_count, - ..RawVectorReadTiming::default() - }); - loop { - let stream_wait_start = timing_enabled.then(Instant::now); - let batch = stream.try_next().await?; - if let (Some(timing), Some(stream_wait_start)) = (&mut timing, stream_wait_start) { - timing.stream_wait = timing - .stream_wait - .saturating_add(stream_wait_start.elapsed()); - } - let Some(batch) = batch else { - break; - }; - if let Some(timing) = &mut timing { - timing.batch_count += 1; - timing.row_count = timing.row_count.saturating_add(batch.num_rows()); - } - let score_start = timing_enabled.then(Instant::now); - collect_raw_batch_vector_batch(&batch, vector_searches, metric, &scoring_plan, &mut top_k)?; - if let (Some(timing), Some(score_start)) = (&mut timing, score_start) { - timing.score_cpu = timing.score_cpu.saturating_add(score_start.elapsed()); - } - } - - if let (Some(timing), Some(total_start)) = (&mut timing, total_start) { - timing.total = total_start.elapsed(); - } - Ok(( - top_k - .into_iter() - .map(RawScoreTopK::into_search_result) - .collect(), - timing, - )) -} - -struct RawScoringPlan { - all_query_indices: Vec, - shared_filter_groups: Vec, - candidate_query_indices: HashMap>, - query_l2_squared_norms: Vec, - dense_query_dimension: Option, - dense_query_matrix: Option>, -} - -struct SharedRawFilterGroup { - include_row_ids: Arc, - query_indices: Vec, -} - -impl RawScoringPlan { - fn new(vector_searches: &[VectorSearch], metric: RawVectorMetric) -> Self { - let mut all_query_indices = Vec::new(); - let mut shared_filter_groups = Vec::new(); - let mut candidate_query_indices: HashMap> = HashMap::new(); - let query_l2_squared_norms = vector_searches - .iter() - .map(|vector_search| match metric { - RawVectorMetric::L2 | RawVectorMetric::Cosine => vector_search - .vector - .iter() - .map(|value| value * value) - .sum::(), - RawVectorMetric::InnerProduct => 0.0, - }) - .collect(); - - if let Some(include_row_ids) = shared_batch_include_row_ids(vector_searches) { - shared_filter_groups.push(SharedRawFilterGroup { - include_row_ids: Arc::clone(include_row_ids), - query_indices: (0..vector_searches.len()).collect(), - }); - } else { - for (query_index, vector_search) in vector_searches.iter().enumerate() { - if let Some(include_row_ids) = vector_search.effective_include_row_ids() { - for row_id in include_row_ids.iter() { - candidate_query_indices - .entry(row_id) - .or_default() - .push(query_index); - } - } else { - all_query_indices.push(query_index); - } - } - } - - let dense_query_dimension = all_query_indices - .first() - .map(|&query_index| vector_searches[query_index].vector.len()); - let dense_query_matrix = dense_query_dimension.and_then(|dimension| { - all_query_indices - .iter() - .all(|&query_index| vector_searches[query_index].vector.len() == dimension) - .then(|| { - let mut matrix = - Vec::with_capacity(all_query_indices.len().saturating_mul(dimension)); - for &query_index in &all_query_indices { - matrix.extend_from_slice(&vector_searches[query_index].vector); - } - matrix - }) - }); - - Self { - all_query_indices, - shared_filter_groups, - candidate_query_indices, - query_l2_squared_norms, - dense_query_dimension, - dense_query_matrix, - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq)] -struct RawScoredRow { - row_id: u64, - score: f32, -} - -impl RawScoredRow { - fn strongest_first(a: &Self, b: &Self) -> Ordering { - b.score - .total_cmp(&a.score) - .then_with(|| a.row_id.cmp(&b.row_id)) - } -} - -struct RawScoreTopK { - limit: usize, - candidates: Vec, -} - -impl RawScoreTopK { - fn new(limit: usize) -> Self { - Self { - limit, - candidates: Vec::with_capacity(limit.min(1024).saturating_add(1)), - } - } - - fn offer(&mut self, row_id: u64, score: f32) { - if self.limit == 0 { - return; - } - self.candidates.push(RawScoredRow { row_id, score }); - if self.candidates.len() >= self.partition_size() { - self.reduce_to_limit(); - } - } - - fn offer_many(&mut self, candidates: I) - where - I: IntoIterator, - { - if self.limit == 0 { - return; - } - self.candidates.extend(candidates); - if self.candidates.len() >= self.partition_size() { - self.reduce_to_limit(); - } - } - - fn partition_size(&self) -> usize { - self.limit - .saturating_mul(2) - .max(RAW_TOP_K_MIN_PARTITION_SIZE) - } - - fn reduce_to_limit(&mut self) { - if self.candidates.len() <= self.limit { - return; - } - // Partition only after a substantial candidate block has accumulated. - // Each partition is linear in its input, so all reductions are O(n) - // amortized; only the final K survivors are fully sorted. - self.candidates - .select_nth_unstable_by(self.limit, RawScoredRow::strongest_first); - self.candidates.truncate(self.limit); - } - - fn into_search_result(mut self) -> SearchResult { - self.reduce_to_limit(); - self.candidates - .sort_unstable_by(RawScoredRow::strongest_first); - let rows = self.candidates; - let mut row_ids = Vec::with_capacity(rows.len()); - let mut scores = Vec::with_capacity(rows.len()); - for row in rows { - row_ids.push(row.row_id); - scores.push(row.score); - } - SearchResult::new(row_ids, scores) - } -} - -fn collect_raw_batch_vector_batch( - batch: &RecordBatch, - vector_searches: &[VectorSearch], - metric: RawVectorMetric, - scoring_plan: &RawScoringPlan, - top_k_out: &mut [RawScoreTopK], -) -> crate::Result<()> { - if vector_searches.is_empty() { - return Ok(()); - } - if top_k_out.len() != vector_searches.len() { - return Err(crate::Error::DataInvalid { - message: "Raw batch vector search output buffers must match query vector count" - .to_string(), - source: None, - }); - } - - let field_name = &vector_searches[0].field_name; - if vector_searches - .iter() - .any(|vector_search| vector_search.field_name != *field_name) - { - return Err(crate::Error::DataInvalid { - message: "Batch vector raw search requires all query vectors to use the same field" - .to_string(), - source: None, - }); - } - - let vector_index = - batch - .schema() - .index_of(field_name) - .map_err(|e| crate::Error::DataInvalid { - message: format!( - "Vector column '{}' not found in raw search batch: {}", - field_name, e - ), - source: None, - })?; - let row_id_index = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|e| crate::Error::DataInvalid { - message: format!("_ROW_ID column not found in raw search batch: {e}"), - source: None, - })?; - - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: "Vector raw search requires non-null Int64 _ROW_ID".to_string(), - source: None, - })?; - - let column = batch.column(vector_index); - enum VectorLayout<'a> { - List(&'a ListArray), - Fixed(&'a FixedSizeListArray), - } - let layout = if let Some(a) = column.as_any().downcast_ref::() { - VectorLayout::List(a) - } else if let Some(a) = column.as_any().downcast_ref::() { - VectorLayout::Fixed(a) - } else { - return Err(crate::Error::DataInvalid { - message: "Vector raw search requires Arrow List or FixedSizeList" - .to_string(), - source: None, - }); - }; - let values = match layout { - VectorLayout::List(a) => a.values(), - VectorLayout::Fixed(a) => a.values(), - } - .as_any() - .downcast_ref::() - .ok_or_else(|| crate::Error::DataInvalid { - message: "Vector raw search requires Float32 vector elements".to_string(), - source: None, - })?; - - let use_dense_matrix = scoring_plan.all_query_indices.len() >= RAW_SCORE_MATRIX_MIN_QUERY_COUNT; - let dense_dimension = use_dense_matrix - .then_some(scoring_plan.dense_query_dimension) - .flatten(); - let mut dense_row_ids = Vec::with_capacity(batch.num_rows()); - let mut dense_vectors = Vec::with_capacity( - batch - .num_rows() - .saturating_mul(dense_dimension.unwrap_or_default()), - ); - for row in 0..batch.num_rows() { - if row_ids.is_null(row) { - return Err(crate::Error::DataInvalid { - message: "Vector raw search found null _ROW_ID".to_string(), - source: None, - }); - } - let row_id = row_id_to_u64(row_ids.value(row))?; - let is_null = match layout { - VectorLayout::List(a) => a.is_null(row), - VectorLayout::Fixed(a) => a.is_null(row), - }; - if is_null { - continue; - } - - let (start, end) = match layout { - VectorLayout::List(a) => { - let offsets = a.value_offsets(); - (offsets[row] as usize, offsets[row + 1] as usize) - } - VectorLayout::Fixed(a) => { - let len = a.value_length() as usize; - let start = a.value_offset(row) as usize; - (start, start + len) - } - }; - ensure_raw_vector_values_not_null(values, start, end)?; - - let raw_row = RawVectorRow { - row_id, - values, - start, - end, - }; - if let Some(dimension) = dense_dimension { - ensure_raw_vector_dimension(end - start, dimension)?; - if scoring_plan.dense_query_matrix.is_none() { - let &query_index = scoring_plan - .all_query_indices - .iter() - .find(|&&query_index| vector_searches[query_index].vector.len() != dimension) - .expect("a missing dense matrix requires inconsistent query dimensions"); - ensure_raw_vector_dimension(dimension, vector_searches[query_index].vector.len())?; - } - dense_row_ids.push(row_id); - dense_vectors.extend_from_slice(&values.values()[start..end]); - } else { - for &query_index in &scoring_plan.all_query_indices { - offer_raw_vector_score( - raw_row, - query_index, - metric, - vector_searches, - scoring_plan, - top_k_out, - )?; - } - } - if let Some(query_indices) = scoring_plan.candidate_query_indices.get(&row_id) { - for &query_index in query_indices { - offer_raw_vector_score( - raw_row, - query_index, - metric, - vector_searches, - scoring_plan, - top_k_out, - )?; - } - } - for group in &scoring_plan.shared_filter_groups { - if group.include_row_ids.contains(row_id) { - for &query_index in &group.query_indices { - offer_raw_vector_score( - raw_row, - query_index, - metric, - vector_searches, - scoring_plan, - top_k_out, - )?; - } - } - } - } - - if !dense_row_ids.is_empty() { - let query_matrix = scoring_plan - .dense_query_matrix - .as_deref() - .expect("dense query dimensions were validated above"); - let dimension = dense_dimension.expect("dense rows require dense queries"); - let queries_per_chunk = (RAW_SCORE_MATRIX_TARGET_ELEMENTS / dense_row_ids.len()) - .max(1) - .min(scoring_plan.all_query_indices.len()); - for (query_chunk_index, query_indices) in scoring_plan - .all_query_indices - .chunks(queries_per_chunk) - .enumerate() - { - let query_start = query_chunk_index * queries_per_chunk * dimension; - let query_end = query_start + query_indices.len() * dimension; - let scores = compute_raw_vector_score_matrix( - &dense_vectors, - dense_row_ids.len(), - &query_matrix[query_start..query_end], - query_indices.len(), - dimension, - &scoring_plan.query_l2_squared_norms, - query_indices, - metric, - )?; - for (matrix_query_index, &query_index) in query_indices.iter().enumerate() { - let query_scores = &scores[matrix_query_index * dense_row_ids.len() - ..(matrix_query_index + 1) * dense_row_ids.len()]; - top_k_out[query_index].offer_many( - dense_row_ids - .iter() - .zip(query_scores) - .map(|(&row_id, &score)| RawScoredRow { row_id, score }), - ); - } - } - } - - Ok(()) -} - -fn ensure_raw_vector_dimension(stored_len: usize, query_len: usize) -> crate::Result<()> { - if stored_len != query_len { - return Err(crate::Error::DataInvalid { - message: format!( - "Query vector dimension mismatch: raw row has {}, but query has {}", - stored_len, query_len - ), - source: None, - }); - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn compute_raw_vector_score_matrix( - stored_vectors: &[f32], - row_count: usize, - query_vectors: &[f32], - query_count: usize, - dimension: usize, - query_l2_squared_norms: &[f32], - query_indices: &[usize], - metric: RawVectorMetric, -) -> crate::Result> { - let score_count = - row_count - .checked_mul(query_count) - .ok_or_else(|| crate::Error::DataInvalid { - message: "Vector raw search score matrix is too large".to_string(), - source: None, - })?; - debug_assert_eq!(stored_vectors.len(), row_count * dimension); - debug_assert_eq!(query_vectors.len(), query_count * dimension); - debug_assert_eq!(query_indices.len(), query_count); - - let mut scores = vec![0.0; score_count]; - // Query × stored-vector^T produces a query-major score matrix. Each query's - // scores are contiguous, which feeds partial Top-K without strided reads. - sgemm_a_bt( - query_count, - row_count, - dimension, - 1.0, - query_vectors, - stored_vectors, - 0.0, - &mut scores, - ); - if metric == RawVectorMetric::InnerProduct { - return Ok(scores); - } - - let stored_l2_squared_norms = stored_vectors - .chunks_exact(dimension) - .map(|vector| vector.iter().map(|value| value * value).sum::()) - .collect::>(); - for (matrix_query_index, &query_index) in query_indices.iter().enumerate() { - for (row_index, &stored_l2_squared_norm) in stored_l2_squared_norms.iter().enumerate() { - let score = &mut scores[matrix_query_index * row_count + row_index]; - let query_l2_squared_norm = query_l2_squared_norms[query_index]; - *score = match metric { - RawVectorMetric::L2 => { - let squared_distance = - stored_l2_squared_norm + query_l2_squared_norm - 2.0 * *score; - // The norm/dot reconstruction loses the low-order difference when two - // large vectors are close. Estimate a conservative accumulation-error - // bound and preserve the former scalar semantics inside that region. - let roundoff_bound = (stored_l2_squared_norm.abs() - + query_l2_squared_norm.abs() - + 2.0 * score.abs()) - * f32::EPSILON - * (dimension as f32 + 2.0) - * 4.0; - if !squared_distance.is_finite() || squared_distance <= roundoff_bound { - let stored = - &stored_vectors[row_index * dimension..(row_index + 1) * dimension]; - let query = &query_vectors - [matrix_query_index * dimension..(matrix_query_index + 1) * dimension]; - compute_raw_vector_l2_score(query, stored) - } else { - 1.0 / (1.0 + squared_distance) - } - } - RawVectorMetric::Cosine => { - let denominator = stored_l2_squared_norm.sqrt() * query_l2_squared_norm.sqrt(); - if denominator == 0.0 { - 0.0 - } else { - *score / denominator - } - } - RawVectorMetric::InnerProduct => unreachable!(), - }; - } - } - Ok(scores) -} - -fn ensure_raw_vector_values_not_null( - values: &Float32Array, - start: usize, - end: usize, -) -> crate::Result<()> { - if values.null_count() == 0 { - return Ok(()); - } - for value_index in start..end { - if values.is_null(value_index) { - return Err(crate::Error::DataInvalid { - message: "Vector raw search found null vector element".to_string(), - source: None, - }); - } - } - Ok(()) -} - -#[derive(Clone, Copy)] -struct RawVectorRow<'a> { - row_id: u64, - values: &'a Float32Array, - start: usize, - end: usize, -} - -fn offer_raw_vector_score( - row: RawVectorRow<'_>, - query_index: usize, - metric: RawVectorMetric, - vector_searches: &[VectorSearch], - scoring_plan: &RawScoringPlan, - top_k_out: &mut [RawScoreTopK], -) -> crate::Result<()> { - let vector_search = &vector_searches[query_index]; - let stored_len = row.end - row.start; - ensure_raw_vector_dimension(stored_len, vector_search.vector.len())?; - let score = compute_raw_vector_score_from_values( - &vector_search.vector, - scoring_plan.query_l2_squared_norms[query_index], - row.values, - row.start, - row.end, - metric, - ); - top_k_out[query_index].offer(row.row_id, score); - Ok(()) -} - -fn compute_raw_vector_score_from_values( - query: &[f32], - query_l2_squared_norm: f32, - values: &Float32Array, - start: usize, - end: usize, - metric: RawVectorMetric, -) -> f32 { - debug_assert_eq!(query.len(), end - start); - match metric { - RawVectorMetric::L2 => compute_raw_vector_l2_score(query, &values.values()[start..end]), - RawVectorMetric::Cosine => { - let mut dot = 0.0; - let mut norm_b = 0.0; - for (q, value_index) in query.iter().zip(start..end) { - let stored = values.value(value_index); - dot += q * stored; - norm_b += stored * stored; - } - let denominator = query_l2_squared_norm.sqrt() * norm_b.sqrt(); - if denominator == 0.0 { - 0.0 - } else { - dot / denominator - } - } - RawVectorMetric::InnerProduct => query - .iter() - .zip(start..end) - .map(|(q, value_index)| q * values.value(value_index)) - .sum(), - } -} - -fn compute_raw_vector_l2_score(query: &[f32], stored: &[f32]) -> f32 { - let squared_distance = query - .iter() - .zip(stored) - .map(|(query_value, stored_value)| { - let difference = query_value - stored_value; - difference * difference - }) - .sum::(); - 1.0 / (1.0 + squared_distance) -} - -fn row_id_to_u64(row_id: i64) -> crate::Result { - u64::try_from(row_id).map_err(|_| crate::Error::DataInvalid { - message: format!("Negative _ROW_ID {row_id} cannot be used for global index search"), - source: None, - }) -} - -#[cfg(test)] -fn compute_raw_vector_score(query: &[f32], stored: &[f32], metric: RawVectorMetric) -> f32 { - match metric { - RawVectorMetric::L2 => compute_raw_vector_l2_score(query, stored), - RawVectorMetric::Cosine => { - let mut dot = 0.0; - let mut norm_a = 0.0; - let mut norm_b = 0.0; - for (q, s) in query.iter().zip(stored.iter()) { - dot += q * s; - norm_a += q * q; - norm_b += s * s; - } - let denominator = norm_a.sqrt() * norm_b.sqrt(); - if denominator == 0.0 { - 0.0 - } else { - dot / denominator - } - } - RawVectorMetric::InnerProduct => query.iter().zip(stored.iter()).map(|(q, s)| q * s).sum(), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::catalog::Identifier; - use crate::io::FileIOBuilder; - use crate::lumina::{LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, LUMINA_IDENTIFIER}; - use crate::spec::stats::BinaryTableStats; - use crate::spec::{ - ArrayType, BinaryRow, DataFileMeta, DataType, Datum, FloatType, GlobalIndexMeta, - IndexFileMeta, IndexManifestEntry, IntType, PredicateBuilder, Schema, TableSchema, - }; - use crate::table::source::DataSplitBuilder; - use crate::table::{TableCommit, TableWrite}; - use crate::vindex::IVF_FLAT_IDENTIFIER; - use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, ListBuilder}; - use arrow_array::ArrayRef; - use arrow_array::Int32Array; - use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; - use std::sync::{Arc, Mutex, Once}; - - const VECTOR_SEARCH_LOG_TARGET: &str = "paimon::vector_search"; - static VECTOR_SEARCH_TEST_LOGGER: VectorSearchTestLogger = VectorSearchTestLogger; - static VECTOR_SEARCH_TEST_LOGS: Mutex> = Mutex::new(Vec::new()); - - struct VectorSearchTestLogger; - - impl log::Log for VectorSearchTestLogger { - fn enabled(&self, metadata: &log::Metadata<'_>) -> bool { - metadata.target() == VECTOR_SEARCH_LOG_TARGET && metadata.level() <= log::Level::Debug - } - - fn log(&self, record: &log::Record<'_>) { - if self.enabled(record.metadata()) { - VECTOR_SEARCH_TEST_LOGS - .lock() - .unwrap() - .push(record.args().to_string()); - } - } - - fn flush(&self) {} - } - - fn reset_vector_search_test_logs() { - static INIT: Once = Once::new(); - INIT.call_once(|| log::set_logger(&VECTOR_SEARCH_TEST_LOGGER).unwrap()); - log::set_max_level(log::LevelFilter::Debug); - VECTOR_SEARCH_TEST_LOGS.lock().unwrap().clear(); - } - - fn l2_score(distance: f32) -> f32 { - VectorSearchMetric::L2.distance_to_score(distance) - } - - #[test] - fn vindex_concurrency_limits_are_independent() { - let default_options = HashMap::new(); - let default_core = CoreOptions::new(&default_options); - assert_eq!( - vindex_concurrency_limits(&default_core, 1, 32).unwrap(), - (1, 64) - ); - assert_eq!( - vindex_concurrency_limits(&default_core, 8, 4).unwrap(), - (4, 64) - ); - - let options = HashMap::from([( - "global-index.vindex.read-thread-num".to_string(), - "48".to_string(), - )]); - let core = CoreOptions::new(&options); - assert_eq!(vindex_concurrency_limits(&core, 1, 32).unwrap(), (1, 48)); - assert_eq!(vindex_concurrency_limits(&core, 8, 4).unwrap(), (4, 48)); - } - - #[test] - fn vindex_array_dimension_accepts_diskann_search_options() { - let field = DataField::new( - 1, - "embedding".to_string(), - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ); - let query_options = HashMap::from([ - ("diskann.dimension".to_string(), "8".to_string()), - ("diskann.l_search".to_string(), "64".to_string()), - ( - "vindex.reader.memory-budget-bytes".to_string(), - "1048576".to_string(), - ), - ]); - - assert_eq!( - pk_vector_query_dimension(&HashMap::new(), &query_options, "diskann", &field).unwrap(), - Some(8) - ); - } - - fn make_field(id: i32, name: &str) -> DataField { - DataField::new(id, name.to_string(), DataType::Int(IntType::default())) - } - - fn vector_test_table() -> Table { - vector_test_table_at("memory:/vector_test") - } - - fn vector_test_table_at(location: &str) -> Table { - vector_test_table_with_file_io(FileIOBuilder::new("memory").build().unwrap(), location) - } - - fn vector_test_table_with_file_io(file_io: FileIO, location: &str) -> Table { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .build() - .unwrap(); - Table::new( - file_io, - Identifier::new("default", "vector_test"), - location.to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - fn eval_context<'a>( - file_io: &'a FileIO, - options: &'a HashMap, - fields: &'a [DataField], - next_row_id: Option, - ) -> VectorSearchEvaluation<'a> { - VectorSearchEvaluation { - table: None, - file_io, - table_path: "memory:///test_table", - table_options: options, - schema_fields: fields, - next_row_id, - } - } - - #[test] - fn test_find_field_id_by_name() { - let fields = vec![make_field(1, "id"), make_field(2, "embedding")]; - assert_eq!(find_field_id_by_name(&fields, "embedding"), Some(2)); - assert_eq!(find_field_id_by_name(&fields, "nonexistent"), None); - } - - #[test] - fn shared_include_filter_is_localized_once_per_index_shard() { - let include_row_ids = RoaringTreemap::from_iter([101, 205, 999]); - let localized = localize_shared_include_row_ids( - &include_row_ids, - &[(100, 109), (200, 209), (300, 309)], - ) - .unwrap(); - - assert_eq!( - localized[0].as_ref().unwrap().iter().collect::>(), - vec![1] - ); - assert_eq!( - localized[1].as_ref().unwrap().iter().collect::>(), - vec![5] - ); - assert!(localized[2].is_none(), "an empty shard must be skipped"); - } - - #[test] - fn shared_batch_include_filter_requires_the_same_arc() { - let shared = Arc::new(RoaringTreemap::from_iter([1, 2, 3])); - let mut shared_searches = vec![ - VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(), - VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(), - ]; - for search in &mut shared_searches { - search.set_shared_include_row_ids(Arc::clone(&shared)); - } - let detected = shared_batch_include_row_ids(&shared_searches).unwrap(); - assert!(Arc::ptr_eq(detected, &shared)); - - let mut equal_but_distinct = shared_searches.clone(); - equal_but_distinct[1] - .set_shared_include_row_ids(Arc::new(RoaringTreemap::from_iter([1, 2, 3]))); - assert!(shared_batch_include_row_ids(&equal_but_distinct).is_none()); - - let mut owned = shared_searches; - owned[1] = owned[1] - .clone() - .with_include_row_ids(RoaringTreemap::from_iter([1, 2, 3])); - assert!(shared_batch_include_row_ids(&owned).is_none()); - } - - #[test] - fn shared_raw_filter_does_not_expand_row_query_associations() { - let shared = Arc::new(RoaringTreemap::from_iter(0..1_000)); - let mut searches = (0..128) - .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap()) - .collect::>(); - for search in &mut searches { - search.set_shared_include_row_ids(Arc::clone(&shared)); - } - - let plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); - let expanded_associations = plan - .candidate_query_indices - .values() - .map(Vec::len) - .sum::(); - - assert_eq!( - expanded_associations, 0, - "one shared bitmap must stay O(B + Q), not expand to O(B * Q)" - ); - assert_eq!(plan.shared_filter_groups.len(), 1); - assert!(Arc::ptr_eq( - &plan.shared_filter_groups[0].include_row_ids, - &shared - )); - assert_eq!(plan.shared_filter_groups[0].query_indices.len(), 128); - } - - #[test] - fn shared_raw_filter_prunes_unindexed_ranges_before_reading() { - let shared = Arc::new(RoaringTreemap::from_iter([7, 1_000, 1_001, 900_000])); - let mut searches = (0..128) - .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap()) - .collect::>(); - for search in &mut searches { - search.set_shared_include_row_ids(Arc::clone(&shared)); - } - - let raw_ranges = vec![RowRange::new(0, 999_999)]; - assert_eq!( - prune_raw_ranges_by_include_row_ids(&raw_ranges, &searches).unwrap(), - vec![ - RowRange::new(7, 7), - RowRange::new(1_000, 1_001), - RowRange::new(900_000, 900_000), - ] - ); - } - - #[test] - fn test_raw_vector_score_matches_java_metric_semantics() { - let l2 = compute_raw_vector_score(&[1.0, 2.0], &[1.0, 4.0], RawVectorMetric::L2); - assert!((l2 - 0.2).abs() < 1e-6); - assert_eq!( - compute_raw_vector_score(&[1.0, 2.0], &[3.0, 4.0], RawVectorMetric::InnerProduct), - 11.0 - ); - let cosine = compute_raw_vector_score(&[1.0, 0.0], &[1.0, 1.0], RawVectorMetric::Cosine); - assert!((cosine - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-6); - assert_eq!( - compute_raw_vector_score(&[0.0, 0.0], &[1.0, 1.0], RawVectorMetric::Cosine), - 0.0 - ); - } - - #[test] - fn test_raw_vector_score_matrix_matches_scalar_metrics() { - let stored = vec![1.0, 2.0, 3.0, 4.0, 0.0, 0.0]; - let queries = vec![1.0, 1.0, -1.0, 2.0]; - let query_indices = vec![0, 1]; - let query_l2_squared_norms = vec![2.0, 5.0]; - - for metric in [ - RawVectorMetric::L2, - RawVectorMetric::Cosine, - RawVectorMetric::InnerProduct, - ] { - let matrix_scores = compute_raw_vector_score_matrix( - &stored, - 3, - &queries, - 2, - 2, - &query_l2_squared_norms, - &query_indices, - metric, - ) - .unwrap(); - for (row_index, stored_vector) in stored.as_chunks::<2>().0.iter().enumerate() { - for (query_index, query) in queries.as_chunks::<2>().0.iter().enumerate() { - let expected = compute_raw_vector_score(query, stored_vector, metric); - let actual = matrix_scores[query_index * 3 + row_index]; - assert!( - (actual - expected).abs() < 1e-5, - "metric={metric:?}, row={row_index}, query={query_index}: {actual} != {expected}" - ); - } - } - } - - let non_finite_score = compute_raw_vector_score_matrix( - &[f32::INFINITY, 0.0], - 1, - &[1.0, 0.0], - 1, - 2, - &[1.0], - &[0], - RawVectorMetric::L2, - ) - .unwrap()[0]; - assert_eq!(non_finite_score, 0.0); - } - - #[test] - fn test_raw_vector_score_matrix_l2_preserves_large_finite_distances() { - let dimension = 128; - let query = vec![1.0e10_f32; dimension]; - let mut nearby = query.clone(); - nearby[0] += 1024.0; - let mut stored = query.clone(); - stored.extend_from_slice(&nearby); - let queries = query.repeat(4); - let query_l2_squared_norm = query.iter().map(|value| value * value).sum::(); - let query_l2_squared_norms = vec![query_l2_squared_norm; 4]; - let query_indices = vec![0, 1, 2, 3]; - - let matrix_scores = compute_raw_vector_score_matrix( - &stored, - 2, - &queries, - 4, - dimension, - &query_l2_squared_norms, - &query_indices, - RawVectorMetric::L2, - ) - .unwrap(); - let exact_score = compute_raw_vector_score(&query, &query, RawVectorMetric::L2); - let nearby_score = compute_raw_vector_score(&query, &nearby, RawVectorMetric::L2); - - for query_index in 0..4 { - assert_eq!(matrix_scores[query_index * 2], exact_score); - assert_eq!(matrix_scores[query_index * 2 + 1], nearby_score); - assert!(matrix_scores[query_index * 2] > matrix_scores[query_index * 2 + 1]); - } - } - - #[test] - fn test_raw_vector_cosine_avoids_squared_norm_product_overflow() { - let query = vec![1.0e15_f32, 0.0]; - let query_l2_squared_norm = query.iter().map(|value| value * value).sum::(); - assert!(query_l2_squared_norm.is_finite()); - let values = Float32Array::from(query.clone()); - let scalar_score = compute_raw_vector_score_from_values( - &query, - query_l2_squared_norm, - &values, - 0, - 2, - RawVectorMetric::Cosine, - ); - assert!((scalar_score - 1.0).abs() < 1e-6); - - let queries = query.repeat(4); - let matrix_scores = compute_raw_vector_score_matrix( - &query, - 1, - &queries, - 4, - 2, - &[query_l2_squared_norm; 4], - &[0, 1, 2, 3], - RawVectorMetric::Cosine, - ) - .unwrap(); - assert!(matrix_scores - .iter() - .all(|score| (*score - 1.0).abs() < 1e-6)); - } - - #[test] - fn test_raw_score_top_k_matches_full_sort_with_linear_partial_selection() { - let limit = 7; - let mut top_k = RawScoreTopK::new(limit); - let mut batched_top_k = RawScoreTopK::new(limit); - let mut expected = Vec::new(); - for row_id in 0..10_000 { - let score = ((row_id * 37) % 101) as f32 / 10.0; - let candidate = RawScoredRow { row_id, score }; - expected.push(candidate); - top_k.offer(row_id, score); - batched_top_k.offer_many(std::iter::once(candidate)); - assert!(top_k.candidates.len() < top_k.partition_size()); - assert!(batched_top_k.candidates.len() < batched_top_k.partition_size()); - } - expected.sort_unstable_by(RawScoredRow::strongest_first); - expected.truncate(limit); - - let result = top_k.into_search_result(); - let batched_result = batched_top_k.into_search_result(); - assert_eq!( - result.row_ids, - expected.iter().map(|row| row.row_id).collect::>() - ); - assert_eq!( - result.scores, - expected.iter().map(|row| row.score).collect::>() - ); - assert_eq!(batched_result.row_ids, result.row_ids); - assert_eq!(batched_result.scores, result.scores); - } - - #[test] - fn test_configured_raw_vector_metric_precedence_and_conflict_default() { - let mut options = HashMap::new(); - options.insert( - "fields.embedding.distance.metric".to_string(), - "inner-product".to_string(), - ); - options.insert("metric".to_string(), "cosine".to_string()); - assert_eq!( - configured_raw_vector_metric(&options, "embedding").unwrap(), - RawVectorMetric::InnerProduct - ); - - options.clear(); - options.insert("foo.metric".to_string(), "cosine".to_string()); - options.insert("bar.distance.metric".to_string(), "l2".to_string()); - assert_eq!( - configured_raw_vector_metric(&options, "embedding").unwrap(), - RawVectorMetric::L2 - ); - } - - #[tokio::test] - async fn test_resolve_raw_vector_metric_uses_vindex_manifest_metadata() { - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let mut entry = make_lumina_entry("missing.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); - let index_meta = serde_json::to_vec(&HashMap::from([( - "metric".to_string(), - "cosine".to_string(), - )])) - .unwrap(); - entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .index_meta = Some(index_meta); - - let metric = resolve_raw_vector_metric( - &file_io, - "memory:///test_table", - &HashMap::new(), - &[entry], - 2, - "embedding", - ) - .await - .unwrap(); - - assert_eq!(metric, RawVectorMetric::Cosine); - } - - #[tokio::test] - async fn test_resolve_raw_vector_metric_falls_back_to_vindex_header() { - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let index = build_vindex_segment_bytes("inner_product"); - file_io - .new_output("memory:///test_table/index/test.idx") - .unwrap() - .write(bytes::Bytes::from(index.clone())) - .await - .unwrap(); - for (file_size, index_meta) in [ - (index.len() as i64, br#"{"metric":"euclidean"}"#.to_vec()), - (0, b"{}".to_vec()), - (-1, b"{}".to_vec()), - ] { - let mut entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); - entry.index_file.file_size = file_size; - entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .index_meta = Some(index_meta); - - let metric = resolve_raw_vector_metric( - &file_io, - "memory:///test_table", - &HashMap::new(), - &[entry], - 2, - "embedding", - ) - .await - .unwrap(); - - assert_eq!(metric, RawVectorMetric::InnerProduct); - } - } - - #[test] - fn test_configured_refine_factor_precedence_and_aliases() { - let table_options = HashMap::from([( - "fields.embedding.ivf.refine-factor".to_string(), - "3".to_string(), - )]); - let search_options = HashMap::from([( - "fields.embedding.ivf_flat.rerank_factor".to_string(), - "2".to_string(), - )]); - assert_eq!( - configured_refine_factor( - &search_options, - &table_options, - "embedding", - IVF_FLAT_IDENTIFIER, - ) - .unwrap(), - 2 - ); - - assert_eq!( - configured_refine_factor( - &HashMap::new(), - &table_options, - "embedding", - IVF_FLAT_IDENTIFIER, - ) - .unwrap(), - 3 - ); - - let global_options = HashMap::from([("rerank-factor".to_string(), "4".to_string())]); - assert_eq!( - configured_refine_factor( - &HashMap::new(), - &global_options, - "embedding", - LUMINA_IDENTIFIER, - ) - .unwrap(), - 4 - ); - } - - #[test] - fn test_configured_refine_factor_rejects_invalid_values() { - let zero_options = HashMap::from([("refine_factor".to_string(), "0".to_string())]); - let err = configured_refine_factor( - &zero_options, - &HashMap::new(), - "embedding", - LUMINA_IDENTIFIER, - ) - .unwrap_err(); - assert!(err.to_string().contains("must be positive")); - - let invalid_options = HashMap::from([("refine_factor".to_string(), "abc".to_string())]); - let err = configured_refine_factor( - &invalid_options, - &HashMap::new(), - "embedding", - LUMINA_IDENTIFIER, - ) - .unwrap_err(); - assert!(err.to_string().contains("Must be an integer")); - - assert!(indexed_search_limit(i32::MAX as usize, 2).is_err()); - } - - #[test] - fn test_collect_raw_batch_vector_batch_preserves_query_order() { - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut builder = - FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); - for vector in [[1.0, 0.0], [0.0, 1.0], [0.8, 0.2]] { - builder.values().append_value(vector[0]); - builder.values().append_value(vector[1]); - builder.append(true); - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList( - Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), - 2, - ), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(builder.finish()) as ArrayRef, - Arc::new(Int64Array::from(vec![Some(10), Some(11), Some(12)])) as ArrayRef, - ], - ) - .unwrap(); - let searches = vec![ - VectorSearch::new(vec![1.0, 0.0], 1, "embedding".to_string()).unwrap(), - VectorSearch::new(vec![0.0, 1.0], 1, "embedding".to_string()).unwrap(), - VectorSearch::new(vec![0.8, 0.2], 1, "embedding".to_string()).unwrap(), - VectorSearch::new(vec![0.5, 0.5], 1, "embedding".to_string()).unwrap(), - ]; - let scoring_plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); - let mut top_k = searches - .iter() - .map(|search| RawScoreTopK::new(search.limit)) - .collect::>(); - - collect_raw_batch_vector_batch( - &batch, - &searches, - RawVectorMetric::L2, - &scoring_plan, - &mut top_k, - ) - .unwrap(); - let results = top_k - .into_iter() - .map(RawScoreTopK::into_search_result) - .collect::>(); - - assert_eq!(results[0].row_ids, vec![10]); - assert_eq!(results[1].row_ids, vec![11]); - assert_eq!(results[2].row_ids, vec![12]); - assert_eq!(results[3].row_ids, vec![12]); - } - - #[test] - fn test_collect_raw_batch_vector_batch_respects_fixed_size_list_offset() { - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut builder = - FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); - for vector in [[1.0, 0.0], [0.0, 1.0], [0.8, 0.2]] { - builder.values().append_value(vector[0]); - builder.values().append_value(vector[1]); - builder.append(true); - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList( - Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), - 2, - ), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(builder.finish()) as ArrayRef, - Arc::new(Int64Array::from(vec![10, 11, 12])) as ArrayRef, - ], - ) - .unwrap() - .slice(1, 2); - let searches = vec![VectorSearch::new(vec![0.0, 1.0], 1, "embedding".to_string()).unwrap()]; - let scoring_plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); - let mut top_k = vec![RawScoreTopK::new(1)]; - - collect_raw_batch_vector_batch( - &batch, - &searches, - RawVectorMetric::L2, - &scoring_plan, - &mut top_k, - ) - .unwrap(); - - assert_eq!(top_k.pop().unwrap().into_search_result().row_ids, vec![11]); - } - - #[test] - fn test_collect_raw_batch_vector_batch_scores_only_include_row_ids() { - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut builder = - FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); - for vector in [[1.0, 0.0], [0.0, 1.0], [0.8, 0.2]] { - builder.values().append_value(vector[0]); - builder.values().append_value(vector[1]); - builder.append(true); - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList( - Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), - 2, - ), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(builder.finish()) as ArrayRef, - Arc::new(Int64Array::from(vec![Some(10), Some(11), Some(12)])) as ArrayRef, - ], - ) - .unwrap(); - let mut include_row_ids = RoaringTreemap::new(); - include_row_ids.insert(12); - let searches = vec![ - VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()) - .unwrap() - .with_include_row_ids(include_row_ids), - ]; - let scoring_plan = RawScoringPlan::new(&searches, RawVectorMetric::L2); - let mut top_k = searches - .iter() - .map(|search| RawScoreTopK::new(search.limit)) - .collect::>(); - - collect_raw_batch_vector_batch( - &batch, - &searches, - RawVectorMetric::L2, - &scoring_plan, - &mut top_k, - ) - .unwrap(); - let results = top_k - .into_iter() - .map(RawScoreTopK::into_search_result) - .collect::>(); - - assert_eq!(results[0].row_ids, vec![12]); - assert_eq!(results[0].scores.len(), 1); - } - - #[tokio::test] - async fn test_batch_vector_search_requires_vectors() { - let table = vector_test_table(); - let err = table - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(Vec::new()) - .with_limit(1) - .execute() - .await - .unwrap_err(); - - assert!( - err.to_string() - .contains("Query vectors must be set via with_query_vectors()"), - "unexpected error: {err}" - ); - } - - #[tokio::test] - async fn test_batch_vector_search_rejects_zero_limit() { - let table = vector_test_table(); - let err = table - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0]]) - .with_limit(0) - .execute() - .await - .unwrap_err(); - - assert!( - err.to_string().contains("Limit must be between 1"), - "unexpected error: {err}" - ); - } - - #[tokio::test] - async fn test_batch_evaluate_no_matching_field_returns_empty_per_query() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let fields = vec![make_field(1, "id")]; - let searches = vec![ - VectorSearch::new(vec![1.0], 10, "embedding".to_string()).unwrap(), - VectorSearch::new(vec![0.0], 10, "embedding".to_string()).unwrap(), - ]; - let options = HashMap::new(); - - let entry = make_lumina_entry( - "test.idx", - LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, - FileKind::Add, - 99, - ); - - let results = evaluate_batch_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &searches, - ) - .await - .unwrap(); - - assert_eq!(results.len(), searches.len()); - assert!(results.iter().all(SearchResult::is_empty)); - } - - #[tokio::test] - async fn test_evaluate_no_matching_entries() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let fields = vec![make_field(1, "id"), make_field(2, "embedding")]; - let vs = VectorSearch::new(vec![1.0, 2.0], 10, "embedding".to_string()).unwrap(); - let options = HashMap::new(); - - let entry = IndexManifestEntry { - kind: FileKind::Add, - partition: vec![], - bucket: 0, - index_file: IndexFileMeta { - index_type: "btree".to_string(), - file_name: "test.idx".to_string(), - file_size: 100, - row_count: 10, - deletion_vectors_ranges: None, - external_path: None, - global_index_meta: None, - }, - version: 1, - }; - - let result = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap(); - assert!(result.is_empty()); - } - - #[tokio::test] - 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(); - let options = HashMap::new(); - - let entry = make_lumina_entry("test.idx", "btree", FileKind::Add, 2); - - let result = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap(); - assert!(result.is_empty()); - } - - #[tokio::test] - async fn test_evaluate_full_mode_without_vector_entries_uses_raw_path() { - 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 options = HashMap::from([("vector-index.search-mode".to_string(), "full".to_string())]); - - let err = evaluate_vector_search( - eval_context(&file_io, &options, &fields, Some(10)), - &[], - &vs, - ) - .await - .unwrap_err(); - assert!( - err.to_string() - .contains("Vector raw search requires table context"), - "unexpected error: {err}" - ); - } - - #[tokio::test] - async fn test_evaluate_no_matching_field() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let fields = vec![make_field(1, "id")]; - let vs = VectorSearch::new(vec![1.0], 10, "embedding".to_string()).unwrap(); - let options = HashMap::new(); - - let entry = make_lumina_entry( - "test.idx", - LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, - FileKind::Add, - 99, - ); - - let result = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap(); - assert!(result.is_empty()); - } - - #[tokio::test] - async fn test_evaluate_skips_delete_entries() { - 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 options = HashMap::new(); - - let entry = make_lumina_entry( - "test.idx", - LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, - FileKind::Delete, - 2, - ); - - let result = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap(); - assert!(result.is_empty()); - } - - #[tokio::test] - async fn test_evaluate_accepts_canonical_lumina_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 options = HashMap::new(); - - let entry = make_lumina_entry("missing.idx", LUMINA_IDENTIFIER, FileKind::Add, 2); - - let err = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap_err(); - assert!( - err.to_string() - .contains("Failed to read Lumina index file 'missing.idx'"), - "unexpected error: {err}" - ); - } - - #[tokio::test] - async fn test_evaluate_accepts_legacy_lumina_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 options = HashMap::new(); - - let entry = make_lumina_entry( - "missing.idx", - LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, - FileKind::Add, - 2, - ); - - let err = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap_err(); - assert!( - err.to_string() - .contains("Failed to read Lumina index file 'missing.idx'"), - "unexpected error: {err}" - ); - } - - #[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 options = HashMap::new(); - - let entry = make_lumina_entry("missing.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); - - let err = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .unwrap_err(); - assert!( - err.to_string() - .contains("Failed to read vindex index file 'missing.idx'"), - "unexpected error: {err}" - ); - assert!( - std::error::Error::source(&err).is_some(), - "wrapped vindex read errors should retain their source: {err:?}" - ); - } - - #[test] - fn test_single_vindex_outside_tokio_returns_error() { - futures::executor::block_on(async { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - file_io - .new_output("memory:///test_table/index/test.idx") - .unwrap() - .write(bytes::Bytes::from_static(b"index")) - .await - .unwrap(); - let fields = vec![make_field(2, "embedding")]; - let vs = VectorSearch::new(vec![1.0], 10, "embedding".to_string()).unwrap(); - let options = HashMap::new(); - let entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); - - let err = evaluate_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &vs, - ) - .await - .expect_err("vindex range reads outside Tokio should fail without panicking"); - - assert!( - matches!(err, crate::Error::UnexpectedError { ref message, .. } - if message.contains("requires a Tokio runtime")), - "unexpected error: {err:?}" - ); - }); - } - - #[test] - fn test_batch_vindex_outside_tokio_uses_buffered_fallback() { - futures::executor::block_on(async { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let index = build_vindex_segment_bytes("l2"); - file_io - .new_output("memory:///test_table/index/test.idx") - .unwrap() - .write(bytes::Bytes::from(index.clone())) - .await - .unwrap(); - let fields = vec![make_field(2, "embedding")]; - let searches = vec![ - VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(), - VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(), - ]; - let options = HashMap::new(); - let mut entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); - entry.index_file.file_size = index.len() as i64; - entry.index_file.row_count = 3; - entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .row_range_end = 2; - - let results = evaluate_batch_vector_search( - eval_context(&file_io, &options, &fields, None), - &[entry], - &searches, - ) - .await - .expect("batch vindex search should fall back to buffered I/O outside Tokio"); - - assert_eq!(results.len(), searches.len()); - assert!(results.iter().all(|result| !result.is_empty())); - }); - } - - #[tokio::test] - async fn test_execute_fails_closed_when_query_auth_enabled() { - let table = crate::table::query_auth_table(); - let err = table - .new_vector_search_builder() - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), - "vector search must fail closed for a query-auth table" - ); - } - - #[tokio::test] - async fn test_batch_execute_fails_closed_when_query_auth_enabled() { - // The batch scored entry returns data-derived row ids/scores outside - // `TableScan`/`TableRead`, so it must fail closed under - // `query-auth.enabled` exactly like the single-query builder. Its config - // is otherwise valid, so without the guard the empty-snapshot fast path - // would return empty results and silently bypass authorization. - let table = crate::table::query_auth_table(); - let err = table - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0, 2.0]]) - .with_limit(5) - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), - "batch vector search must fail closed for a query-auth table, got: {err:?}" - ); - } - - #[tokio::test] - async fn prepared_filter_cannot_bypass_builder_target_query_auth() { - let source = vector_test_table(); - let prepared = source - .prepare_vector_search_filter(id_gt_filter(&source, 0)) - .await - .unwrap(); - let target = source.copy_with_options(HashMap::from([( - "query-auth.enabled".to_string(), - "true".to_string(), - )])); - - let err = target - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0, 0.0]]) - .with_limit(1) - .with_prepared_filter(prepared) - .execute() - .await - .expect_err("a stale prepared filter must not bypass current target authorization"); - - assert!( - matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), - "builder target authorization must remain authoritative, got: {err:?}" - ); - } - - fn pk_data_file(name: &str, row_count: i64, first_row_id: Option) -> DataFileMeta { - DataFileMeta { - file_name: name.to_string(), - file_size: 1, - row_count, - min_key: Vec::new(), - max_key: Vec::new(), - key_stats: BinaryTableStats::empty(), - value_stats: BinaryTableStats::empty(), - min_sequence_number: 0, - max_sequence_number: 0, - schema_id: 1, - level: 0, - extra_files: Vec::new(), - creation_time: None, - delete_row_count: None, - embedded_index: None, - file_source: None, - value_stats_cols: None, - external_path: None, - first_row_id, - write_cols: None, - column_max_sequence_numbers: None, - } - } - - fn pk_search_split(bucket: i32, files: Vec) -> PkVectorSearchSplit { - PkVectorSearchSplit { - data_split: DataSplitBuilder::new() - .with_snapshot(1) - .with_partition(BinaryRow::new(0)) - .with_bucket(bucket) - .with_bucket_path(format!("memory:/t/bucket-{bucket}")) - .with_total_buckets(1) - .with_data_files(files) - .build() - .unwrap(), - ann_segments: Vec::new(), - active_files: Vec::new(), - } - } - - fn pk_candidate( - split_index: usize, - bucket: i32, - file: &str, - pos: i64, - distance: f32, - ) -> PkVectorCandidate { - PkVectorCandidate { - split_index, - partition: BinaryRow::new(0), - bucket, - data_file_name: file.to_string(), - row_position: pos, - distance, - } - } - - // Candidate with a fixed empty (arity-0) partition and bucket 0, keyed only by - // (split_index, file, position) — the dimensions the rerank core groups on. - fn cand_at(split_index: usize, file: &str, pos: i64, dist: f32) -> PkVectorCandidate { - pk_candidate(split_index, 0, file, pos, dist) - } - - /// The single data-file name every rerank fixture writes. - const RERANK_FILE: &str = "part-0.parquet"; - - /// Serialize a Paimon deletion-vector blob covering `deleted_rows` and write it - /// at `path`, returning the matching `DeletionFile`. Byte layout mirrors the - /// position-read tests: `[length][magic][roaring bitmap][0]`. - async fn write_deletion_blob( - file_io: &FileIO, - path: &str, - deleted_rows: &[u32], - ) -> crate::table::source::DeletionFile { - use roaring::RoaringBitmap; - - const MAGIC_NUMBER: i32 = 1581511376; - let mut bitmap = RoaringBitmap::new(); - for row in deleted_rows { - bitmap.insert(*row); - } - let mut bitmap_bytes = Vec::new(); - bitmap.serialize_into(&mut bitmap_bytes).unwrap(); - let bitmap_length = 4 + bitmap_bytes.len() as i32; - let mut blob = Vec::new(); - blob.extend_from_slice(&bitmap_length.to_be_bytes()); - blob.extend_from_slice(&MAGIC_NUMBER.to_be_bytes()); - blob.extend_from_slice(&bitmap_bytes); - blob.extend_from_slice(&0i32.to_be_bytes()); - file_io - .new_output(path) - .unwrap() - .write(bytes::Bytes::from(blob)) - .await - .unwrap(); - crate::table::source::DeletionFile::new( - path.to_string(), - 0, - bitmap_length as i64, - Some(deleted_rows.len() as i64), - ) - } - - /// Write a single-file vector data file (`FixedSizeList` of width - /// `dim`) holding `rows` (a `None` entry is a NULL vector row) as Parquet, and - /// return a vector-only `DataFileReader`, the enclosing `PkVectorSearchSplit`, - /// and the vector `DataField`. When `deleted_rows` is non-empty a deletion - /// vector covering those physical positions is attached to the split, so the - /// position read drops them exactly as `PkVectorIndexedSplitRead::read` does. - /// - /// This is the position-only analogue of the old `ArrayReader`: rerank now - /// re-reads real stored rows through `PkVectorPositionRead`, so the fixtures - /// exercise that path rather than an in-memory preloaded column. - async fn vector_rerank_fixture( - table_path: &str, - dim: u32, - rows: &[Option>], - deleted_rows: &[u32], - ) -> (DataFileReader, PkVectorSearchSplit, DataField) { - use crate::arrow::build_target_arrow_schema; - use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; - use crate::spec::VectorType; - use crate::table::schema_manager::SchemaManager; - - let vector_type = - VectorType::try_new(true, dim, DataType::Float(FloatType::new())).unwrap(); - let vector_field = - DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)); - let read_fields = vec![vector_field.clone()]; - let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); - - let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), dim as i32).with_field( - Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), - ); - for row in rows { - match row { - Some(values) => { - for v in values { - builder.values().append_value(*v); - } - builder.append(true); - } - None => { - for _ in 0..dim { - builder.values().append_value(0.0); - } - builder.append(false); - } - } - } - let vec_array = builder.finish(); - let batch = - arrow_array::RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]) - .unwrap(); - - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let bucket_path = format!("{table_path}/bucket-0"); - let output = file_io - .new_output(&format!("{bucket_path}/{RERANK_FILE}")) - .unwrap(); - let mut writer: Box = Box::new( - ParquetFormatWriter::new( - &output, - arrow_schema.clone(), - "zstd", - 1, - None, - &HashMap::new(), - ) - .await - .unwrap(), - ); - writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap().file_size; - - let schema_id = 1; - let file_meta = pk_data_file(RERANK_FILE, rows.len() as i64, Some(0)); - let file_meta = DataFileMeta { - file_size: file_size as i64, - schema_id, - ..file_meta - }; - - let mut split_builder = DataSplitBuilder::new() - .with_snapshot(1) - .with_partition(BinaryRow::new(0)) - .with_bucket(0) - .with_bucket_path(bucket_path) - .with_total_buckets(1) - .with_data_files(vec![file_meta]); - if !deleted_rows.is_empty() { - let df = - write_deletion_blob(&file_io, &format!("{table_path}/index/dv-0"), deleted_rows) - .await; - split_builder = split_builder.with_data_deletion_files(vec![Some(df)]); - } - let data_split = split_builder.build().unwrap(); - let split = PkVectorSearchSplit { - data_split, - ann_segments: Vec::new(), - active_files: Vec::new(), - }; - - let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); - let reader = DataFileReader::new( - file_io, - schema_manager, - schema_id, - read_fields.clone(), - read_fields, - Vec::new(), - ); - (reader, split, vector_field) - } - - #[tokio::test] - async fn rerank_aligns_recomputed_distance_by_position_column() { - use crate::arrow::build_target_arrow_schema; - use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; - use crate::spec::VectorType; - use crate::table::schema_manager::SchemaManager; - - // A vector data file with 4 physical rows: positions 0,1,3 hold vectors - // and position 2 (a NON-candidate) holds a NULL vector. Candidates sit at - // non-contiguous positions {1, 3}. The ANN-reported distances are - // deliberately reversed relative to the true stored vectors; after rerank - // each candidate must carry compute_distance(query, vec_at_its_position), - // proving alignment is by the _PKEY_VECTOR_POSITION column value, not batch - // order. Position 2's NULL is never read (it is not a candidate), so it - // cannot trip the null-vector guard. - let vector_type = VectorType::try_new(true, 2, DataType::Float(FloatType::new())).unwrap(); - let vector_field = - DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)); - let read_fields = vec![vector_field.clone()]; - let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); - - // pos0=[7,0], pos1=[1,0], pos2=NULL, pos3=[4,0]. - let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new( - ArrowField::new("element", ArrowDataType::Float32, true), - )); - for row in [ - Some([7.0f32, 0.0]), - Some([1.0, 0.0]), - None, - Some([4.0, 0.0]), - ] { - match row { - Some([a, b]) => { - builder.values().append_value(a); - builder.values().append_value(b); - builder.append(true); - } - None => { - builder.values().append_value(0.0); - builder.values().append_value(0.0); - builder.append(false); - } - } - } - let vec_array = builder.finish(); - let batch = - arrow_array::RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]) - .unwrap(); - - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let table_path = "memory:/rerank_positional"; - let bucket_path = format!("{table_path}/bucket-0"); - let file_name = "part-0.parquet"; - let output = file_io - .new_output(&format!("{bucket_path}/{file_name}")) - .unwrap(); - let mut writer: Box = Box::new( - ParquetFormatWriter::new( - &output, - arrow_schema.clone(), - "zstd", - 1, - None, - &HashMap::new(), - ) - .await - .unwrap(), - ); - writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap().file_size; - - let schema_id = 1; - let file_meta = pk_data_file(file_name, 4, Some(0)); - let file_meta = DataFileMeta { - file_size: file_size as i64, - schema_id, - ..file_meta - }; - let data_split = DataSplitBuilder::new() - .with_snapshot(1) - .with_partition(BinaryRow::new(0)) - .with_bucket(0) - .with_bucket_path(bucket_path) - .with_total_buckets(1) - .with_data_files(vec![file_meta]) - .build() - .unwrap(); - let split = PkVectorSearchSplit { - data_split, - ann_segments: Vec::new(), - active_files: Vec::new(), - }; - - let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); - let reader = DataFileReader::new( - file_io, - schema_manager, - schema_id, - read_fields.clone(), - read_fields.clone(), - Vec::new(), - ); - - let query = vec![1.0f32, 0.0]; - // ANN-reported distances reversed vs. truth: pos1 reported worse (0.9) than - // pos3 (0.1), but the true L2 distances are pos1=0 and pos3=9. - let indexed = vec![cand_at(0, file_name, 1, 0.9), cand_at(0, file_name, 3, 0.1)]; - - let out = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 2, - &vector_field, - ) - .await - .unwrap(); - - // Best-first after exact recompute: pos1 (d=0) then pos3 (d=9), each - // carrying the distance computed from its OWN position's stored vector. - assert_eq!(out.len(), 2); - assert_eq!(out[0].row_position, 1); - assert_eq!(out[0].distance, 0.0); - assert_eq!(out[1].row_position, 3); - assert_eq!(out[1].distance, 9.0); - } - - #[tokio::test] - async fn rerank_recomputes_distance_and_reorders() { - // pos0=[9,0], pos1=[1,0]; query=[1,0]. The ANN-reported distances are - // reversed relative to the truth (pos0 reported best at 0.1, pos1 worst at - // 0.9), so an implementation that trusted the ANN order would emit pos0 - // first. Exact L2 recompute yields pos0=64, pos1=0, so the output must - // reorder to pos1-then-pos0 with the recomputed distances. - let (reader, split, vector_field) = vector_rerank_fixture( - "memory:/rerank_reorder", - 2, - &[Some(vec![9.0, 0.0]), Some(vec![1.0, 0.0])], - &[], - ) - .await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![ - cand_at(0, RERANK_FILE, 0, 0.1), - cand_at(0, RERANK_FILE, 1, 0.9), - ]; - - let out = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 2, - &vector_field, - ) - .await - .unwrap(); - - assert_eq!(out.len(), 2); - assert_eq!(out[0].row_position, 1); - assert_eq!(out[0].distance, 0.0); - assert_eq!(out[1].row_position, 0); - assert_eq!(out[1].distance, 64.0); - // Order genuinely changed vs. the ANN-reported best-first (which was pos0). - assert!(out[0].distance < out[1].distance); - } - - #[tokio::test] - async fn rerank_is_independent_of_fast_mode_reranks_indexed() { - // The rerank core takes only the indexed (fast-path) candidates and always - // recomputes their true distance; there is no fast/exact switch that can - // skip it. The single candidate carries a bogus ANN distance (0.42) but its - // stored vector equals the query, so the recomputed L2 distance is exactly - // 0.0 — proving the indexed candidate WAS reranked rather than passed - // through with its ANN distance. - let (reader, split, vector_field) = - vector_rerank_fixture("memory:/rerank_indexed", 2, &[Some(vec![1.0, 0.0])], &[]).await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.42)]; - - let out = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 1, - &vector_field, - ) - .await - .unwrap(); - - assert_eq!(out.len(), 1); - assert_eq!(out[0].row_position, 0); - assert_ne!(out[0].distance, 0.42); - assert_eq!(out[0].distance, 0.0); - } - - #[tokio::test] - async fn rerank_fails_loud_on_null_vector() { - // A NULL vector stored AT a candidate position must fail loud rather than - // silently scoring it: the candidate genuinely has no vector to rerank on. - let (reader, split, vector_field) = - vector_rerank_fixture("memory:/rerank_null", 2, &[None], &[]).await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; - - let err = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 1, - &vector_field, - ) - .await - .err() - .expect("null vector at a candidate position must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("null vector")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn rerank_fails_loud_on_leftover_candidate() { - // pos1 is deleted by the deletion vector, so the position read returns no - // row for it. The search path already DV-filters, so a deleted candidate - // reaching rerank is a real inconsistency: the leftover guard must fail - // loud rather than silently dropping the candidate. - let (reader, split, vector_field) = vector_rerank_fixture( - "memory:/rerank_leftover", - 2, - &[Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])], - &[1], - ) - .await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![ - cand_at(0, RERANK_FILE, 0, 0.1), - cand_at(0, RERANK_FILE, 1, 0.9), - ]; - - let err = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 2, - &vector_field, - ) - .await - .err() - .expect("a candidate returning no row must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("failed to read")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn rerank_fails_loud_on_dimension_mismatch() { - // Stored vectors are 3-dimensional but the query is 2-dimensional. The - // vector extraction validates each stored row against the query dimension - // and fails loud, so the recompute never runs against mismatched vectors. - let (reader, split, vector_field) = - vector_rerank_fixture("memory:/rerank_dim", 3, &[Some(vec![1.0, 0.0, 0.0])], &[]).await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; - - let err = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 1, - &vector_field, - ) - .await - .err() - .expect("dimension mismatch must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("dimension")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn rerank_fails_loud_on_duplicate_candidate_position() { - // Two candidates addressing the same (split_index, file, position) is a - // programming error upstream: the dedup guard fires before any read. - let (reader, split, vector_field) = - vector_rerank_fixture("memory:/rerank_dup", 2, &[Some(vec![1.0, 0.0])], &[]).await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![ - cand_at(0, RERANK_FILE, 0, 0.1), - cand_at(0, RERANK_FILE, 0, 0.9), - ]; - - let err = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 2, - &vector_field, - ) - .await - .err() - .expect("duplicate candidate position must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("duplicate")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn rerank_fails_loud_on_unexpected_position() { - // Every position the read surfaces must resolve to a candidate keyed by - // (split_index, file, position). Here the plan carries two splits for the - // SAME (partition, bucket, file), so `split_index_of` resolves the file to - // the LAST plan index (1). The single candidate is tagged with split_index - // 0, so its by_key entry is (0, file, 0) while the read looks up - // (1, file, 0). The lookup misses and the unexpected-position guard fires - // rather than silently dropping the surfaced row. - let (reader, split, vector_field) = - vector_rerank_fixture("memory:/rerank_unexpected", 2, &[Some(vec![1.0, 0.0])], &[]) - .await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; - - // Two plan entries for the same file: split_index_of ends up mapping the - // file to plan index 1, not the candidate's split_index 0. - let dup = PkVectorSearchSplit { - data_split: split.data_split.clone(), - ann_segments: Vec::new(), - active_files: Vec::new(), - }; - let plan = vec![dup, split]; - - let err = rerank_indexed_positional( - &reader, - indexed, - &plan, - &query, - VectorSearchMetric::L2, - 1, - &vector_field, - ) - .await - .err() - .expect("a read position absent from the candidate map must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("unexpected position")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn rerank_fails_loud_on_file_not_in_plan() { - // A candidate references a (partition, bucket, file) that is absent from - // plan_splits. build_indexed_splits groups it into an indexed split, but the - // split_index_of lookup — built only from plan_splits — has no entry, so the - // kernel fails loud rather than reading an unplanned file. - let (reader, _split, vector_field) = - vector_rerank_fixture("memory:/rerank_noplan", 2, &[Some(vec![1.0, 0.0])], &[]).await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; - - // Empty plan: the candidate's file resolves in no plan split. - let err = rerank_indexed_positional( - &reader, - indexed, - &[], - &query, - VectorSearchMetric::L2, - 1, - &vector_field, - ) - .await - .err() - .expect("a candidate file absent from the plan must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("not found in plan")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn rerank_reads_only_candidate_positions_not_whole_column() { - // A 6-row file where every NON-candidate position (0, 2, 4, 5) holds a NULL - // vector "poison" and only the two candidate positions (1, 3) hold real - // vectors. The rerank read is told to fetch only positions {1, 3}; every - // row it surfaces is looked up in the candidate map, and any position not in - // the map trips the "unexpected position" guard (a surfaced NULL row would - // additionally trip the null-vector guard). So if the read had surfaced any - // of the poison rows, rerank would fail. It succeeds and returns exactly the - // two candidates at positions {1, 3}, which proves the position selection - // reaching the read contained only the candidate positions (not the whole - // column). - let rows = &[ - None, // pos0 poison (non-candidate) - Some(vec![1.0, 0.0]), // pos1 candidate - None, // pos2 poison (non-candidate) - Some(vec![3.0, 0.0]), // pos3 candidate - None, // pos4 poison (non-candidate) - None, // pos5 poison (non-candidate) - ]; - let (reader, split, vector_field) = - vector_rerank_fixture("memory:/rerank_spy", 2, rows, &[]).await; - let query = vec![1.0f32, 0.0]; - let indexed = vec![ - cand_at(0, RERANK_FILE, 1, 0.9), - cand_at(0, RERANK_FILE, 3, 0.1), - ]; - - let out = rerank_indexed_positional( - &reader, - indexed, - &[split], - &query, - VectorSearchMetric::L2, - 2, - &vector_field, - ) - .await - .unwrap_or_else(|e| { - panic!("only candidate positions are read, so the poison NULLs never decode: {e:?}") - }); - - assert_eq!(out.len(), 2, "exactly the candidate count of rows was read"); - let mut positions: Vec = out.iter().map(|c| c.row_position).collect(); - positions.sort_unstable(); - assert_eq!( - positions, - vec![1, 3], - "only candidate positions reached the read" - ); - // Recomputed distances confirm each surviving row is its own candidate's vector. - assert_eq!(out[0].row_position, 1); - assert_eq!(out[0].distance, 0.0); - assert_eq!(out[1].row_position, 3); - assert_eq!(out[1].distance, 4.0); - } - - /// Build a real vindex IVF-flat segment trained with `metric`, returning the - /// serialized bytes. `nlist = 1` keeps training trivial and deterministic; the - /// only thing the metric check cares about is the persisted metadata metric. - fn build_vindex_segment_bytes(metric: &str) -> Vec { - use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; - use paimon_vindex_core::io::PosWriter; - - const DIM: usize = 2; - let vectors: Vec = vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; - let n = vectors.len() / DIM; - let ids: Vec = (0..n as i64).collect(); - let options = HashMap::from([ - ("index.type".to_string(), "ivf_flat".to_string()), - ("dimension".to_string(), DIM.to_string()), - ("nlist".to_string(), "1".to_string()), - ("metric".to_string(), metric.to_string()), - ]); - let config = VectorIndexConfig::from_options(&options).unwrap(); - let training = VectorIndexTrainer::train(config, &vectors, n).unwrap(); - let mut writer = VectorIndexWriter::new(training); - writer.add_vectors(&ids, &vectors, n).unwrap(); - let mut bytes = Vec::new(); - { - let mut output = PosWriter::new(&mut bytes); - writer.write(&mut output).unwrap(); - } - bytes - } - - fn pk_split_with_lumina_segment(path: &str, metric: &str) -> PkVectorSearchSplit { - let mut split = pk_search_split(0, vec![pk_data_file("file-a", 3, Some(0))]); - let source_meta = crate::spec::PrimaryKeyIndexSourceMeta::new( - 1, - vec![crate::spec::PrimaryKeyIndexSourceFile::new("file-a".to_string(), 3).unwrap()], - ) - .unwrap(); - let mut segment = BucketAnnSegment::for_test(source_meta); - segment.path = path.to_string(); - // Lumina stores its metric in the serialized index metadata blob, not in - // the segment file bytes. `deserialize` requires both keys present. - let meta = crate::lumina::LuminaIndexMeta::new(HashMap::from([ - ("index.dimension".to_string(), "2".to_string()), - ("distance.metric".to_string(), metric.to_string()), - ])); - segment.index_meta = meta.serialize().unwrap(); - split.ann_segments = vec![segment]; - split - } - - #[test] - fn verify_segment_metric_accepts_matching_lumina_metric() { - // Lumina segment metadata says cosine; configured cosine => Ok. No segment - // file bytes are needed on the Lumina path. - let split = pk_split_with_lumina_segment("seg-lumina", "cosine"); - let segment = &split.ann_segments[0]; - let lumina_metric = LuminaIndexMeta::deserialize(&segment.index_meta) - .unwrap() - .metric() - .unwrap(); - verify_segment_metric( - VectorSearchMetric::Cosine, - VectorSearchMetric::from_lumina(lumina_metric), - ) - .expect("matching lumina metric must pass"); - } - - #[test] - fn verify_segment_metric_rejects_mismatched_lumina_metric() { - // Lumina segment metadata says l2; configured inner_product => fail loud, - // naming both metrics. - let split = pk_split_with_lumina_segment("seg-lumina", "l2"); - let segment = &split.ann_segments[0]; - let lumina_metric = LuminaIndexMeta::deserialize(&segment.index_meta) - .unwrap() - .metric() - .unwrap(); - let err = verify_segment_metric( - VectorSearchMetric::InnerProduct, - VectorSearchMetric::from_lumina(lumina_metric), - ) - .expect_err("mismatched lumina metric must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("does not match configured metric") - && message.contains("l2") - && message.contains("inner_product")), - "unexpected error: {err:?}" - ); - } - - #[test] - fn from_index_type_classifies_lumina_and_vindex() { - assert_eq!( - VectorIndexBackend::from_index_type("lumina"), - Some(VectorIndexBackend::Lumina) - ); - assert_eq!( - VectorIndexBackend::from_index_type("lumina-vector-ann"), - Some(VectorIndexBackend::Lumina) - ); - assert_eq!( - VectorIndexBackend::from_index_type("ivf-flat"), - Some(VectorIndexBackend::Vindex) - ); - for index_type in ["ivf-sq", "ivf-rq", "diskann"] { - assert_eq!( - VectorIndexBackend::from_index_type(index_type), - Some(VectorIndexBackend::Vindex) - ); - } - } - - #[test] - fn verify_segment_metric_accepts_matching_vindex_metric() { - // Real IVF segment trained with L2; configured metric L2 => Ok. - let bytes = bytes::Bytes::from(build_vindex_segment_bytes("l2")); - let reader = VIndexReader::open(Cursor::new(bytes)).unwrap(); - verify_segment_metric( - VectorSearchMetric::L2, - VectorSearchMetric::from_vindex(reader.metadata().metric), - ) - .expect("matching metric must pass"); - } - - #[test] - fn verify_segment_metric_rejects_mismatched_vindex_metric() { - // Real IVF segment trained with L2; configured metric Cosine => fail loud. - let bytes = bytes::Bytes::from(build_vindex_segment_bytes("l2")); - let reader = VIndexReader::open(Cursor::new(bytes)).unwrap(); - let err = verify_segment_metric( - VectorSearchMetric::Cosine, - VectorSearchMetric::from_vindex(reader.metadata().metric), - ) - .expect_err("mismatched metric must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("does not match configured metric") - && message.contains("l2") - && message.contains("cosine")), - "unexpected error: {err:?}" - ); - } - - fn pk_vector_table(options: &[(&str, &str)]) -> Table { - let mut builder = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ); - if options - .iter() - .any(|(key, _)| *key == "pk-vector.index.columns") - { - builder = builder.primary_key(["id"]).option("bucket", "1"); - } - let schema = builder.build().unwrap(); - // Runtime validation must remain defensive for schemas committed by old - // or external writers, including malformed configurations which the - // current Schema builder rejects at commit time. - let runtime_options = options - .iter() - .map(|(key, value)| ((*key).to_string(), (*value).to_string())) - .collect(); - let table_schema = TableSchema::new(0, &schema).copy_with_options(runtime_options); - Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "pk_vector_test"), - "memory:/pk_vector_test".to_string(), - table_schema, - None, - ) - } - - /// A data-evolution (global-index) vector table with a committed IVF-flat - /// index over the `embedding` column: row-tracking + data-evolution + - /// global-index enabled so committed data files carry `first_row_id` and the - /// search returns global row-ids that `execute_read` can materialize. The - /// returned table has one committed batch of `(id, embedding)` rows and a real - /// vindex index built end-to-end. - async fn de_vector_table() -> Table { - let table_path = "memory:/de_vector_search_test"; - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .option("row-tracking.enabled", "true") - .option("data-evolution.enabled", "true") - .option("global-index.enabled", "true") - .option("global-index.row-count-per-shard", "10") - .option("ivf-flat.dimension", "2") - .option("ivf-flat.nlist", "2") - .build() - .unwrap(); - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let table = Table::new( - file_io.clone(), - Identifier::new("default", "de_vector_test"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ); - file_io - .mkdirs(&format!("{table_path}/snapshot/")) - .await - .unwrap(); - file_io - .mkdirs(&format!("{table_path}/manifest/")) - .await - .unwrap(); - - let ids = vec![1, 2, 3]; - let vectors = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]]; - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut vector_builder = - ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); - for vector in vectors { - for value in vector { - vector_builder.values().append_value(value); - } - vector_builder.append(true); - } - let arrow_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("embedding", ArrowDataType::List(element_field), true), - ])); - let batch = RecordBatch::try_new( - arrow_schema, - vec![ - Arc::new(Int32Array::from(ids)) as ArrayRef, - Arc::new(vector_builder.finish()) as ArrayRef, - ], - ) - .unwrap(); - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write.write_arrow_batch(&batch).await.unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await - .unwrap(); - assert!(built > 0, "DE fixture must build a global vector index"); - let built = table - .new_sorted_global_index_build_builder() - .with_index_column("id") - .with_index_type("btree") - .execute() - .await - .unwrap(); - assert!(built > 0, "DE fixture must build a scalar BTree index"); - table - } - - /// One Java `DataOutput#writeUTF` value (u16-BE length + modified UTF-8), used - /// to assemble the `PrimaryKeyIndexSourceMeta` frame below. - fn java_write_utf(s: &str) -> Vec { - let mut body = Vec::new(); - for c in s.encode_utf16() { - if (0x0001..=0x007F).contains(&c) { - body.push(c as u8); - } else if c > 0x07FF { - body.push(0xE0 | (c >> 12) as u8); - body.push(0x80 | ((c >> 6) & 0x3F) as u8); - body.push(0x80 | (c & 0x3F) as u8); - } else { - body.push(0xC0 | (c >> 6) as u8); - body.push(0x80 | (c & 0x3F) as u8); - } - } - let mut out = (body.len() as u16).to_be_bytes().to_vec(); - out.extend_from_slice(&body); - out - } - - /// The Java `PrimaryKeyIndexSourceMeta` frame: `i32-BE version=1`, `i32-BE - /// data_level`, `i32-BE count`, then per source file a `writeUTF` name and an - /// `i64-BE` row count. - fn pk_source_meta_bytes(data_level: i32, files: &[(&str, i64)]) -> Vec { - let mut out = Vec::new(); - out.extend_from_slice(&1i32.to_be_bytes()); - out.extend_from_slice(&data_level.to_be_bytes()); - out.extend_from_slice(&(files.len() as i32).to_be_bytes()); - for (name, rows) in files { - out.extend_from_slice(&java_write_utf(name)); - out.extend_from_slice(&rows.to_be_bytes()); - } - out - } - - /// Build a committed primary-key vector table (memory FS) over `vectors` - /// (dimension 2): write a real data file via the write path, promote its meta - /// to a compacted, non-level-0 file (the PK index-source precondition), then - /// build + commit a real vindex IVF-flat ANN segment naming that file. Single - /// bucket, `nlist = 1`, so the ANN search is exact. Returns the opened table, - /// ready for `search_pk_route`. - async fn build_committed_pk_vector_table(vectors: &[[f32; 2]]) -> Table { - use crate::spec::{GlobalIndexMeta, IndexFileMeta, VectorType}; - use crate::table::CommitMessage; - use bytes::Bytes; - use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; - use paimon_vindex_core::io::PosWriter; - - const DIM: usize = 2; - let table_path = "memory:/pk_vector_route_test"; - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Vector( - VectorType::try_new(true, DIM as u32, DataType::Float(FloatType::new())) - .unwrap(), - ), - ) - .primary_key(["id"]) - .option("bucket", "1") - .option("deletion-vectors.enabled", "true") - .option("pk-vector.index.columns", "embedding") - .option("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER) - .option("fields.embedding.pk-vector.distance.metric", "l2") - .build() - .unwrap(); - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let table = Table::new( - file_io.clone(), - Identifier::new("default", "pk_vector_route_test"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ); - for dir in ["snapshot", "manifest", "index"] { - file_io - .mkdirs(&format!("{table_path}/{dir}")) - .await - .unwrap(); - } - - // id + FixedSizeList batch matching the table's target schema. - let ids: Vec = (0..vectors.len() as i32).collect(); - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut vec_builder = FixedSizeListBuilder::new(Float32Builder::new(), DIM as i32) - .with_field(element_field.clone()); - for v in vectors { - for &x in v { - vec_builder.values().append_value(x); - } - vec_builder.append(true); - } - let arrow_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList(element_field, DIM as i32), - true, - ), - ])); - let batch = RecordBatch::try_new( - arrow_schema, - vec![ - Arc::new(Int32Array::from(ids)) as ArrayRef, - Arc::new(vec_builder.finish()) as ArrayRef, - ], - ) - .unwrap(); - - // Real data-file meta via the write path (these messages are not committed - // as-is; the meta is promoted below and committed with the index). - let mut writer = TableWrite::new(&table, "route-test".to_string()).unwrap(); - writer.write_arrow_batch(&batch).await.unwrap(); - let messages = writer.prepare_commit().await.unwrap(); - let base = &messages[0]; - let base_meta = base.new_files[0].clone(); - let bucket = base.bucket; - let partition = base.partition.clone(); - let data_file_name = base_meta.file_name.clone(); - let row_count = base_meta.row_count; - - // PK index-source precondition: compacted, non-level-0, first_row_id pinned. - let indexed_meta = DataFileMeta { - level: 1, - file_source: Some(1), - first_row_id: Some(0), - ..base_meta - }; - - // Real vindex IVF-flat segment (nlist=1 -> exact) over the vectors. - let n = vectors.len(); - let flat: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); - let seg_ids: Vec = (0..n as i64).collect(); - let native = HashMap::from([ - ("index.type".to_string(), "ivf_flat".to_string()), - ("dimension".to_string(), DIM.to_string()), - ("nlist".to_string(), "1".to_string()), - ("metric".to_string(), "l2".to_string()), - ]); - let config = VectorIndexConfig::from_options(&native).unwrap(); - let training = VectorIndexTrainer::train(config, &flat, n).unwrap(); - let mut ann_writer = VectorIndexWriter::new(training); - ann_writer.add_vectors(&seg_ids, &flat, n).unwrap(); - let mut seg_bytes = Vec::new(); - { - let mut out = PosWriter::new(&mut seg_bytes); - ann_writer.write(&mut out).unwrap(); - } - let index_file_name = "vector-ivf-flat-route.index".to_string(); - let index_file_size = seg_bytes.len() as u64; - file_io - .new_output(&format!("{table_path}/index/{index_file_name}")) - .unwrap() - .write(Bytes::from(seg_bytes)) - .await - .unwrap(); - - let vector_field_id = schema - .fields() - .iter() - .find(|f| f.name() == "embedding") - .unwrap() - .id(); - let index_file = IndexFileMeta { - index_type: IVF_FLAT_IDENTIFIER.to_string(), - file_name: index_file_name, - file_size: i64::try_from(index_file_size).unwrap(), - row_count, - deletion_vectors_ranges: None, - external_path: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: 0, - row_range_end: row_count - 1, - index_field_id: vector_field_id, - extra_field_ids: None, - source_meta: Some(pk_source_meta_bytes(1, &[(&data_file_name, row_count)])), - index_meta: None, - }), - }; - - let mut message = CommitMessage::new(partition, bucket, vec![indexed_meta]); - message.new_index_files = vec![index_file]; - TableCommit::new(table.clone(), "route-test".to_string()) - .commit(vec![message]) - .await - .unwrap(); - table - } - - // ---- search_pk_route: candidate-only producer returns candidates + context ---- - #[tokio::test] - async fn search_pk_route_returns_candidates_and_publishes_diagnostics() { - reset_vector_search_test_logs(); - let _timing = crate::vindex::enable_vector_search_timing_for_test(); - // query [0,1]: squared-L2 distances pos1=0 < pos2=1 < pos0=2, so the - // strict-gap top-2 is [pos1, pos2] (best-first, not physical order). - let table = build_committed_pk_vector_table(&[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]).await; - let core = CoreOptions::new(table.schema().options()); - let builder = table.new_vector_search_builder(); - let route = builder - .search_pk_route(&core, "embedding", &[0.0, 1.0], 2) - .await - .unwrap(); - - // Two nearest neighbours recalled, best-first, without materialization. - assert_eq!(route.candidates.len(), 2, "top-2 candidates expected"); - assert_eq!(route.candidates[0].row_position, 1, "nearest is position 1"); - assert_eq!( - route.candidates[1].row_position, 2, - "second nearest is position 2" - ); - assert!( - route.candidates[0].distance <= route.candidates[1].distance, - "candidates must be best-first by distance" - ); - - // Source context present for a non-empty plan: the snapshot the plan - // pinned during planning (a real id, not None/0), per-bucket source - // splits, and the resolved metric. - assert_eq!(route.snapshot_id, 1, "first commit -> snapshot 1"); - assert!(!route.splits.is_empty(), "non-empty source splits expected"); - assert_eq!(route.metric, VectorSearchMetric::L2); - assert!( - route - .candidates - .iter() - .all(|c| c.split_index < route.splits.len()), - "candidate split_index must refer into the returned splits" - ); - - let logs = VECTOR_SEARCH_TEST_LOGS.lock().unwrap(); - assert!( - logs.iter().any(|entry| { - entry.contains("event=paimon_vindex_reader") - && entry.contains("vector-ivf-flat-route.index") - }), - "PK vector search must publish vindex reader timing" - ); - assert!( - logs.iter().any(|entry| { - entry.contains("event=paimon_vector_range_io") - && entry.contains("vector-ivf-flat-route.index") - }), - "PK vector search must publish range-I/O timing" - ); - } - - /// A table with no snapshot at all (never written) yields empty candidates - /// and empty source splits; with no snapshot to pin the id is `0`, and the - /// metric still resolves. - #[tokio::test] - async fn search_pk_route_empty_plan_yields_empty_context() { - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let core = CoreOptions::new(table.schema().options()); - let builder = table.new_vector_search_builder(); - let route = builder - .search_pk_route(&core, "embedding", &[1.0f32; 128], 3) - .await - .unwrap(); - assert!(route.candidates.is_empty(), "no data -> no candidates"); - assert!(route.splits.is_empty(), "no data -> no source splits"); - assert_eq!(route.snapshot_id, 0, "no snapshot -> zero id"); - assert_eq!(route.metric, VectorSearchMetric::L2); - } - - #[tokio::test] - async fn pk_branch_disabled_falls_through_to_de_path() { - // No pk-vector.index.columns: behaves exactly as the DE path. With no - // snapshot the DE path returns an empty result; the PK branch must not - // intercept it. - let table = pk_vector_table(&[]); - let result = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_scored() - .await - .unwrap(); - assert!(result.is_empty()); - } - - #[tokio::test] - async fn pk_branch_execute_scored_fails_loud() { - // On a PK-vector table `execute_scored` reports global row ids, which the - // PK path cannot produce (physical (file, position) coords, no global ids). - // It must fail loud rather than fabricate ids; callers use `execute_read`. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let err = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_scored() - .await - .map(|_| ()) - .expect_err("execute_scored on a PK-vector column must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("does not produce global row ids")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn pk_branch_other_column_falls_through_to_de_path() { - // pk-vector index configured for "embedding", but the query targets a - // different column -> the PK branch must not intercept; DE path (no - // snapshot) yields empty. Discriminator: the PK column carries a - // DELIBERATELY INVALID distance metric, which the PK branch parses eagerly - // (`VectorSearchMetric::parse`) and would fail on. So a regression that - // dropped the `pk_col == vector_column` guard and ran the PK branch for - // "other" would surface as Err here, not Ok(empty) -- the assertion - // therefore proves the DE path ran, not merely that the result is empty. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ( - "fields.embedding.pk-vector.distance.metric", - "not-a-real-metric", - ), - ]); - let result = table - .new_vector_search_builder() - .with_vector_column("other") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_scored() - .await - .unwrap(); - assert!(result.is_empty()); - } - - #[tokio::test] - async fn pk_branch_multi_column_config_does_not_break_unrelated_de_query() { - // A malformed multi-column PK-vector config ("a,b") must not abort an - // unrelated DE vector query. The query targets a column NOT among the - // configured PK-vector columns, so membership resolution short-circuits - // before the exactly-one-column rule fires -- the query falls through to - // the DE path (no snapshot -> empty) instead of surfacing the "must name - // exactly one column" error. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "a,b"), - ("fields.a.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.a.pk-vector.distance.metric", "l2"), - ]); - let result = table - .new_vector_search_builder() - .with_vector_column("other") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_scored() - .await; - match result { - Ok(search) => assert!(search.is_empty()), - Err(err) => panic!( - "unrelated DE query must not error on a malformed multi-column PK config: {err}" - ), - } - } - - /// `id > threshold` built against the table's user fields (leaf index resolves - /// against `table.schema().fields()`). - fn id_gt_filter(table: &Table, threshold: i32) -> Predicate { - PredicateBuilder::new(table.schema().fields()) - .greater_than("id", Datum::Int(threshold)) - .unwrap() - } - - /// The vector residual is derived from the DATA conjuncts of the filter: - /// partition-only conjuncts are enforced by scan planning (`PkVectorScan` - /// pushes the whole filter through the normal scan) and must not enter the - /// per-row residual, so a partition-only filter yields no residual at all. - #[test] - fn residual_uses_only_data_conjuncts_of_the_filter() { - use crate::spec::VarCharType; - use crate::table::bucket_filter::split_partition_and_data_predicates; - - // Partitioned table: `dt` (partition key) + `id`. - let schema = Schema::builder() - .column("dt", DataType::VarChar(VarCharType::string_type())) - .column("id", DataType::Int(IntType::new())) - .partition_keys(["dt"]) - .build() - .unwrap(); - let ts = TableSchema::new(0, &schema); - let fields = ts.fields(); - let partition_keys = ts.partition_keys(); - let pb = PredicateBuilder::new(fields); - - // Partition-only `dt = 'a'` -> no residual data predicate (residual skipped; - // the partition is enforced by planning alone). - let (_p, data) = split_partition_and_data_predicates( - pb.equal("dt", Datum::String("a".to_string())).unwrap(), - fields, - partition_keys, - ); - assert!( - data.is_empty(), - "partition-only filter must leave no residual data predicate" - ); - - // Data-only `id > 5` -> kept as the residual. - let (_p, data) = split_partition_and_data_predicates( - pb.greater_than("id", Datum::Int(5)).unwrap(), - fields, - partition_keys, - ); - assert_eq!(data.len(), 1, "data-only filter must remain the residual"); - - // `dt = 'a' AND id > 5` -> only the data conjunct enters the residual. - let (_p, data) = split_partition_and_data_predicates( - Predicate::and(vec![ - pb.equal("dt", Datum::String("a".to_string())).unwrap(), - pb.greater_than("id", Datum::Int(5)).unwrap(), - ]), - fields, - partition_keys, - ); - assert_eq!( - data.len(), - 1, - "AND(partition, data) residual must drop the partition conjunct" - ); - - // `dt = 'a' OR id > 5` is a single mixed conjunct: it is NOT partition-only, - // so it stays whole in the residual (evaluated against the materialized - // partition column), rather than being dropped or split. - let mixed = Predicate::or(vec![ - pb.equal("dt", Datum::String("a".to_string())).unwrap(), - pb.greater_than("id", Datum::Int(5)).unwrap(), - ]); - let (_p, data) = split_partition_and_data_predicates(mixed.clone(), fields, partition_keys); - assert_eq!( - data, - vec![mixed], - "a mixed partition/data conjunct must stay whole in the residual" - ); - } - - #[tokio::test] - async fn execute_read_filter_without_deletion_vectors_fails_loud() { - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let filter = id_gt_filter(&table, 2); - let err = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .with_filter(filter) - .execute_read() - .await - .map(|_| ()) - .expect_err("read filter without deletion vectors must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("deletion vectors without merge-on-read")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn execute_scored_filter_on_empty_de_path_returns_empty() { - // No PK-vector index and no snapshot: the request follows the - // data-evolution path. Scalar pre-filter support must not turn an empty - // table into an error. - let table = pk_vector_table(&[]); - let filter = id_gt_filter(&table, 2); - let result = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .with_filter(filter) - .execute_scored() - .await - .expect("an empty data-evolution search with a filter should succeed"); - assert!(result.is_empty()); - } - - #[tokio::test] - async fn execute_read_filter_with_merge_on_read_fails_loud() { - // Deletion vectors enabled BUT merge-on-read on: still rejected, because a - // merge-on-read scan can surface stale key versions that a physical-row - // filter cannot reconcile. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ("deletion-vectors.enabled", "true"), - ("deletion-vectors.merge-on-read", "true"), - ]); - let filter = id_gt_filter(&table, 2); - let err = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .with_filter(filter) - .execute_read() - .await - .map(|_| ()) - .expect_err("merge-on-read filter must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("deletion vectors without merge-on-read")), - "unexpected error: {err:?}" - ); - } - - #[tokio::test] - async fn execute_read_filter_with_deletion_vectors_passes_guard() { - // Deletion vectors enabled, merge-on-read off (default): the residual guard - // passes. With no snapshot the plan is empty, so the (guarded) filter path - // simply yields an empty stream rather than erroring — proving the guard - // admits a legal filtered query. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ("deletion-vectors.enabled", "true"), - // Pin the index dimension so the query vector below matches it; the - // up-front dimension guard runs before this test's residual guard. - ("fields.embedding.dimension", "4"), - ]); - let filter = id_gt_filter(&table, 2); - let mut stream = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .with_filter(filter) - .execute_read() - .await - .expect("guarded filter query must be admitted"); - assert!(stream.try_next().await.unwrap().is_none()); - } - - /// A partition-only `with_filter` needs no per-row residual (partition pruning - /// happens in scan planning), so the deletion-vector pre-filter guard must NOT - /// reject it even when deletion vectors are off. Mirrors Java, where a - /// partition-only filter leaves `this.filter == null` and the scan guard is - /// skipped. Regression test for the guard keying on the whole filter rather - /// than its data conjuncts. - #[tokio::test] - async fn execute_read_partition_only_filter_without_deletion_vectors_passes_guard() { - use crate::spec::VarCharType; - - // Partitioned PK-vector table from old metadata, deletion vectors OFF. - let schema = Schema::builder() - .column("dt", DataType::VarChar(VarCharType::string_type())) - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .partition_keys(["dt"]) - .primary_key(["id"]) - .option("bucket", "1") - .build() - .unwrap(); - let table_schema = TableSchema::new(0, &schema).copy_with_options(HashMap::from([ - ( - "pk-vector.index.columns".to_string(), - "embedding".to_string(), - ), - ( - "fields.embedding.pk-vector.index.type".to_string(), - IVF_FLAT_IDENTIFIER.to_string(), - ), - ( - "fields.embedding.pk-vector.distance.metric".to_string(), - "l2".to_string(), - ), - ("fields.embedding.dimension".to_string(), "4".to_string()), - ])); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "pk_vector_partitioned"), - "memory:/pk_vector_partitioned".to_string(), - table_schema, - None, - ); - - // Partition-only `dt = 'a'`: no data residual, so the guard admits it and - // (with no snapshot) the query yields an empty stream instead of the - // deletion-vector error. - let filter = PredicateBuilder::new(table.schema().fields()) - .equal("dt", Datum::String("a".to_string())) - .unwrap(); - let mut stream = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .with_filter(filter) - .execute_read() - .await - .expect("partition-only filter must be admitted without deletion vectors"); - assert!(stream.try_next().await.unwrap().is_none()); - - // But a DATA conjunct (`id > 2`) on the same non-DV table must still fail - // loud — the guard now keys on data predicates, not the whole filter. - let data_filter = id_gt_filter(&table, 2); - let err = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .with_filter(data_filter) - .execute_read() - .await - .map(|_| ()) - .expect_err("data filter without deletion vectors must still fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("deletion vectors without merge-on-read")), - "unexpected error: {err:?}" - ); - - // `AND(partition, data)` still has a data conjunct after the split, so it - // must fail loud on the non-DV table just like the data-only filter. - let pb = PredicateBuilder::new(table.schema().fields()); - let and_filter = Predicate::and(vec![ - pb.equal("dt", Datum::String("a".to_string())).unwrap(), - pb.greater_than("id", Datum::Int(2)).unwrap(), - ]); - let err = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .with_filter(and_filter) - .execute_read() - .await - .map(|_| ()) - .expect_err("AND(partition, data) without deletion vectors must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("deletion vectors without merge-on-read")), - "unexpected error: {err:?}" - ); - - // A mixed `OR(partition, data)` conjunct is not partition-only, so it stays - // whole as a data predicate and must also fail loud without deletion vectors. - let or_filter = Predicate::or(vec![ - pb.equal("dt", Datum::String("a".to_string())).unwrap(), - pb.greater_than("id", Datum::Int(2)).unwrap(), - ]); - let err = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .with_filter(or_filter) - .execute_read() - .await - .map(|_| ()) - .expect_err("mixed OR(partition, data) without deletion vectors must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("deletion vectors without merge-on-read")), - "unexpected error: {err:?}" - ); - } - - fn make_lumina_entry( - file_name: &str, - index_type: &str, - kind: FileKind, - index_field_id: i32, - ) -> IndexManifestEntry { - IndexManifestEntry { - kind, - partition: vec![], - bucket: 0, - index_file: IndexFileMeta { - index_type: index_type.to_string(), - file_name: file_name.to_string(), - file_size: 100, - row_count: 10, - deletion_vectors_ranges: None, - external_path: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: 0, - row_range_end: 9, - index_field_id, - extra_field_ids: None, - source_meta: None, - index_meta: None, - }), - }, - version: 1, - } - } - - // ---- Task B: search-and-read (`execute_read`) tests ---- - - /// Build a small materialization batch: user column `id: Int32`, the internal - /// `_PKEY_VECTOR_POSITION: Int64`, and `__paimon_search_score: Float32` (mirroring - /// what `PkVectorIndexedSplitRead` emits for a single file). - fn materialized_batch(rows: &[(i32, i64, f32)]) -> RecordBatch { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new(PKEY_VECTOR_POSITION_COLUMN, ArrowDataType::Int64, false), - ArrowField::new(SEARCH_SCORE_COLUMN, ArrowDataType::Float32, false), - ])); - let ids = Int32Array::from(rows.iter().map(|(id, _, _)| *id).collect::>()); - let positions = Int64Array::from(rows.iter().map(|(_, pos, _)| *pos).collect::>()); - let scores = Float32Array::from(rows.iter().map(|(_, _, s)| *s).collect::>()); - RecordBatch::try_new( - schema, - vec![Arc::new(ids), Arc::new(positions), Arc::new(scores)], - ) - .unwrap() - } - - fn i32_col(batch: &RecordBatch, name: &str) -> Vec { - let idx = batch.schema().index_of(name).unwrap(); - batch - .column(idx) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec() - } - fn f32_col(batch: &RecordBatch, name: &str) -> Vec { - let idx = batch.schema().index_of(name).unwrap(); - batch - .column(idx) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec() - } - - #[test] - fn reorder_and_strip_position_recovers_best_first_and_drops_position() { - // Single file, one bucket. The materialization reader emits rows in - // ascending physical position [pos0, pos1, pos2] -> ids [40,41,42]. The - // search candidates ranked them best-first as pos1(rank0), pos2(rank1), - // pos0(rank2), which is NEITHER position order nor score order-by-batch. - // The reorder must yield ids [41,42,40] and drop _PKEY_VECTOR_POSITION. - let batch = materialized_batch(&[ - (40, 0, l2_score(9.0)), - (41, 1, l2_score(1.0)), - (42, 2, l2_score(4.0)), - ]); - let batches = vec![batch]; - let part = BinaryRow::new(0).to_serialized_bytes(); - let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); - rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 1), 0); - rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 2), 1); - rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 0), 2); - - let mut ranked = Vec::new(); - collect_ranked_rows(&batches[0], 0, &part, 0, "o.mosaic", &rank_of, &mut ranked).unwrap(); - let out = reorder_and_strip_position(&batches, ranked).unwrap(); - assert_eq!(out.len(), 1); - let out = &out[0]; - - // Best-first row order, not ascending position order. - assert_eq!(i32_col(out, "id"), vec![41, 42, 40]); - // Score column preserved and aligned to the reordered rows. - assert_eq!( - f32_col(out, SEARCH_SCORE_COLUMN), - vec![l2_score(1.0), l2_score(4.0), l2_score(9.0)] - ); - // Position column dropped; _ROW_ID never present. - assert!(out.schema().index_of(PKEY_VECTOR_POSITION_COLUMN).is_err()); - assert!(out.schema().index_of("_ROW_ID").is_err()); - } - - #[test] - fn reorder_and_strip_position_merges_rows_across_files() { - // Two files (two materialization batches). Best-first interleaves them: - // file-b pos0 (rank0), file-a pos1 (rank1), file-a pos0 (rank2). The - // reorder must pull rows from both batches into one best-first output. - let batch_a = materialized_batch(&[(10, 0, l2_score(9.0)), (11, 1, l2_score(1.0))]); - let batch_b = materialized_batch(&[(20, 0, l2_score(0.5))]); - let batches = vec![batch_a, batch_b]; - let part = BinaryRow::new(0).to_serialized_bytes(); - let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); - rank_of.insert((part.clone(), 0, "b".to_string(), 0), 0); - rank_of.insert((part.clone(), 0, "a".to_string(), 1), 1); - rank_of.insert((part.clone(), 0, "a".to_string(), 0), 2); - - let mut ranked = Vec::new(); - collect_ranked_rows(&batches[0], 0, &part, 0, "a", &rank_of, &mut ranked).unwrap(); - collect_ranked_rows(&batches[1], 1, &part, 0, "b", &rank_of, &mut ranked).unwrap(); - let out = reorder_and_strip_position(&batches, ranked).unwrap(); - assert_eq!(i32_col(&out[0], "id"), vec![20, 11, 10]); - assert_eq!( - f32_col(&out[0], SEARCH_SCORE_COLUMN), - vec![l2_score(0.5), l2_score(1.0), l2_score(9.0)] - ); - } - - #[test] - fn reorder_and_strip_position_empty_yields_no_batches() { - let out = reorder_and_strip_position(&[], Vec::new()).unwrap(); - assert!(out.is_empty()); - } - - #[test] - fn collect_ranked_rows_missing_candidate_fails_loud() { - // A materialized position with no candidate rank must fail loud rather than - // silently drop the row. - let batch = materialized_batch(&[(40, 7, l2_score(1.0))]); - let part = BinaryRow::new(0).to_serialized_bytes(); - let rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); - let mut ranked = Vec::new(); - let err = collect_ranked_rows(&batch, 0, &part, 0, "f", &rank_of, &mut ranked) - .expect_err("missing candidate must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("no matching search candidate")), - "unexpected error: {err:?}" - ); - } - - #[test] - fn attach_scores_reorders_by_rank_not_score() { - use arrow_array::{Int32Array, Int64Array, RecordBatch}; - use arrow_schema::{DataType, Field, Schema}; - use std::sync::Arc; - - // Two rows materialized in row-id order [10, 20]; ranks say 20 is best (rank 0), - // 10 is rank 1. Scores tie at 0.5 to prove ordering follows rank, not score. - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new(ROW_ID_FIELD_NAME, DataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![100, 200])), - Arc::new(Int64Array::from(vec![10, 20])), - ], - ) - .unwrap(); - let mut map = HashMap::new(); - map.insert(20i64, (0usize, 0.5f32)); - map.insert(10i64, (1usize, 0.5f32)); - - let out = attach_scores_by_row_id(&[batch], &map, 2).unwrap(); - assert_eq!(out.len(), 1); - let b = &out[0]; - // _ROW_ID stripped, score appended. - assert!(b.schema().index_of(ROW_ID_FIELD_NAME).is_err()); - let score_idx = b.schema().index_of("__paimon_search_score").unwrap(); - assert_eq!( - b.schema().field(score_idx).data_type(), - &arrow_schema::DataType::Float32 - ); - // Row order is rank order: id 200 (rank 0) first, then id 100 (rank 1). - let ids = b.column(0).as_any().downcast_ref::().unwrap(); - assert_eq!(ids.values(), &[200, 100]); - } - - #[test] - fn attach_scores_fails_on_unknown_row_id() { - use arrow_array::{Int32Array, Int64Array, RecordBatch}; - use arrow_schema::{DataType, Field, Schema}; - use std::sync::Arc; - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new(ROW_ID_FIELD_NAME, DataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![1])), - Arc::new(Int64Array::from(vec![99])), - ], - ) - .unwrap(); - let map: HashMap = HashMap::new(); // no entry for 99 - let err = attach_scores_by_row_id(&[batch], &map, 1).unwrap_err(); - assert!(matches!(err, crate::Error::DataInvalid { .. })); - } - - #[test] - fn attach_scores_fails_on_count_mismatch() { - use arrow_array::{Int32Array, Int64Array, RecordBatch}; - use arrow_schema::{DataType, Field, Schema}; - use std::sync::Arc; - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new(ROW_ID_FIELD_NAME, DataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![1])), - Arc::new(Int64Array::from(vec![10])), - ], - ) - .unwrap(); - let mut map = HashMap::new(); - map.insert(10i64, (0usize, 0.5f32)); - // expected_len 2 but only 1 row materialized. - let err = attach_scores_by_row_id(&[batch], &map, 2).unwrap_err(); - assert!(matches!(err, crate::Error::DataInvalid { .. })); - } - - #[test] - fn attach_scores_fails_on_null_row_id() { - use arrow_array::{Int32Array, Int64Array, RecordBatch}; - use arrow_schema::{DataType, Field, Schema}; - use std::sync::Arc; - // _ROW_ID column has a NULL at row 1; the map contains the non-null id, so - // the failure is specifically the null (not an unknown id). - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new(ROW_ID_FIELD_NAME, DataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(Int64Array::from(vec![Some(10i64), None])), - ], - ) - .unwrap(); - let mut map = HashMap::new(); - map.insert(10i64, (0usize, 0.5f32)); - let err = attach_scores_by_row_id(&[batch], &map, 2).unwrap_err(); - assert!(matches!(err, crate::Error::DataInvalid { .. })); - } - - #[test] - fn attach_scores_fails_on_wrong_type_row_id() { - use arrow_array::{Int32Array, RecordBatch}; - use arrow_schema::{DataType, Field, Schema}; - use std::sync::Arc; - // _ROW_ID column is Int32, not Int64: the downcast fails loud. - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new(ROW_ID_FIELD_NAME, DataType::Int32, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![1])), - Arc::new(Int32Array::from(vec![10])), - ], - ) - .unwrap(); - let mut map = HashMap::new(); - map.insert(10i64, (0usize, 0.5f32)); - let err = attach_scores_by_row_id(&[batch], &map, 1).unwrap_err(); - assert!(matches!(err, crate::Error::DataInvalid { .. })); - } - - #[tokio::test] - async fn execute_read_de_table_empty_snapshot_yields_empty_stream() { - // No pk-vector index configured and no snapshot: execute_read routes to the - // data-evolution path, whose search finds nothing and returns an empty - // stream (not an error). - let table = pk_vector_table(&[]); - let mut stream = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_read() - .await - .expect("DE read over an empty table must succeed with no rows"); - let mut rows = 0usize; - while let Some(batch) = stream.try_next().await.unwrap() { - rows += batch.num_rows(); - } - assert_eq!(rows, 0, "empty DE table must yield no rows"); - } - - #[tokio::test] - async fn execute_read_unknown_column_fails_loud() { - // pk-vector index configured for "embedding", but the query targets a - // column that does not exist. The read path must fail loud rather than - // fall through to the data-evolution path and return an empty stream (a - // typo must not look like a normal empty read through the C API). - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let err = match table - .new_vector_search_builder() - .with_vector_column("other") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_read() - .await - { - Ok(_) => panic!("unknown vector column must fail loud on execute_read"), - Err(e) => e, - }; - assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("does not exist")), - "expected a does-not-exist error, got: {err}" - ); - } - - #[tokio::test] - async fn execute_read_scalar_column_fails_loud() { - // A scalar (non-vector) column targeted by a vector read must fail loud, - // not return an empty data-evolution stream. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let err = match table - .new_vector_search_builder() - .with_vector_column("id") // scalar Int column - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_read() - .await - { - Ok(_) => panic!("scalar vector column must fail loud on execute_read"), - Err(e) => e, - }; - assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("must be a FLOAT vector column")), - "expected a not-a-vector-column error, got: {err}" - ); - } - - #[tokio::test] - async fn execute_read_non_float_vector_column_fails_loud() { - // An ARRAY column is not a searchable vector column (the index/search - // operates on FLOAT elements). It must fail loud rather than fall through - // to the DE path and return an empty stream. - use crate::spec::{ArrayType, IntType, Schema, TableSchema}; - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), - ) - .build() - .unwrap(); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "de_non_float_vector"), - "memory:/de_non_float_vector".to_string(), - TableSchema::new(0, &schema), - None, - ); - let err = match table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_read() - .await - { - Ok(_) => panic!("ARRAY vector column must fail loud on execute_read"), - Err(e) => e, - }; - assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("must be a FLOAT vector column")), - "expected a FLOAT-vector-column error, got: {err}" - ); - } - - #[tokio::test] - async fn execute_read_empty_plan_reserved_projection_fails_loud() { - // Empty plan (no snapshot) must still fail loud on a reserved-name - // projection: projection validity does not depend on whether the search - // matched any rows. A regression that resolved the projection only after - // the `candidates.is_empty()` early return would yield an empty stream here - // instead of an error. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - // Pin the index dimension so the query vector below matches it; the - // up-front dimension guard runs before this test's reserved-projection - // guard, so a mismatched query would mask the error under test. - ("fields.embedding.dimension", "4"), - ]); - for reserved in [ - ROW_ID_FIELD_NAME, - PKEY_VECTOR_POSITION_COLUMN, - SEARCH_SCORE_COLUMN, - ] { - let mut builder = table.new_vector_search_builder(); - builder - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .with_projection(&["id", reserved]); - let err = builder - .execute_read() - .await - .map(|_| ()) - .expect_err("empty plan + reserved projection must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("reserved column")), - "unexpected error for {reserved}: {err:?}" - ); - } - } - - #[tokio::test] - async fn execute_read_empty_plan_lumina_array_float_is_admitted() { - // A Lumina PK-vector `ARRAY` column is a valid configuration, but - // batch query dimension validation routed every `ARRAY` column - // through the vindex resolver, which rejects `lumina` as an unsupported - // index type before planning — failing even an empty table. The - // dimension must be resolved per the configured backend, so a - // well-formed Lumina query is admitted and (with no snapshot) yields an - // empty stream rather than an "Unsupported vindex index type" error. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ( - "fields.embedding.pk-vector.index.type", - crate::lumina::LUMINA_IDENTIFIER, - ), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ("lumina.index.dimension", "4"), - ]); - let mut stream = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0; 4]) - .with_limit(5) - .execute_read() - .await - .expect( - "Lumina ARRAY query must be admitted, not rejected as unsupported vindex", - ); - assert!(stream.try_next().await.unwrap().is_none()); - } - - #[tokio::test] - async fn execute_read_projection_reserved_name_fails_loud() { - // Projecting a reserved metadata / row-id column must fail loud. The guard - // lives in `resolve_materialize_read_type`, which `execute_read` invokes - // before the empty-plan early return; assert on the resolver directly here. - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - for reserved in [ - ROW_ID_FIELD_NAME, - PKEY_VECTOR_POSITION_COLUMN, - SEARCH_SCORE_COLUMN, - ] { - let mut builder = table.new_vector_search_builder(); - builder - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .with_projection(&["id", reserved]); - let err = builder - .resolve_materialize_read_type() - .expect_err("reserved projection must fail loud"); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("reserved column")), - "unexpected error for {reserved}: {err:?}" - ); - } - } - - #[test] - fn resolve_materialize_read_type_default_is_all_user_columns() { - // No with_projection -> every user table column (id + embedding). - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let builder = table.new_vector_search_builder(); - let fields = builder.resolve_materialize_read_type().unwrap(); - let names: Vec<&str> = fields.iter().map(|f| f.name()).collect(); - assert_eq!(names, vec!["id", "embedding"]); - } - - /// A PK-vector table whose user schema carries an extra column named - /// `reserved`, used to prove reserved metadata names are rejected even when - /// they arrive via the default (all-columns) projection. - fn pk_vector_table_with_extra_column(reserved: &str) -> Table { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .column(reserved, DataType::Int(IntType::new())) - .primary_key(["id"]) - .option("bucket", "1") - .option("deletion-vectors.enabled", "true") - .option("pk-vector.index.columns", "embedding") - .option("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER) - .option("fields.embedding.pk-vector.distance.metric", "l2") - .build() - .unwrap(); - Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "reserved_col_test"), - "memory:/reserved_col_test".to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - #[test] - fn resolve_materialize_read_type_default_rejects_reserved_user_column() { - // The default (all-columns) projection must reject a user column whose - // name collides with an injected metadata column, not only columns named - // in an explicit projection. Otherwise it silently passes on an empty - // result and collides with the metadata columns the read attaches. - let table = pk_vector_table_with_extra_column(SEARCH_SCORE_COLUMN); - let builder = table.new_vector_search_builder(); - let err = builder.resolve_materialize_read_type().unwrap_err(); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("reserved column")), - "single-query default projection must reject reserved user column, got: {err:?}" - ); - } - - #[test] - fn batch_resolve_materialize_read_type_default_rejects_reserved_user_column() { - // Same guard on the batch resolver. - let table = pk_vector_table_with_extra_column(PKEY_VECTOR_POSITION_COLUMN); - let builder = table.new_batch_vector_search_builder(); - let err = builder.resolve_materialize_read_type().unwrap_err(); - assert!( - matches!(err, crate::Error::DataInvalid { ref message, .. } - if message.contains("reserved column")), - "batch default projection must reject reserved user column, got: {err:?}" - ); - } - - #[test] - fn resolve_materialize_read_type_projection_selects_named_columns() { - let table = pk_vector_table(&[ - ("pk-vector.index.columns", "embedding"), - ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), - ("fields.embedding.pk-vector.distance.metric", "l2"), - ]); - let mut builder = table.new_vector_search_builder(); - builder.with_projection(&["id"]); - let fields = builder.resolve_materialize_read_type().unwrap(); - let names: Vec<&str> = fields.iter().map(|f| f.name()).collect(); - assert_eq!(names, vec!["id"]); - } - - #[tokio::test] - async fn de_execute_read_materializes_rows_with_score() { - // A data-evolution vector table with a committed global index: execute_read - // must materialize one row per scored hit and carry the unified score - // column, in best-first rank order. - let table = de_vector_table().await; - let query = vec![1.0, 0.0]; - - let scored = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(query.clone()) - .with_limit(3) - .execute_scored() - .await - .unwrap(); - assert!(!scored.is_empty(), "DE search must return hits"); - - let mut stream = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(query) - .with_limit(3) - .execute_read() - .await - .unwrap(); - - let mut rows = 0usize; - let mut saw_score = false; - while let Some(batch) = stream.try_next().await.unwrap() { - rows += batch.num_rows(); - saw_score |= batch.schema().index_of(SEARCH_SCORE_COLUMN).is_ok(); - } - assert_eq!( - rows, - scored.len(), - "DE read must emit exactly the scored result count" - ); - assert!( - saw_score, - "DE read output must carry the search score column" - ); - } - - #[tokio::test] - async fn de_vector_search_uses_time_travel_snapshot() { - let table = de_vector_table().await; - let latest = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(3) - .execute_scored() - .await - .unwrap(); - assert!( - !latest.is_empty(), - "latest snapshot should contain the committed vector index" - ); - - let traveled = table - .copy_with_time_travel(HashMap::from([( - crate::spec::SCAN_VERSION_OPTION.to_string(), - "1".to_string(), - )])) - .await - .unwrap(); - assert_eq!( - traveled.travel_snapshot().map(|snapshot| snapshot.id()), - Some(1) - ); - - let historical = traveled - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(3) - .execute_scored() - .await - .unwrap(); - assert!( - historical.is_empty(), - "snapshot 1 predates the vector index and should return no hits" - ); - } - - #[tokio::test] - async fn resolved_vector_snapshot_can_be_reused_by_all_read_stages() { - let table = de_vector_table().await; - let snapshot = crate::table::time_travel::resolve_snapshot(&table) - .await - .unwrap() - .unwrap(); - let pinned = table.copy_with_resolved_snapshot(&snapshot).await.unwrap(); - - assert_eq!( - pinned.travel_snapshot().map(|snapshot| snapshot.id()), - Some(snapshot.id()) - ); - let options = CoreOptions::new(pinned.schema().options()); - let selector = options.try_time_travel_selector().unwrap().unwrap(); - assert!(matches!( - selector, - crate::spec::TimeTravelSelector::SnapshotId { - value, - option_name: crate::spec::SCAN_SNAPSHOT_ID_OPTION, - } if value == snapshot.id().to_string() - )); - } - - #[tokio::test] - async fn de_execute_read_applies_scalar_filter_before_top_k() { - // Row id=1 is the closest vector to [1, 0], but the scalar filter excludes - // it. Filter-before-Top-K must return the best rows among ids > 1 instead - // of recalling id=1 first and filtering it after the search. - let table = de_vector_table().await; - let filter = id_gt_filter(&table, 1); - let mut stream = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(2) - .with_filter(filter) - .execute_read() - .await - .expect("DE vector search should support a scalar pre-filter"); - - let mut ids = Vec::new(); - while let Some(batch) = stream.try_next().await.unwrap() { - let id = batch - .column_by_name("id") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - ids.extend((0..id.len()).map(|row| id.value(row))); - } - - assert_eq!(ids, vec![3, 2]); - } - - #[tokio::test] - async fn de_scalar_filter_with_no_matching_rows_returns_empty() { - let table = de_vector_table().await; - let filter = id_gt_filter(&table, 99); - - let result = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(2) - .with_filter(filter.clone()) - .execute_scored() - .await - .unwrap(); - assert!(result.is_empty()); - - let results = table - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) - .with_limit(2) - .with_filter(filter) - .execute() - .await - .unwrap(); - assert_eq!(results.len(), 2); - assert!(results.iter().all(SearchResult::is_empty)); - } - - #[tokio::test] - async fn prepared_de_scalar_filter_can_be_reused_by_batch_search() { - let table = de_vector_table().await; - let prepared = table - .prepare_vector_search_filter(id_gt_filter(&table, 1)) - .await - .unwrap(); - let results = table - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) - .with_limit(2) - .with_prepared_filter(prepared) - .execute() - .await - .unwrap(); - - assert_eq!(results.len(), 2); - assert_eq!(results[0].row_ids, vec![2, 1]); - assert_eq!(results[1].row_ids, vec![1, 2]); - } - - #[tokio::test] - async fn prepared_filter_from_different_table_is_rejected() { - let prepared = PreparedVectorSearchFilter { - table: vector_test_table_at("memory:/prepared_filter_source"), - include_row_ids: Arc::new(RoaringTreemap::from_iter([1])), - }; - let target = vector_test_table_at("memory:/prepared_filter_target"); - - let error = target - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0, 0.0]]) - .with_limit(1) - .with_prepared_filter(prepared) - .execute() - .await - .expect_err("a prepared filter must not retarget the builder to another table"); - - assert!( - error.to_string().contains("different table"), - "unexpected error: {error}" - ); - } - - #[tokio::test] - async fn de_scalar_filter_applies_to_unindexed_raw_fallback() { - let table = de_vector_table().await; - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut vector_builder = - ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); - vector_builder.values().append_value(1.0); - vector_builder.values().append_value(0.0); - vector_builder.append(true); - let batch = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("embedding", ArrowDataType::List(element_field), true), - ])), - vec![ - Arc::new(Int32Array::from(vec![4])) as ArrayRef, - Arc::new(vector_builder.finish()) as ArrayRef, - ], - ) - .unwrap(); - let mut writer = TableWrite::new(&table, "test-user".to_string()).unwrap(); - writer.write_arrow_batch(&batch).await.unwrap(); - let messages = writer.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let table = table.copy_with_options(HashMap::from([ - ("vector-index.search-mode".to_string(), "full".to_string()), - ("scalar-index.search-mode".to_string(), "full".to_string()), - ])); - let result = table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(1) - .with_filter(id_gt_filter(&table, 3)) - .execute_scored() - .await - .unwrap(); - - assert_eq!(result.row_ids, vec![3]); - } -} - -/// Tests for [`residual_positions_by_file`]: the residual predicate is applied at -/// the Arrow level (no pushdown) against the predicate columns, and each surviving -/// row's file-local physical position is recovered from its ordinal in the -/// unfiltered scan (no `_ROW_ID`, no `first_row_id`). #[cfg(test)] -mod residual_positions_tests { - use super::*; - use crate::arrow::build_target_arrow_schema; - use crate::arrow::format::FilePredicates; - use crate::io::FileIOBuilder; - use crate::spec::stats::BinaryTableStats; - use crate::spec::{ - BigIntType, BinaryRow, DataField, DataFileMeta, DataType, Datum, IntType, PredicateBuilder, - ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, - }; - use crate::table::data_file_reader::DataFileReader; - use crate::table::schema_manager::SchemaManager; - use crate::table::source::{DataSplit, DataSplitBuilder}; - use arrow_array::{Int32Array, RecordBatch}; - use bytes::Bytes; - use paimon_mosaic_core::spec::COMPRESSION_NONE; - use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions}; - use std::io; - use std::sync::Arc; - - struct MemOutputFile { - data: Vec, - } - - impl OutputFile for MemOutputFile { - fn write(&mut self, data: &[u8]) -> io::Result<()> { - self.data.extend_from_slice(data); - Ok(()) - } - fn flush(&mut self) -> io::Result<()> { - Ok(()) - } - fn pos(&self) -> u64 { - self.data.len() as u64 - } - } - - fn id_field() -> DataField { - DataField::new(0, "id".to_string(), DataType::Int(IntType::new())) - } - - fn row_id_field() -> DataField { - DataField::new( - ROW_ID_FIELD_ID, - ROW_ID_FIELD_NAME.to_string(), - DataType::BigInt(BigIntType::new()), - ) - } - - fn id_batch(ids: Vec) -> RecordBatch { - let schema = build_target_arrow_schema(&[id_field()]).unwrap(); - RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(ids))]).unwrap() - } - - fn write_mosaic(batch: &RecordBatch) -> Bytes { - let mut writer = MosaicWriter::new( - MemOutputFile { data: Vec::new() }, - batch.schema().as_ref(), - WriterOptions { - compression: COMPRESSION_NONE, - num_buckets: 2, - row_group_max_size: u64::MAX, - ..Default::default() - }, - ) - .unwrap(); - writer.write_batch(batch).unwrap(); - writer.close().unwrap(); - Bytes::from(writer.output().data.to_vec()) - } - - fn data_file( - file_name: &str, - file_size: i64, - row_count: i64, - first_row_id: Option, - ) -> DataFileMeta { - DataFileMeta { - file_name: file_name.to_string(), - file_size, - row_count, - min_key: Vec::new(), - max_key: Vec::new(), - key_stats: BinaryTableStats::empty(), - value_stats: BinaryTableStats::empty(), - min_sequence_number: 0, - max_sequence_number: 0, - schema_id: 1, - level: 0, - extra_files: Vec::new(), - creation_time: None, - delete_row_count: None, - embedded_index: None, - file_source: None, - value_stats_cols: None, - external_path: None, - first_row_id, - write_cols: None, - column_max_sequence_numbers: None, - } - } - - /// Build a predicate-free reader (read_type = `id` + `_ROW_ID`) over a split - /// containing `files` (each `(name, ids, first_row_id)`), written as Mosaic - /// data files in the same bucket. The returned active-file list covers every - /// file (all files active). - async fn build_reader_and_split( - table_path: &str, - files: &[(&str, Vec, i64)], - ) -> (DataFileReader, DataSplit, Vec) { - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let bucket_path = format!("{table_path}/bucket-0"); - let mut metas = Vec::new(); - let mut active_files = Vec::new(); - for (name, ids, first_row_id) in files { - let data = write_mosaic(&id_batch(ids.clone())); - file_io - .new_output(&format!("{bucket_path}/{name}")) - .unwrap() - .write(data.clone()) - .await - .unwrap(); - metas.push(data_file( - name, - data.len() as i64, - ids.len() as i64, - Some(*first_row_id), - )); - active_files.push(BucketActiveFile { - file_name: name.to_string(), - row_count: ids.len() as i64, - }); - } - let split = DataSplitBuilder::new() - .with_snapshot(1) - .with_partition(BinaryRow::new(0)) - .with_bucket(0) - .with_bucket_path(bucket_path) - .with_total_buckets(1) - .with_data_files(metas) - .build() - .unwrap(); - let reader = DataFileReader::new( - file_io.clone(), - SchemaManager::new(file_io, table_path.to_string()), - 1, - vec![id_field()], - vec![id_field(), row_id_field()], - Vec::new(), - ); - (reader, split, active_files) - } - - /// `id > threshold`, with `file_fields` = `[id]` so the leaf index resolves. - fn residual_id_gt(threshold: i32) -> FilePredicates { - let pred = PredicateBuilder::new(&[id_field()]) - .greater_than("id", Datum::Int(threshold)) - .unwrap(); - FilePredicates { - predicates: vec![pred], - row_filter_factory: None, - file_fields: vec![id_field()], - } - } - - fn sorted(t: &roaring::RoaringTreemap) -> Vec { - t.iter().collect() - } - - #[tokio::test] - async fn test_residual_selects_matching_positions() { - // ids [1,2,3,4,5] at first_row_id 0; id > 2 -> ids 3,4,5 -> positions 2,3,4. - let (reader, split, active) = build_reader_and_split( - "memory:/rpf_basic", - &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], - ) - .await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), None) - .await - .unwrap(); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]); - } - - #[tokio::test] - async fn test_residual_only_evaluates_the_rows_the_plan_allows() { - // ids [1,2,3,4,5]; the plan allows positions 3-4 only. `id > 2` matches 2,3,4 - // over the whole file, so a result of 3,4 is the plan's restriction taking - // effect *before* evaluation: position 2 is never seen. - // - // This also cannot pass under a full read. The scan walks the selection in - // step with the emitted rows, so a read that emitted all five would run the - // selection dry and fail loudly rather than return a filtered answer. - let (reader, split, active) = build_reader_and_split( - "memory:/rpf_plan_ranges", - &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], - ) - .await; - let allowed = HashMap::from([("part-0.mosaic".to_string(), vec![RowRange::new(3, 4)])]); - let map = residual_positions_by_file( - &reader, - &split, - &active, - &residual_id_gt(2), - Some(&allowed), - ) - .await - .unwrap(); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]); - } - - #[tokio::test] - async fn test_residual_does_not_read_a_file_the_plan_excludes() { - // An EMPTY range list is how a plan says "no rows of this file": it is - // registered empty and never opened. Absence means the opposite -- the plan - // narrowed nothing there -- so the residual reads the whole file. - let (reader, split, active) = build_reader_and_split( - "memory:/rpf_plan_excludes", - &[("part-0.mosaic", vec![1, 2, 3], 0)], - ) - .await; - - let excluded = HashMap::from([("part-0.mosaic".to_string(), Vec::new())]); - let map = residual_positions_by_file( - &reader, - &split, - &active, - &residual_id_gt(0), - Some(&excluded), - ) - .await - .unwrap(); - assert!(map.contains_key("part-0.mosaic")); - assert!(sorted(&map["part-0.mosaic"]).is_empty()); - - let unrestricted = HashMap::new(); - let map = residual_positions_by_file( - &reader, - &split, - &active, - &residual_id_gt(0), - Some(&unrestricted), - ) - .await - .unwrap(); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); - } - - #[tokio::test] - async fn test_residual_matches_none_yields_empty_entry() { - // id > 100 matches nothing; the file still gets a (present, empty) entry. - let (reader, split, active) = - build_reader_and_split("memory:/rpf_none", &[("part-0.mosaic", vec![1, 2, 3], 0)]) - .await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(100), None) - .await - .unwrap(); - assert!(map.contains_key("part-0.mosaic")); - assert!(map["part-0.mosaic"].is_empty()); - } - - #[tokio::test] - async fn test_residual_matches_all_yields_full_set() { - let (reader, split, active) = - build_reader_and_split("memory:/rpf_all", &[("part-0.mosaic", vec![1, 2, 3], 0)]).await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(0), None) - .await - .unwrap(); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); - } - - #[tokio::test] - async fn test_residual_positions_are_file_local_across_files() { - // Two files with distinct first_row_id; positions must be 0-based within - // each file, not global. id > 3 keeps ids 4,5 in both -> positions {3,4}. - let (reader, split, active) = build_reader_and_split( - "memory:/rpf_multi", - &[ - ("part-0.mosaic", vec![1, 2, 3, 4, 5], 0), - ("part-1.mosaic", vec![1, 2, 3, 4, 5], 100), - ], - ) - .await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(3), None) - .await - .unwrap(); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]); - assert_eq!(sorted(&map["part-1.mosaic"]), vec![3, 4]); - } - - #[tokio::test] - async fn test_non_active_files_are_skipped() { - // Two files in the split, but only `part-0.mosaic` is active. The bucket - // search never recalls from `part-1.mosaic` (level-0 / non-active), so it - // must not appear in the residual map — and even though it lacks a - // `first_row_id`, the query still succeeds because non-active files are - // skipped before the guard. - let (reader, split, mut active) = build_reader_and_split( - "memory:/rpf_nonactive", - &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], - ) - .await; - // Append a non-active file (missing first_row_id) directly to the split's - // data files, but leave it out of the active list. - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let bucket_path = "memory:/rpf_nonactive/bucket-0"; - let data = write_mosaic(&id_batch(vec![9, 9, 9])); - file_io - .new_output(&format!("{bucket_path}/part-1.mosaic")) - .unwrap() - .write(data.clone()) - .await - .unwrap(); - let mut metas = split.data_files().to_vec(); - metas.push(data_file("part-1.mosaic", data.len() as i64, 3, None)); - // `active` already lists only part-0.mosaic; keep it that way. - let _ = &mut active; - let split = DataSplitBuilder::new() - .with_snapshot(1) - .with_partition(BinaryRow::new(0)) - .with_bucket(0) - .with_bucket_path(bucket_path.to_string()) - .with_total_buckets(1) - .with_data_files(metas) - .build() - .unwrap(); - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), None) - .await - .unwrap(); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]); - assert!( - !map.contains_key("part-1.mosaic"), - "non-active file must be skipped" - ); - } - - #[tokio::test] - async fn test_missing_first_row_id_recovers_local_positions() { - // Real primary-key data files carry no `first_row_id`. Positions are - // recovered from each row's ordinal in the scan, so the residual still - // works: ids [1,2,3] with id > 0 -> all match -> local positions [0,1,2]. - let (reader, split, active) = build_reader_and_split_no_first_row_id().await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(0), None) - .await - .expect("missing first_row_id must not fail the residual read"); - assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); - } - - async fn build_reader_and_split_no_first_row_id( - ) -> (DataFileReader, DataSplit, Vec) { - let table_path = "memory:/rpf_nofrid"; - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let bucket_path = format!("{table_path}/bucket-0"); - let data = write_mosaic(&id_batch(vec![1, 2, 3])); - file_io - .new_output(&format!("{bucket_path}/part-0.mosaic")) - .unwrap() - .write(data.clone()) - .await - .unwrap(); - let split = DataSplitBuilder::new() - .with_snapshot(1) - .with_partition(BinaryRow::new(0)) - .with_bucket(0) - .with_bucket_path(bucket_path) - .with_total_buckets(1) - .with_data_files(vec![data_file("part-0.mosaic", data.len() as i64, 3, None)]) - .build() - .unwrap(); - let reader = DataFileReader::new( - file_io.clone(), - SchemaManager::new(file_io, table_path.to_string()), - 1, - vec![id_field()], - vec![id_field(), row_id_field()], - Vec::new(), - ); - // The lone file is active and carries no first_row_id, exercising the - // ordinal-based position recovery. - let active = vec![BucketActiveFile { - file_name: "part-0.mosaic".to_string(), - row_count: 3, - }]; - (reader, split, active) - } - - // ---- combining the plan's positional restriction with the residual ---- - - fn allow_list(entries: &[(&str, &[u64])]) -> HashMap { - entries - .iter() - .map(|(file, positions)| ((*file).to_string(), positions.iter().copied().collect())) - .collect() - } - - /// The plan side carries ranges, so its fixtures are built from the positions - /// each file allows and coalesced the way the planner normalizes them. - fn range_allow_list(entries: &[(&str, &[u64])]) -> HashMap> { - entries - .iter() - .map(|(file, positions)| { - let ranges = positions - .iter() - .map(|p| RowRange::new(*p as i64, *p as i64)) - .collect(); - ((*file).to_string(), merge_row_ranges(ranges)) - }) - .collect() - } - - /// The positions a merged selection allows, expanded for readable assertions. - /// Test-only: the production path never expands a range. - fn listed(map: &FileRowSelections, file: &str) -> Vec { - match map.get(file) { - None => Vec::new(), - Some(FileRowSelection::Positions(positions)) => positions.iter().collect(), - Some(FileRowSelection::Ranges(ranges)) => ranges - .iter() - .flat_map(|range| (range.from() as u64)..=(range.to() as u64)) - .collect(), - } - } - - #[test] - fn no_restriction_on_either_side_stays_unrestricted() { - assert!(intersect_row_allow_lists(None, None, 1).unwrap().is_none()); - } - - #[test] - fn one_side_alone_passes_through() { - let physical = vec![range_allow_list(&[("d0", &[1, 2])])]; - let only_physical = intersect_row_allow_lists(Some(&physical), None, 1) - .unwrap() - .expect("a plan restriction survives on its own"); - assert_eq!(listed(&only_physical[0], "d0"), vec![1, 2]); - // Still intervals. Expanding them here is the unbounded step the plan side - // must never take, and the positions above cannot tell the two apart. - assert!( - matches!(only_physical[0]["d0"], FileRowSelection::Ranges(_)), - "the plan's ranges must reach the search as ranges" - ); - - let residual = vec![allow_list(&[("d0", &[3])])]; - let only_residual = intersect_row_allow_lists(None, Some(residual), 1) - .unwrap() - .expect("a residual survives on its own"); - assert_eq!(listed(&only_residual[0], "d0"), vec![3]); - } - - #[test] - fn both_sides_intersect_and_the_residual_stays_fail_closed() { - // `d0`: both restrict it, so only the shared positions survive. `d1`: the - // residual says nothing about it. The residual registers EVERY file the - // search can read from, so its silence is "no rows" -- the plan's ranges - // must not resurrect the file, and neither may its absence make it - // unrestricted. - let physical = vec![range_allow_list(&[("d0", &[1, 2, 3]), ("d1", &[0, 1])])]; - let residual = vec![allow_list(&[("d0", &[2, 3, 4])])]; - let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) - .unwrap() - .expect("both sides restrict"); - assert_eq!(listed(&combined[0], "d0"), vec![2, 3]); - assert!( - combined[0]["d1"].is_excluded(), - "a file the residual omits must stay excluded" - ); - } - - #[test] - fn a_file_neither_side_restricts_stays_absent() { - // Absence is how "every row" is spelled. A merged map must not invent an - // entry for a file no one narrowed, or the ANN backend takes the filtered - // path for a query that filters nothing. - let physical = vec![range_allow_list(&[("d0", &[1])])]; - let combined = intersect_row_allow_lists(Some(&physical), None, 1) - .unwrap() - .expect("the plan restricts d0"); - assert!(!combined[0].contains_key("d1")); - - let residual = vec![allow_list(&[("d0", &[1])])]; - let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) - .unwrap() - .expect("both restrict d0"); - assert!(!combined[0].contains_key("d1")); - assert!(!combined[0].contains_key("d2")); - } - - #[test] - fn a_plan_that_restricts_nothing_produces_an_empty_selection_map() { - // The no-pre-filter split: the plan carries a map with no entries at all, - // and that must survive the merge as an empty map (which the ANN layer reads - // as "nothing to mask"), not become a per-file all-permitting mask. - let physical = vec![HashMap::new()]; - let combined = intersect_row_allow_lists(Some(&physical), None, 1) - .unwrap() - .expect("a split-driven plan is always Some"); - assert!(combined[0].is_empty()); - } - - /// The batch terminals here are handed exactly one query, so a result vector of any - /// other length means the batch ran the wrong number of searches. The - /// `debug_assert_eq!` this replaced was compiled out of release builds, where an - /// empty vector panicked on `remove(0)` and a longer one silently returned another - /// query's result. - #[test] - fn take_only_result_rejects_bad_batch_arity() { - assert_eq!(take_only_result(vec![7], "test").unwrap(), 7); - assert!(take_only_result::(Vec::new(), "test").is_err()); - assert!(take_only_result(vec![1, 2], "test").is_err()); - } - - #[test] - fn rejects_allow_lists_that_do_not_cover_every_split() { - let physical = vec![range_allow_list(&[("d0", &[1])])]; - let error = intersect_row_allow_lists(Some(&physical), None, 2) - .map(|_| ()) - .expect_err("an allow-list per split is what makes the index meaningful"); - assert!(error.to_string().contains("for 2 splits"), "{error}"); - - let residual = vec![allow_list(&[("d0", &[1])])]; - let error = intersect_row_allow_lists(Some(&physical), Some(residual), 2) - .map(|_| ()) - .expect_err("the residual must cover every split too"); - assert!(error.to_string().contains("for 2 splits"), "{error}"); - } -} +mod tests; diff --git a/crates/paimon/src/table/vector_search_builder/tests.rs b/crates/paimon/src/table/vector_search_builder/tests.rs new file mode 100644 index 000000000..3638254ee --- /dev/null +++ b/crates/paimon/src/table/vector_search_builder/tests.rs @@ -0,0 +1,678 @@ +// 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::catalog::Identifier; +use crate::io::FileIOBuilder; +use crate::spec::{ + ArrayType, DataType, Datum, FloatType, IntType, Predicate, PredicateBuilder, Schema, + TableSchema, ROW_ID_FIELD_NAME, +}; +use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; +use crate::table::vector_search_common::resolve_materialize_read_type; +use crate::table::vector_search_test_utils::{ + id_gt_filter, pk_vector_table, pk_vector_table_with_extra_column, +}; +use crate::table::Table; +use crate::vindex::IVF_FLAT_IDENTIFIER; +use futures::TryStreamExt; +use std::collections::HashMap; + +#[tokio::test] +async fn test_execute_fails_closed_when_query_auth_enabled() { + let table = crate::table::query_auth_table(); + let err = table + .new_vector_search_builder() + .execute() + .await + .unwrap_err(); + assert!( + matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), + "vector search must fail closed for a query-auth table" + ); +} + +#[tokio::test] +async fn pk_branch_disabled_falls_through_to_de_path() { + // No pk-vector.index.columns: behaves exactly as the DE path. With no + // snapshot the DE path returns an empty result; the PK branch must not + // intercept it. + let table = pk_vector_table(&[]); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await + .unwrap(); + assert!(result.is_empty()); +} + +#[tokio::test] +async fn pk_branch_execute_returns_physical_positions() { + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ("fields.embedding.dimension", "4"), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .execute() + .await + .unwrap(); + assert!(result.positions().unwrap().is_empty()); + assert_eq!(result.snapshot_id(), None); + assert!(result.row_ids().is_err()); +} + +#[tokio::test] +async fn pk_branch_other_column_falls_through_to_de_path() { + // pk-vector index configured for "embedding", but the query targets a + // different column -> the PK branch must not intercept; DE path (no + // snapshot) yields empty. Discriminator: the PK column carries a + // DELIBERATELY INVALID distance metric, which the PK branch parses eagerly + // (`VectorSearchMetric::parse`) and would fail on. So a regression that + // dropped the `pk_col == vector_column` guard and ran the PK branch for + // "other" would surface as Err here, not Ok(empty) -- the assertion + // therefore proves the DE path ran, not merely that the result is empty. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ( + "fields.embedding.pk-vector.distance.metric", + "not-a-real-metric", + ), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("other") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await + .unwrap(); + assert!(result.is_empty()); +} + +#[tokio::test] +async fn pk_branch_multi_column_config_does_not_break_unrelated_de_query() { + // A malformed multi-column PK-vector config ("a,b") must not abort an + // unrelated DE vector query. The query targets a column NOT among the + // configured PK-vector columns, so membership resolution short-circuits + // before the exactly-one-column rule fires -- the query falls through to + // the DE path (no snapshot -> empty) instead of surfacing the "must name + // exactly one column" error. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "a,b"), + ("fields.a.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.a.pk-vector.distance.metric", "l2"), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("other") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await; + match result { + Ok(search) => assert!(search.is_empty()), + Err(err) => { + panic!("unrelated DE query must not error on a malformed multi-column PK config: {err}") + } + } +} + +#[tokio::test] +async fn result_read_filter_without_deletion_vectors_fails_loud() { + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let filter = id_gt_filter(&table, 2); + let err = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .with_filter(filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .map(|_| ()) + .expect_err("read filter without deletion vectors must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("deletion vectors without merge-on-read")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn result_read_filter_with_merge_on_read_fails_loud() { + // Deletion vectors enabled BUT merge-on-read on: still rejected, because a + // merge-on-read scan can surface stale key versions that a physical-row + // filter cannot reconcile. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ("deletion-vectors.enabled", "true"), + ("deletion-vectors.merge-on-read", "true"), + ]); + let filter = id_gt_filter(&table, 2); + let err = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .with_filter(filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .map(|_| ()) + .expect_err("merge-on-read filter must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("deletion vectors without merge-on-read")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn result_read_filter_with_deletion_vectors_passes_guard() { + // Deletion vectors enabled, merge-on-read off (default): the residual guard + // passes. With no snapshot the plan is empty, so the (guarded) filter path + // simply yields an empty stream rather than erroring — proving the guard + // admits a legal filtered query. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ("deletion-vectors.enabled", "true"), + // Pin the index dimension so the query vector below matches it; the + // up-front dimension guard runs before this test's residual guard. + ("fields.embedding.dimension", "4"), + ]); + let filter = id_gt_filter(&table, 2); + let mut stream = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .with_filter(filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .expect("guarded filter query must be admitted"); + assert!(stream.try_next().await.unwrap().is_none()); +} + +/// A partition-only `with_filter` needs no per-row residual (partition pruning +/// happens in scan planning), so the deletion-vector pre-filter guard must NOT +/// reject it even when deletion vectors are off. Mirrors Java, where a +/// partition-only filter leaves `this.filter == null` and the scan guard is +/// skipped. Regression test for the guard keying on the whole filter rather +/// than its data conjuncts. +#[tokio::test] +async fn result_read_partition_only_filter_without_deletion_vectors_passes_guard() { + use crate::spec::VarCharType; + + // Partitioned PK-vector table from old metadata, deletion vectors OFF. + let schema = Schema::builder() + .column("dt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .partition_keys(["dt"]) + .primary_key(["id"]) + .option("bucket", "1") + .build() + .unwrap(); + let table_schema = TableSchema::new(0, &schema).copy_with_options(HashMap::from([ + ( + "pk-vector.index.columns".to_string(), + "embedding".to_string(), + ), + ( + "fields.embedding.pk-vector.index.type".to_string(), + IVF_FLAT_IDENTIFIER.to_string(), + ), + ( + "fields.embedding.pk-vector.distance.metric".to_string(), + "l2".to_string(), + ), + ("fields.embedding.dimension".to_string(), "4".to_string()), + ])); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "pk_vector_partitioned"), + "memory:/pk_vector_partitioned".to_string(), + table_schema, + None, + ); + + // Partition-only `dt = 'a'`: no data residual, so the guard admits it and + // (with no snapshot) the query yields an empty stream instead of the + // deletion-vector error. + let filter = PredicateBuilder::new(table.schema().fields()) + .equal("dt", Datum::String("a".to_string())) + .unwrap(); + let mut stream = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .with_filter(filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .expect("partition-only filter must be admitted without deletion vectors"); + assert!(stream.try_next().await.unwrap().is_none()); + + // But a DATA conjunct (`id > 2`) on the same non-DV table must still fail + // loud — the guard now keys on data predicates, not the whole filter. + let data_filter = id_gt_filter(&table, 2); + let err = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .with_filter(data_filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .map(|_| ()) + .expect_err("data filter without deletion vectors must still fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("deletion vectors without merge-on-read")), + "unexpected error: {err:?}" + ); + + // `AND(partition, data)` still has a data conjunct after the split, so it + // must fail loud on the non-DV table just like the data-only filter. + let pb = PredicateBuilder::new(table.schema().fields()); + let and_filter = Predicate::and(vec![ + pb.equal("dt", Datum::String("a".to_string())).unwrap(), + pb.greater_than("id", Datum::Int(2)).unwrap(), + ]); + let err = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .with_filter(and_filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .map(|_| ()) + .expect_err("AND(partition, data) without deletion vectors must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("deletion vectors without merge-on-read")), + "unexpected error: {err:?}" + ); + + // A mixed `OR(partition, data)` conjunct is not partition-only, so it stays + // whole as a data predicate and must also fail loud without deletion vectors. + let or_filter = Predicate::or(vec![ + pb.equal("dt", Datum::String("a".to_string())).unwrap(), + pb.greater_than("id", Datum::Int(2)).unwrap(), + ]); + let err = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .with_filter(or_filter) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .map(|_| ()) + .expect_err("mixed OR(partition, data) without deletion vectors must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("deletion vectors without merge-on-read")), + "unexpected error: {err:?}" + ); +} + +#[tokio::test] +async fn result_read_de_table_empty_snapshot_yields_empty_stream() { + // No pk-vector index configured and no snapshot: result_read routes to the + // data-evolution path, whose search finds nothing and returns an empty + // stream (not an error). + let table = pk_vector_table(&[]); + let mut stream = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .expect("DE read over an empty table must succeed with no rows"); + let mut rows = 0usize; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows(); + } + assert_eq!(rows, 0, "empty DE table must yield no rows"); +} + +#[tokio::test] +async fn result_read_unknown_column_fails_loud() { + // pk-vector index configured for "embedding", but the query targets a + // column that does not exist. The read path must fail loud rather than + // fall through to the data-evolution path and return an empty stream (a + // typo must not look like a normal empty read through the C API). + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let err = match async { + table + .new_vector_search_builder() + .with_vector_column("other") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + { + Ok(_) => panic!("unknown vector column must fail loud on result_read"), + Err(e) => e, + }; + assert!( + matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("does not exist")), + "expected a does-not-exist error, got: {err}" + ); +} + +#[tokio::test] +async fn result_read_scalar_column_fails_loud() { + // A scalar (non-vector) column targeted by a vector read must fail loud, + // not return an empty data-evolution stream. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let err = match async { + table + .new_vector_search_builder() + .with_vector_column("id") // scalar Int column + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + { + Ok(_) => panic!("scalar vector column must fail loud on result_read"), + Err(e) => e, + }; + assert!( + matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("must be a FLOAT vector column")), + "expected a not-a-vector-column error, got: {err}" + ); +} + +#[tokio::test] +async fn result_read_non_float_vector_column_fails_loud() { + // An ARRAY column is not a searchable vector column (the index/search + // operates on FLOAT elements). It must fail loud rather than fall through + // to the DE path and return an empty stream. + use crate::spec::{ArrayType, IntType, Schema, TableSchema}; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), + ) + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "de_non_float_vector"), + "memory:/de_non_float_vector".to_string(), + TableSchema::new(0, &schema), + None, + ); + let err = match async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + { + Ok(_) => panic!("ARRAY vector column must fail loud on result_read"), + Err(e) => e, + }; + assert!( + matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("must be a FLOAT vector column")), + "expected a FLOAT-vector-column error, got: {err}" + ); +} + +#[tokio::test] +async fn result_read_empty_plan_reserved_projection_fails_loud() { + // Empty plan (no snapshot) must still fail loud on a reserved-name + // projection: projection validity does not depend on whether the search + // matched any rows. A regression that resolved the projection only after + // the `candidates.is_empty()` early return would yield an empty stream here + // instead of an error. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + // Pin the index dimension so the query vector below matches it; the + // up-front dimension guard runs before this test's reserved-projection + // guard, so a mismatched query would mask the error under test. + ("fields.embedding.dimension", "4"), + ]); + for reserved in [ + ROW_ID_FIELD_NAME, + PKEY_VECTOR_POSITION_COLUMN, + SEARCH_SCORE_COLUMN, + ] { + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5); + let err = async { + builder + .execute() + .await? + .new_read_builder() + .with_projection(&["id", reserved]) + .read() + .await + } + .await + .map(|_| ()) + .expect_err("empty plan + reserved projection must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "unexpected error for {reserved}: {err:?}" + ); + } +} + +#[tokio::test] +async fn result_read_empty_plan_lumina_array_float_is_admitted() { + // A Lumina PK-vector `ARRAY` column is a valid configuration, but + // batch query dimension validation routed every `ARRAY` column + // through the vindex resolver, which rejects `lumina` as an unsupported + // index type before planning — failing even an empty table. The + // dimension must be resolved per the configured backend, so a + // well-formed Lumina query is admitted and (with no snapshot) yields an + // empty stream rather than an "Unsupported vindex index type" error. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ( + "fields.embedding.pk-vector.index.type", + crate::lumina::LUMINA_IDENTIFIER, + ), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ("lumina.index.dimension", "4"), + ]); + let mut stream = async { + table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .execute() + .await? + .new_read_builder() + .read() + .await + } + .await + .expect("Lumina ARRAY query must be admitted, not rejected as unsupported vindex"); + assert!(stream.try_next().await.unwrap().is_none()); +} + +#[tokio::test] +async fn result_read_projection_reserved_name_fails_loud() { + // Projecting a reserved metadata / row-id column must fail loud. The guard + // lives in `resolve_materialize_read_type`, which `SearchResultReadBuilder::read` invokes + // before the empty-plan early return; assert on the resolver directly here. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + for reserved in [ + ROW_ID_FIELD_NAME, + PKEY_VECTOR_POSITION_COLUMN, + SEARCH_SCORE_COLUMN, + ] { + let err = + resolve_materialize_read_type(&table, Some(&["id".to_string(), reserved.to_string()])) + .expect_err("reserved projection must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "unexpected error for {reserved}: {err:?}" + ); + } +} + +#[test] +fn resolve_materialize_read_type_default_is_all_user_columns() { + // No with_projection -> every user table column (id + embedding). + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let fields = resolve_materialize_read_type(&table, None).unwrap(); + let names: Vec<&str> = fields.iter().map(|f| f.name()).collect(); + assert_eq!(names, vec!["id", "embedding"]); +} + +#[test] +fn resolve_materialize_read_type_default_rejects_reserved_user_column() { + // The default (all-columns) projection must reject a user column whose + // name collides with an injected metadata column, not only columns named + // in an explicit projection. Otherwise it silently passes on an empty + // result and collides with the metadata columns the read attaches. + let table = pk_vector_table_with_extra_column(SEARCH_SCORE_COLUMN); + let err = resolve_materialize_read_type(&table, None).unwrap_err(); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "single-query default projection must reject reserved user column, got: {err:?}" + ); +} + +#[test] +fn resolve_materialize_read_type_projection_selects_named_columns() { + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let fields = resolve_materialize_read_type(&table, Some(&["id".to_string()])).unwrap(); + let names: Vec<&str> = fields.iter().map(|f| f.name()).collect(); + assert_eq!(names, vec!["id"]); +} diff --git a/crates/paimon/src/table/vector_search_common.rs b/crates/paimon/src/table/vector_search_common.rs new file mode 100644 index 000000000..61535a8ad --- /dev/null +++ b/crates/paimon/src/table/vector_search_common.rs @@ -0,0 +1,401 @@ +// 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. + +//! Shared vector-search option validation, index I/O helpers, and result ordering. + +use crate::lumina::is_lumina_index_type; +use crate::spec::{CoreOptions, DataField, ROW_ID_FIELD_NAME}; +use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; +use crate::table::read_builder::resolve_projected_fields; +use crate::table::Table; +use crate::vindex::is_vindex_index_type; +use crate::vindex::range_reader::RangeIoStats; +use arrow_array::{Int64Array, RecordBatch}; +use arrow_select::interleave::interleave_record_batch; +use std::collections::HashMap; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum VectorIndexBackend { + Lumina, + Vindex, +} + +impl VectorIndexBackend { + pub(super) fn from_index_type(index_type: &str) -> Option { + if is_lumina_index_type(index_type) { + Some(Self::Lumina) + } else if is_vindex_index_type(index_type) { + Some(Self::Vindex) + } else { + None + } + } + + pub(super) fn error_name(self) -> &'static str { + match self { + Self::Lumina => "Lumina", + Self::Vindex => "vindex", + } + } +} + +pub(super) fn current_tokio_runtime_handle() -> crate::Result { + tokio::runtime::Handle::try_current().map_err(|error| crate::Error::UnexpectedError { + message: "Vector index range reader requires a Tokio runtime".to_string(), + source: Some(Box::new(error)), + }) +} + +fn vindex_index_parallelism(entry_count: usize, max_concurrency: usize) -> usize { + entry_count.min(max_concurrency).max(1) +} + +pub(super) fn log_vindex_range_io_stats(file: &str, query_count: usize, stats: &RangeIoStats) { + let stats = stats.snapshot(); + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_range_io file={} nq={} logical_ranges={} requested_bytes={} file_read_calls={} returned_bytes={} read_ahead_hits={} io_wait_sum_ms={:.3} range_permit_wait_sum_ms={:.3} peak_in_flight_reads={} read_many_merged_ranges={} read_many_chunks={} read_many_chunk_size_sum={} read_many_chunk_size_min={} read_many_chunk_size_max={}", + file, + query_count, + stats.logical_ranges, + stats.requested_bytes, + stats.file_read_calls, + stats.returned_bytes, + stats.read_ahead_hits, + stats.io_wait_nanos as f64 / 1_000_000.0, + stats.range_permit_wait_nanos as f64 / 1_000_000.0, + stats.peak_in_flight_reads, + stats.read_many_merged_ranges, + stats.read_many_chunks, + stats.read_many_chunk_size_sum, + stats.read_many_chunk_size_min, + stats.read_many_chunk_size_max, + ); +} + +pub(super) fn vindex_concurrency_limits( + core_options: &CoreOptions<'_>, + entry_count: usize, + max_concurrency: usize, +) -> crate::Result<(usize, usize)> { + Ok(( + vindex_index_parallelism(entry_count, max_concurrency), + core_options.global_index_vindex_read_thread_num()?, + )) +} + +/// Unwrap a single-query result from a batch entry point that must return exactly +/// one element per input query. +/// +/// The batch terminals below are handed one query, so their result vector holds +/// exactly one entry. A `debug_assert_eq!(len, 1)` followed by `remove(0)` checked +/// that only in debug builds, where a release build would instead panic on an index +/// out of bounds for an empty vector -- or SILENTLY return the first of several, +/// pairing the caller's single query with another query's result. A length that is +/// wrong means the batch ran the wrong number of searches, which is a programming +/// error in this crate rather than bad input, so it is reported as one. +pub(super) fn take_only_result(results: Vec, operation: &str) -> crate::Result { + let mut results = results.into_iter(); + let result = results + .next() + .ok_or_else(|| crate::Error::UnexpectedError { + message: format!("{operation} returned no result for one query"), + source: None, + })?; + if results.next().is_some() { + return Err(crate::Error::UnexpectedError { + message: format!("{operation} returned more than one result for one query"), + source: None, + }); + } + Ok(result) +} + +/// Resolve the projected fields for the materialization read-type. Default +/// (no projection set) is all user table fields; otherwise the requested +/// names resolved via `resolve_projected_fields`. Rejects reserved metadata +/// names and `_ROW_ID` so a user cannot request a hidden column. +pub(super) fn resolve_materialize_read_type( + table: &Table, + projection: Option<&[String]>, +) -> crate::Result> { + let fields = match projection { + None => table.schema().fields().to_vec(), + Some(names) => { + for name in names { + if is_reserved_read_column(name) { + return Err(crate::Error::DataInvalid { + message: format!( + "vector search read projection must not request reserved column '{name}'" + ), + source: None, + }); + } + } + resolve_projected_fields( + table.identifier().full_name(), + table.schema().fields(), + names, + true, + )? + } + }; + // The default projection returns every user column, so a user column + // whose name collides with an injected metadata column must be rejected + // on the resolved field list too — not only when explicitly requested. + ensure_no_reserved_read_columns(&fields)?; + Ok(fields) +} + +/// Names a read injects as metadata columns — `__paimon_search_score`, +/// `_PKEY_VECTOR_POSITION`, and `_ROW_ID` — that a materialized read type must +/// not reuse for a user column. +fn is_reserved_read_column(name: &str) -> bool { + name == PKEY_VECTOR_POSITION_COLUMN || name == SEARCH_SCORE_COLUMN || name == ROW_ID_FIELD_NAME +} + +/// Reject a materialized read type whose resolved fields contain a reserved +/// metadata column name. Applied to the RESOLVED field list so the default +/// (all user columns) projection is covered, not only an explicit one. +pub(crate) fn ensure_no_reserved_read_columns(fields: &[DataField]) -> crate::Result<()> { + for field in fields { + if is_reserved_read_column(field.name()) { + return Err(crate::Error::DataInvalid { + message: format!( + "search read must not include reserved column '{}'", + field.name() + ), + source: None, + }); + } + } + Ok(()) +} + +/// One materialized row tagged with its best-first `rank` and its `(batch_index, +/// row_index)` location in the retained materialization batches. +pub(crate) struct RankedRow { + rank: usize, + batch_index: usize, + row_index: usize, +} + +/// For each row in a materialized batch, look up its best-first rank via the +/// `(partition bytes, bucket, file, position)` key and record its location. The +/// `_PKEY_VECTOR_POSITION` column supplies the physical position; every row must +/// map to a candidate rank (the batch came from that candidate's file), so a miss +/// fails loud rather than silently dropping a row. +#[allow(clippy::too_many_arguments)] +pub(crate) fn collect_ranked_rows( + batch: &RecordBatch, + batch_index: usize, + partition_bytes: &[u8], + bucket: i32, + file_name: &str, + rank_of: &HashMap<(Vec, i32, String, i64), usize>, + out: &mut Vec, +) -> crate::Result<()> { + let position_idx = batch + .schema() + .index_of(PKEY_VECTOR_POSITION_COLUMN) + .map_err(|_| crate::Error::DataInvalid { + message: format!("materialized batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), + source: None, + })?; + let positions = batch + .column(position_idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is not Int64"), + source: None, + })?; + for row_index in 0..batch.num_rows() { + let position = positions.value(row_index); + let key = ( + partition_bytes.to_vec(), + bucket, + file_name.to_string(), + position, + ); + let rank = *rank_of.get(&key).ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "materialized row (file {file_name}, position {position}) has no matching search candidate" + ), + source: None, + })?; + out.push(RankedRow { + rank, + batch_index, + row_index, + }); + } + Ok(()) +} + +/// Reorder the materialized rows into best-first order and drop the internal +/// `_PKEY_VECTOR_POSITION` column, yielding a single output batch (empty input +/// yields no batches). The projected user columns and `__paimon_search_score` are +/// retained. +pub(crate) fn reorder_and_strip_position( + batches: &[RecordBatch], + mut ranked: Vec, +) -> crate::Result> { + if ranked.is_empty() { + return Ok(Vec::new()); + } + ranked.sort_by_key(|r| r.rank); + let indices: Vec<(usize, usize)> = ranked + .iter() + .map(|r| (r.batch_index, r.row_index)) + .collect(); + let refs: Vec<&RecordBatch> = batches.iter().collect(); + let reordered = + interleave_record_batch(&refs, &indices).map_err(|e| crate::Error::DataInvalid { + message: format!("failed to reorder vector search read rows: {e}"), + source: None, + })?; + + // Drop the internal position column; keep every other column (projected user + // columns + __paimon_search_score) in order. + let position_idx = reordered + .schema() + .index_of(PKEY_VECTOR_POSITION_COLUMN) + .map_err(|_| crate::Error::DataInvalid { + message: format!("reordered batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), + source: None, + })?; + let keep: Vec = (0..reordered.num_columns()) + .filter(|i| *i != position_idx) + .collect(); + let projected = reordered + .project(&keep) + .map_err(|e| crate::Error::DataInvalid { + message: format!("failed to drop position column: {e}"), + source: None, + })?; + Ok(vec![projected]) +} + +pub(super) fn indexed_search_limit(limit: usize, refine_factor: usize) -> crate::Result { + if refine_factor == 0 { + return Ok(limit); + } + let search_limit = + limit + .checked_mul(refine_factor) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: format!( + "Vector search limit overflow: limit={limit}, refine factor={refine_factor}" + ), + })?; + if search_limit > i32::MAX as usize { + return Err(crate::Error::ConfigInvalid { + message: format!( + "Vector search limit overflow: limit={limit}, refine factor={refine_factor}" + ), + }); + } + Ok(search_limit) +} + +pub(super) fn normalize_metric(metric: &str) -> String { + metric.to_ascii_lowercase().replace('-', "_") +} + +fn indexed_type_prefixes(field_name: &str, index_type: &str) -> Vec { + let mut prefixes = Vec::new(); + add_refine_prefixes(&mut prefixes, &format!("fields.{field_name}."), index_type); + add_refine_prefixes(&mut prefixes, "", index_type); + prefixes +} + +fn add_refine_prefixes(prefixes: &mut Vec, base: &str, index_type: &str) { + if !index_type.is_empty() { + prefixes.push(format!("{base}{index_type}.")); + let normalized = normalize_metric(index_type); + if normalized != index_type { + prefixes.push(format!("{base}{normalized}.")); + } + if normalized.starts_with("ivf") { + prefixes.push(format!("{base}ivf.")); + } + } + prefixes.push(base.to_string()); +} + +pub(super) fn configured_refine_factor( + search_options: &HashMap, + table_options: &HashMap, + field_name: &str, + index_type: &str, +) -> crate::Result { + if let Some(value) = + configured_refine_factor_from_options(search_options, field_name, index_type) + { + return parse_refine_factor(&value); + } + if let Some(value) = + configured_refine_factor_from_options(table_options, field_name, index_type) + { + return parse_refine_factor(&value); + } + Ok(0) +} + +fn configured_refine_factor_from_options( + options: &HashMap, + field_name: &str, + index_type: &str, +) -> Option { + for prefix in indexed_type_prefixes(field_name, index_type) { + for suffix in [ + "refine_factor", + "refine-factor", + "rerank_factor", + "rerank-factor", + ] { + if let Some(value) = options.get(&(prefix.clone() + suffix)) { + return Some(value.trim().to_string()); + } + } + } + None +} + +fn parse_refine_factor(value: &str) -> crate::Result { + let factor = value + .parse::() + .map_err(|_| crate::Error::ConfigInvalid { + message: format!("Invalid vector refine factor: {value}. Must be an integer."), + })?; + if factor == 0 { + return Err(crate::Error::ConfigInvalid { + message: format!("Vector refine factor must be positive, got: {value}"), + }); + } + Ok(factor) +} + +/// A malformed PK configuration must not reject an unrelated DE query. +pub(super) fn targets_primary_key_column(core: &CoreOptions<'_>, column: &str) -> bool { + core.primary_key_vector_index_enabled() + && core + .primary_key_vector_index_columns() + .ok() + .is_some_and(|columns| columns.iter().any(|c| c == column)) +} diff --git a/crates/paimon/src/table/vector_search_result.rs b/crates/paimon/src/table/vector_search_result.rs new file mode 100644 index 000000000..5ceae40b6 --- /dev/null +++ b/crates/paimon/src/table/vector_search_result.rs @@ -0,0 +1,231 @@ +// 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. + +//! Snapshot-scoped search results for global row IDs and primary-key file positions. + +use crate::spec::{BinaryRow, CoreOptions}; +use crate::table::data_file_reader::DataFileReader; +use crate::table::de_vector_read::{materialize_row_ids, DeVectorRead}; +use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit; +use crate::table::pk_vector_orchestrator::{ + build_indexed_splits, PkVectorCandidate, PkVectorSearchSplit, +}; +use crate::table::pk_vector_read::materialize_positions; +use crate::table::vector_search_common::resolve_materialize_read_type; +use crate::table::{ArrowRecordBatchStream, Table}; +use crate::vector_search::ScoredRowIds; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use std::sync::Arc; + +/// One scored physical row in a primary-key table. Positions are local to the +/// named file; neither the primary key nor a global row ID is synthesized. +#[derive(Debug, Clone)] +pub struct PrimaryKeySearchPosition { + pub partition: BinaryRow, + pub bucket: i32, + pub data_file_name: String, + pub row_position: i64, + pub score: f32, +} + +#[derive(Debug, Clone)] +enum SearchHits { + DataEvolution { + vector_column: String, + hits: ScoredRowIds, + }, + PrimaryKey { + positions: Vec, + splits: Vec, + }, +} + +/// Ranked vector-search hits together with the source snapshot needed to read them. +/// +/// Both single and batch searches return this type. DE hits use global row IDs; +/// PK hits use physical file positions, as in Java's `PrimaryKeyScoredResult`. +/// Searching does not materialize projected user columns. Use +/// [`new_read_builder`](Self::new_read_builder) when rows are needed. +#[derive(Debug, Clone)] +pub struct SearchResult { + table: Arc, + snapshot_id: Option, + hits: SearchHits, +} + +impl SearchResult { + pub(super) fn from_row_ids( + table: Arc
, + vector_column: String, + hits: ScoredRowIds, + ) -> Self { + let snapshot_id = table.travel_snapshot.as_ref().map(|snapshot| snapshot.id()); + Self { + table, + snapshot_id, + hits: SearchHits::DataEvolution { + vector_column, + hits, + }, + } + } + + pub(super) fn from_primary_key( + table: Arc
, + snapshot_id: i64, + candidates: Vec, + source_splits: &[PkVectorSearchSplit], + metric: VectorSearchMetric, + ) -> crate::Result { + let positions = candidates + .iter() + .map(|c| PrimaryKeySearchPosition { + partition: c.partition.clone(), + bucket: c.bucket, + data_file_name: c.data_file_name.clone(), + row_position: c.row_position, + score: metric.distance_to_score(c.distance), + }) + .collect(); + let splits = build_indexed_splits(candidates, source_splits, metric)?; + Ok(Self { + table, + snapshot_id: (snapshot_id != 0).then_some(snapshot_id), + hits: SearchHits::PrimaryKey { positions, splits }, + }) + } + + /// Source snapshot, or `None` when the table had no resolved snapshot. + pub fn snapshot_id(&self) -> Option { + self.snapshot_id + } + + /// The source table retained by the search; DE scans are pinned to this snapshot. + pub fn table(&self) -> &Table { + &self.table + } + + pub fn len(&self) -> usize { + match &self.hits { + SearchHits::DataEvolution { hits, .. } => hits.len(), + SearchHits::PrimaryKey { positions, .. } => positions.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Scored global row IDs, in relevance order. PK positions cannot be converted + /// into this address space, including when a PK search has no hits. + pub fn row_ids(&self) -> crate::Result<&ScoredRowIds> { + match &self.hits { + SearchHits::DataEvolution { hits, .. } => Ok(hits), + SearchHits::PrimaryKey { .. } => Err(no_global_row_ids()), + } + } + + pub fn into_row_ids(self) -> crate::Result { + match self.hits { + SearchHits::DataEvolution { hits, .. } => Ok(hits), + SearchHits::PrimaryKey { .. } => Err(no_global_row_ids()), + } + } + + /// Scored physical positions in relevance order, scoped to this result's snapshot. + pub fn positions(&self) -> crate::Result<&[PrimaryKeySearchPosition]> { + match &self.hits { + SearchHits::PrimaryKey { positions, .. } => Ok(positions), + SearchHits::DataEvolution { .. } => Err(crate::Error::Unsupported { + message: "data-evolution search results use global row IDs, not primary-key file positions".to_string(), + }), + } + } + + /// Snapshot-scoped file metadata and selections, reused when hybrid fusion + /// builds its final selections without scanning the source table again. + pub(super) fn indexed_splits(&self) -> crate::Result<&[PkVectorIndexedSplit]> { + match &self.hits { + SearchHits::PrimaryKey { splits, .. } => Ok(splits), + SearchHits::DataEvolution { .. } => Err(crate::Error::Unsupported { + message: "data-evolution search results do not carry primary-key indexed splits" + .to_string(), + }), + } + } + + pub fn new_read_builder(&self) -> SearchResultReadBuilder<'_> { + SearchResultReadBuilder { + result: self, + projection: None, + } + } +} + +fn no_global_row_ids() -> crate::Error { + crate::Error::Unsupported { + message: "primary-key search results use physical file positions, not global row IDs" + .to_string(), + } +} + +/// Reads selected user columns and `__paimon_search_score`, in relevance order. +/// Internal global row IDs and PK positions are not output. Reading a result +/// never reruns vector search or replans its primary-key source files. +#[derive(Debug, Clone)] +pub struct SearchResultReadBuilder<'a> { + result: &'a SearchResult, + projection: Option>, +} + +impl SearchResultReadBuilder<'_> { + /// Select user columns in this order. Defaults to every user table column. + pub fn with_projection(&mut self, columns: &[&str]) -> &mut Self { + self.projection = Some(columns.iter().map(|name| name.to_string()).collect()); + self + } + + pub async fn read(&self) -> crate::Result { + let table = &self.result.table; + CoreOptions::new(table.schema().options()).ensure_read_authorized()?; + match &self.result.hits { + SearchHits::DataEvolution { + vector_column, + hits, + } => { + let read_type = + DeVectorRead::read_type(table, vector_column, self.projection.as_deref())?; + materialize_row_ids(table, hits, read_type).await + } + SearchHits::PrimaryKey { positions, splits } => { + let read_type = resolve_materialize_read_type(table, self.projection.as_deref())?; + let reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + read_type, + Vec::new(), + ); + materialize_positions(positions, splits, &reader).await + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/vector_search_result/tests.rs b/crates/paimon/src/table/vector_search_result/tests.rs new file mode 100644 index 000000000..c38eb8008 --- /dev/null +++ b/crates/paimon/src/table/vector_search_result/tests.rs @@ -0,0 +1,211 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::table::vector_search_test_utils::{de_vector_table, id_gt_filter, pk_vector_table}; +use crate::vindex::IVF_FLAT_IDENTIFIER; +use arrow_array::{Float32Array, Int32Array, RecordBatch}; +use futures::TryStreamExt; +use roaring::RoaringTreemap; + +#[tokio::test] +async fn de_batch_results_read_projection_and_scores_from_the_search_snapshot() { + let table = de_vector_table().await; + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .execute() + .await + .unwrap(); + assert_eq!(results.len(), 2); + let snapshot_id = results[0].snapshot_id().unwrap(); + assert_eq!(results[1].snapshot_id(), Some(snapshot_id)); + assert_eq!(results[0].row_ids().unwrap().row_ids, vec![0, 2]); + assert_eq!(results[1].row_ids().unwrap().row_ids, vec![1, 2]); + assert!(results[0].positions().is_err()); + + // The result owns the resolved snapshot, so reading does not resolve latest. + let manager = table.snapshot_manager(); + for id in manager.list_all_ids().await.unwrap() { + table + .file_io() + .delete_file(&manager.snapshot_path(id)) + .await + .unwrap(); + } + for (result, expected_ids) in results.iter().zip([vec![1, 3], vec![2, 3]]) { + let batches: Vec = result + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids = Vec::new(); + let mut scores = Vec::new(); + for batch in batches { + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + vec!["id", "__paimon_search_score"] + ); + ids.extend_from_slice( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + scores.extend_from_slice( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + assert_eq!(ids, expected_ids); + assert_eq!(scores, result.row_ids().unwrap().scores); + } +} + +#[tokio::test] +async fn de_empty_filter_result_retains_snapshot_and_validates_projection() { + let table = de_vector_table().await; + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(2) + .with_filter(id_gt_filter(&table, 100)) + .execute() + .await + .unwrap(); + assert!(result.is_empty()); + assert!(result.snapshot_id().is_some()); + assert!(result + .new_read_builder() + .read() + .await + .unwrap() + .try_next() + .await + .unwrap() + .is_none()); + assert!(result + .new_read_builder() + .with_projection(&["_ROW_ID"]) + .read() + .await + .is_err()); +} + +#[tokio::test] +async fn pk_rejects_global_row_id_allow_lists_even_for_an_empty_table() { + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ("fields.embedding.dimension", "4"), + ]); + let err = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0; 4]]) + .with_limit(2) + .with_include_row_ids(RoaringTreemap::from_iter([0])) + .execute() + .await + .unwrap_err(); + assert!(err.to_string().contains("global row-ID filters")); +} + +#[tokio::test] +async fn public_plan_can_be_reused_by_single_and_batch_readers_after_builder_drop() { + let table = de_vector_table().await; + let mut builder = table.new_vector_search_builder(); + builder.with_vector_column("embedding"); + // Scan configuration does not depend on a query vector or Top-K. + let scan = builder.new_scan().unwrap(); + assert!(builder.new_read().is_err()); + builder.with_query_vector(vec![1.0, 0.0]).with_limit(2); + let read = builder.new_read().unwrap(); + drop(builder); + let plan = scan.plan().await.unwrap(); + let snapshot_id = plan.snapshot_id(); + drop(scan); + let single = read.read(plan.clone()).await.unwrap(); + assert_eq!(single.row_ids().unwrap().row_ids, vec![0, 2]); + assert_eq!(single.snapshot_id(), snapshot_id); + let batch = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![-1.0, 1.0], vec![1.0, -1.0]]) + .with_limit(1) + .new_read() + .unwrap(); + let results = batch.read(plan).await.unwrap(); + assert_eq!(results.len(), 2); + for (result, expected) in results.iter().zip([vec![1], vec![0]]) { + assert_eq!(result.row_ids().unwrap().row_ids, expected); + assert_eq!(result.snapshot_id(), snapshot_id); + } +} + +#[tokio::test] +async fn public_reader_rejects_a_plan_for_another_table_or_filter() { + let table = de_vector_table().await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1); + let plan = builder.new_scan().unwrap().plan().await.unwrap(); + builder.with_filter(id_gt_filter(&table, 1)); + let error = builder + .new_read() + .unwrap() + .read(plan.clone()) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("same table, column and pre-filter")); + let other = crate::table::vector_search_test_utils::vector_test_table_at( + "memory:/different_plan_table", + ); + let read = other + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1) + .new_read() + .unwrap(); + let error = read.read(plan).await.unwrap_err(); + assert!(error + .to_string() + .contains("same table, column and pre-filter")); +} diff --git a/crates/paimon/src/table/vector_search_test_utils.rs b/crates/paimon/src/table/vector_search_test_utils.rs new file mode 100644 index 000000000..65c41c5e8 --- /dev/null +++ b/crates/paimon/src/table/vector_search_test_utils.rs @@ -0,0 +1,241 @@ +// 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::catalog::Identifier; +use crate::io::{FileIO, FileIOBuilder}; +use crate::spec::{ + ArrayType, DataType, Datum, FloatType, IntType, Predicate, PredicateBuilder, Schema, + TableSchema, +}; +use crate::table::{Table, TableCommit, TableWrite}; +use crate::vindex::IVF_FLAT_IDENTIFIER; +use arrow_array::builder::{Float32Builder, ListBuilder}; +use arrow_array::{ArrayRef, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use std::collections::HashMap; +use std::sync::Arc; + +pub(super) fn vector_test_table() -> Table { + vector_test_table_at("memory:/vector_test") +} + +pub(super) fn vector_test_table_at(location: &str) -> Table { + vector_test_table_with_file_io(FileIOBuilder::new("memory").build().unwrap(), location) +} + +pub(super) fn vector_test_table_with_file_io(file_io: FileIO, location: &str) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .build() + .unwrap(); + Table::new( + file_io, + Identifier::new("default", "vector_test"), + location.to_string(), + TableSchema::new(0, &schema), + None, + ) +} + +/// Build a real vindex IVF-flat segment trained with `metric`, returning the +/// serialized bytes. `nlist = 1` keeps training trivial and deterministic; the +/// only thing the metric check cares about is the persisted metadata metric. +pub(super) fn build_vindex_segment_bytes(metric: &str) -> Vec { + use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; + use paimon_vindex_core::io::PosWriter; + + const DIM: usize = 2; + let vectors: Vec = vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let n = vectors.len() / DIM; + let ids: Vec = (0..n as i64).collect(); + let options = HashMap::from([ + ("index.type".to_string(), "ivf_flat".to_string()), + ("dimension".to_string(), DIM.to_string()), + ("nlist".to_string(), "1".to_string()), + ("metric".to_string(), metric.to_string()), + ]); + let config = VectorIndexConfig::from_options(&options).unwrap(); + let training = VectorIndexTrainer::train(config, &vectors, n).unwrap(); + let mut writer = VectorIndexWriter::new(training); + writer.add_vectors(&ids, &vectors, n).unwrap(); + let mut bytes = Vec::new(); + { + let mut output = PosWriter::new(&mut bytes); + writer.write(&mut output).unwrap(); + } + bytes +} + +pub(super) fn pk_vector_table(options: &[(&str, &str)]) -> Table { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + if options + .iter() + .any(|(key, _)| *key == "pk-vector.index.columns") + { + builder = builder.primary_key(["id"]).option("bucket", "1"); + } + let schema = builder.build().unwrap(); + // Runtime validation must remain defensive for schemas committed by old + // or external writers, including malformed configurations which the + // current Schema builder rejects at commit time. + let runtime_options = options + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect(); + let table_schema = TableSchema::new(0, &schema).copy_with_options(runtime_options); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "pk_vector_test"), + "memory:/pk_vector_test".to_string(), + table_schema, + None, + ) +} + +/// A data-evolution (global-index) vector table with a committed IVF-flat +/// index over the `embedding` column: row-tracking + data-evolution + +/// global-index enabled so committed data files carry `first_row_id` and the +/// search returns global row-ids that `execute_read` can materialize. The +/// returned table has one committed batch of `(id, embedding)` rows and a real +/// vindex index built end-to-end. +pub(super) async fn de_vector_table() -> Table { + let table_path = "memory:/de_vector_search_test"; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .option("global-index.enabled", "true") + .option("global-index.row-count-per-shard", "10") + .option("ivf-flat.dimension", "2") + .option("ivf-flat.nlist", "2") + .build() + .unwrap(); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "de_vector_test"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + file_io + .mkdirs(&format!("{table_path}/snapshot/")) + .await + .unwrap(); + file_io + .mkdirs(&format!("{table_path}/manifest/")) + .await + .unwrap(); + + let ids = vec![1, 2, 3]; + let vectors = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]]; + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = + ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); + for vector in vectors { + for value in vector { + vector_builder.values().append_value(value); + } + vector_builder.append(true); + } + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element_field), true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap(); + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let built = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .unwrap(); + assert!(built > 0, "DE fixture must build a global vector index"); + let built = table + .new_sorted_global_index_build_builder() + .with_index_column("id") + .with_index_type("btree") + .execute() + .await + .unwrap(); + assert!(built > 0, "DE fixture must build a scalar BTree index"); + table +} + +/// `id > threshold` built against the table's user fields (leaf index resolves +/// against `table.schema().fields()`). +pub(super) fn id_gt_filter(table: &Table, threshold: i32) -> Predicate { + PredicateBuilder::new(table.schema().fields()) + .greater_than("id", Datum::Int(threshold)) + .unwrap() +} + +/// A PK-vector table whose user schema carries an extra column named +/// `reserved`, used to prove reserved metadata names are rejected even when +/// they arrive via the default (all-columns) projection. +pub(super) fn pk_vector_table_with_extra_column(reserved: &str) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .column(reserved, DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "1") + .option("deletion-vectors.enabled", "true") + .option("pk-vector.index.columns", "embedding") + .option("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER) + .option("fields.embedding.pk-vector.distance.metric", "l2") + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "reserved_col_test"), + "memory:/reserved_col_test".to_string(), + TableSchema::new(0, &schema), + None, + ) +} diff --git a/crates/paimon/src/vector_search.rs b/crates/paimon/src/vector_search.rs index e1a3de10d..0599b3f49 100644 --- a/crates/paimon/src/vector_search.rs +++ b/crates/paimon/src/vector_search.rs @@ -15,6 +15,10 @@ // specific language governing permissions and limitations // under the License. +pub use crate::table::vector_search_result::{ + PrimaryKeySearchPosition, SearchResult, SearchResultReadBuilder, +}; + use std::cmp::Ordering; use std::collections::{BinaryHeap, HashMap}; use std::sync::Arc; @@ -111,8 +115,10 @@ impl GlobalIndexIOMeta { } } +/// Global row IDs and aligned scores used by DE index evaluation and ranking. +/// Table searches return [`SearchResult`], which also represents PK file positions. #[derive(Debug, Clone)] -pub struct SearchResult { +pub struct ScoredRowIds { pub row_ids: Vec, pub scores: Vec, } @@ -157,7 +163,7 @@ fn sort_scored_rows_by_rank(rows: &mut [ScoredRow]) { }); } -impl SearchResult { +impl ScoredRowIds { pub fn new(row_ids: Vec, scores: Vec) -> Self { assert_eq!(row_ids.len(), scores.len()); Self { row_ids, scores } @@ -209,7 +215,7 @@ impl SearchResult { } } - pub fn or(&self, other: &SearchResult) -> Self { + pub fn or(&self, other: &ScoredRowIds) -> Self { let mut row_ids = self.row_ids.clone(); let mut scores = self.scores.clone(); row_ids.extend_from_slice(&other.row_ids); @@ -361,13 +367,13 @@ mod tests { let mut map = HashMap::new(); map.insert(1u64, 0.9f32); map.insert(2, 0.5); - let result = SearchResult::from_scored_map(map); + let result = ScoredRowIds::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 result = ScoredRowIds::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)); @@ -376,8 +382,8 @@ mod tests { #[test] fn test_search_result_top_k_deduplicates_overlapping_rows() { - let indexed = SearchResult::new(vec![1], vec![0.9]); - let fallback = SearchResult::new(vec![1, 2], vec![0.8, 0.7]); + let indexed = ScoredRowIds::new(vec![1], vec![0.9]); + let fallback = ScoredRowIds::new(vec![1, 2], vec![0.8, 0.7]); let merged = indexed.or(&fallback); assert_eq!(merged.row_ids, vec![1, 1, 2]); @@ -390,7 +396,7 @@ mod tests { #[test] fn test_search_result_top_k_keeps_highest_duplicate_score() { - let result = SearchResult::new(vec![1, 2, 1, 3], vec![0.5, 0.8, 0.9, 0.7]); + let result = ScoredRowIds::new(vec![1, 2, 1, 3], vec![0.5, 0.8, 0.9, 0.7]); let top = result.top_k(2); assert_eq!(top.row_ids, vec![1, 2]); @@ -401,7 +407,7 @@ mod tests { fn test_search_result_top_k_sorts_best_first_without_truncation() { // Even when k >= candidate count (no truncation), results must be returned // best-first by score, not in the input/insertion order. - let result = SearchResult::new(vec![3, 1, 2], vec![0.1, 0.9, 0.5]); + let result = ScoredRowIds::new(vec![3, 1, 2], vec![0.1, 0.9, 0.5]); let top = result.top_k(3); assert_eq!(top.row_ids, vec![1, 2, 3]); @@ -410,7 +416,7 @@ mod tests { #[test] fn test_search_result_top_k_tie_breaks_by_smaller_row_id() { - let result = SearchResult::new(vec![30, 10, 20], vec![0.9, 0.9, 0.9]); + let result = ScoredRowIds::new(vec![30, 10, 20], vec![0.9, 0.9, 0.9]); let top = result.top_k(2); assert_eq!(top.row_ids, vec![10, 20]); assert_eq!(top.scores, vec![0.9, 0.9]); @@ -418,7 +424,7 @@ mod tests { #[test] fn test_search_result_filters_deleted_row_ranges() { - let result = SearchResult::new(vec![1, 2, 3, 4], vec![0.1, 0.9, 0.8, 0.2]); + let result = ScoredRowIds::new(vec![1, 2, 3, 4], vec![0.1, 0.9, 0.8, 0.2]); let deleted = crate::table::global_index_scanner::RowRangeIndex::create(vec![ crate::table::RowRange::new(2, 3), ]); @@ -434,7 +440,7 @@ mod tests { #[test] fn test_search_result_offset() { - let result = SearchResult::new(vec![0, 1], vec![0.5, 0.6]); + let result = ScoredRowIds::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]); @@ -442,15 +448,15 @@ mod tests { #[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 a = ScoredRowIds::new(vec![1, 2], vec![0.5, 0.6]); + let b = ScoredRowIds::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 result = ScoredRowIds::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); @@ -463,7 +469,7 @@ mod tests { #[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 result = ScoredRowIds::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"), diff --git a/crates/paimon/src/vindex/executor.rs b/crates/paimon/src/vindex/executor.rs index b604ae611..38c6278ca 100644 --- a/crates/paimon/src/vindex/executor.rs +++ b/crates/paimon/src/vindex/executor.rs @@ -321,9 +321,12 @@ pub(crate) async fn drain_indexed_jobs( where F: std::future::Future>, { + // Materialize the wrappers before awaiting so their borrowing iterator + // closures do not prevent callers from exposing a Send future. let indexed = jobs .enumerate() - .map(|(index, job)| async move { (index, job.await) }); + .map(|(index, job)| async move { (index, job.await) }) + .collect::>(); let mut collected: Vec<(usize, crate::Result)> = stream::iter(indexed) .buffer_unordered(concurrency.max(1)) .collect() diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index 7e0a0627b..a724544f0 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -77,6 +77,7 @@ impl BucketAnnSegment { /// A data file participating in the bucket search, with its row count. Used by /// the bucket kernel to plan exact vs. ANN search over active files. +#[derive(Clone)] pub(crate) struct BucketActiveFile { pub file_name: String, pub row_count: i64, @@ -527,7 +528,8 @@ pub(crate) async fn bucket_search( // Eligible uncovered exact files (active-file order) with their exclusion // predicate; a file with no residual-allowed rows is skipped without reading. #[allow(clippy::type_complexity)] - let mut exact_tasks: Vec<(&BucketActiveFile, Box bool + Sync>)> = Vec::new(); + let mut exact_tasks: Vec<(&BucketActiveFile, Box bool + Send + Sync>)> = + Vec::new(); if !skip_exact_fallback { for file in active_files { if covered.contains(&file.file_name) { @@ -816,7 +818,8 @@ pub(crate) async fn bucket_search_batch( }); #[allow(clippy::type_complexity)] - let mut exact_tasks: Vec<(&BucketActiveFile, Box bool + Sync>)> = Vec::new(); + let mut exact_tasks: Vec<(&BucketActiveFile, Box bool + Send + Sync>)> = + Vec::new(); if !skip_exact_fallback { for file in active_files { if covered.contains(&file.file_name) { diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs b/crates/paimon/tests/pk_vector_baseline_test.rs index 8b24b2b56..667cc7bcc 100644 --- a/crates/paimon/tests/pk_vector_baseline_test.rs +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -21,8 +21,8 @@ //! entirely from Rust — data file, a real vindex IVF-flat ANN index segment, and //! the snapshot/manifest/index-manifest metadata — then reads it back through the //! public `new_vector_search_builder()` API and asserts both the search result -//! (`execute_scored()` -> `row_ids`/`scores`) and the materialized rows -//! (`execute_read()` -> Arrow batches, best-first order, `__paimon_search_score`). +//! (`execute()` -> physical positions and scores) and the materialized rows +//! (result reader -> Arrow batches, best-first order, `__paimon_search_score`). //! //! Why Rust-built rather than a committed cross-language fixture: the Java //! primary-key vector ANN segment is an opaque native Lumina format that cannot @@ -452,7 +452,7 @@ async fn build_table_with_first_row_id( (tmp, table) } -/// Run `execute_read()` and flatten the stream into per-row `(id, score)` tuples +/// Run `new_read_builder().read()` and flatten the stream into per-row `(id, score)` tuples /// in emission order (best-first), returning the collected batches too for /// schema / row-content assertions. async fn read_id_and_scores( @@ -466,11 +466,13 @@ async fn read_id_and_scores( .with_vector_column(VECTOR_COLUMN) .with_query_vector(query) .with_limit(limit); + let result = builder.execute().await.unwrap(); + let mut reader = result.new_read_builder(); if let Some(cols) = projection { - builder.with_projection(cols); + reader.with_projection(cols); } - let batches = builder - .execute_read() + let batches = reader + .read() .await .expect("primary-key vector read failed") .try_collect::>() @@ -585,23 +587,36 @@ async fn pk_vector_end_to_end_returns_expected_row_ids_and_scores() { let expected_row_ids: Vec = expected.iter().map(|(id, _)| *id).collect(); let expected_scores: Vec = expected.iter().map(|(_, d)| l2_score(*d)).collect(); - // A primary-key vector table exposes no global row ids, so the search-only - // `execute_scored()` path is unsupported and must fail loud, directing callers - // to the materialized `execute_read()` path exercised below. - let scored_err = table + // Search-only PK results preserve local positions and metric scores. + let scored = table .new_vector_search_builder() .with_vector_column(VECTOR_COLUMN) .with_query_vector(query.to_vec()) .with_limit(3) - .execute_scored() + .execute() .await - .expect_err("primary-key execute_scored must fail loud"); - assert!( - format!("{scored_err:?}").contains("execute_read"), - "primary-key execute_scored should point at execute_read, got: {scored_err:?}" + .unwrap(); + assert!(scored.row_ids().is_err()); + assert_eq!( + scored + .positions() + .unwrap() + .iter() + .map(|p| p.row_position as u64) + .collect::>(), + expected_row_ids + ); + assert_eq!( + scored + .positions() + .unwrap() + .iter() + .map(|p| p.score) + .collect::>(), + expected_scores ); - // Search-and-read: execute_read() materializes the matching rows best-first + // Search-and-read: the result reader materializes the matching rows best-first // with a `__paimon_search_score` column, hiding `_ROW_ID`/`_PKEY_VECTOR_POSITION`. // Projection ['id'] excludes the vector column. let (ids, scores, batches) = read_id_and_scores(&table, query.to_vec(), 3, Some(&["id"])).await; @@ -775,7 +790,7 @@ async fn assert_discriminating_local_read(first_row_id: Option) { // Gated off Windows for the same `file://` tempdir reason as the tests above. #[cfg(not(windows))] #[tokio::test] -async fn execute_read_without_first_row_id_selects_local_positions() { +async fn result_read_without_first_row_id_selects_local_positions() { assert_discriminating_local_read(None).await; } @@ -788,7 +803,7 @@ async fn execute_read_without_first_row_id_selects_local_positions() { // Gated off Windows for the same `file://` tempdir reason as the tests above. #[cfg(not(windows))] #[tokio::test] -async fn execute_read_ignores_nonzero_first_row_id() { +async fn result_read_ignores_nonzero_first_row_id() { assert_discriminating_local_read(Some(100)).await; } @@ -823,7 +838,7 @@ fn fixture_residual() -> ([f32; DIM], Vec<[f32; DIM]>) { (query, vectors) } -/// Run `execute_read()` with a residual `filter` attached via `with_filter` and +/// Run `new_read_builder().read()` with a residual `filter` attached via `with_filter` and /// flatten the stream into per-row `(id, score)` tuples in emission order /// (best-first), returning the collected batches too for schema / row-content /// assertions. Mirrors `read_id_and_scores` but exercises the residual path. @@ -840,8 +855,7 @@ async fn read_id_and_scores_filtered( .with_query_vector(query) .with_limit(limit) .with_filter(filter); - let batches = builder - .execute_read() + let batches = async { builder.execute().await?.new_read_builder().read().await } .await .expect("primary-key vector residual read failed") .try_collect::>() @@ -924,34 +938,44 @@ async fn pk_vector_residual_filter_excludes_non_matching_rows() { .greater_or_equal("id", Datum::Int(residual_threshold as i32)) .expect("build residual predicate on id"); - // A primary-key vector table exposes no global row ids, so `execute_scored()` - // is unsupported on this path — with or without a residual filter — and must - // fail loud, directing callers to the materialized `execute_read()` used below. - let unfiltered_err = table + let unfiltered_result = table .new_vector_search_builder() .with_vector_column(VECTOR_COLUMN) .with_query_vector(query.to_vec()) .with_limit(3) - .execute_scored() + .execute() .await - .expect_err("primary-key execute_scored must fail loud"); - assert!( - format!("{unfiltered_err:?}").contains("execute_read"), - "got: {unfiltered_err:?}" + .unwrap(); + assert_eq!( + unfiltered_result + .positions() + .unwrap() + .iter() + .map(|p| p.row_position as i64) + .collect::>(), + unfiltered_ids + .iter() + .map(|id| *id as i64) + .collect::>() ); - let residual_scored_err = table + let residual_result = table .new_vector_search_builder() .with_vector_column(VECTOR_COLUMN) .with_query_vector(query.to_vec()) .with_limit(3) .with_filter(residual.clone()) - .execute_scored() + .execute() .await - .expect_err("primary-key execute_scored must fail loud with a residual too"); - assert!( - format!("{residual_scored_err:?}").contains("execute_read"), - "got: {residual_scored_err:?}" + .unwrap(); + assert_eq!( + residual_result + .positions() + .unwrap() + .iter() + .map(|p| p.row_position as i64) + .collect::>(), + expected_ids.iter().map(|id| *id as i64).collect::>() ); // Search-and-read with the residual: default projection materializes id + @@ -1118,7 +1142,7 @@ async fn write_schema_and_data( ) } -/// Read `execute_read()` into `(id, score)` tuples (best-first) plus the batches, +/// Read `new_read_builder().read()` into `(id, score)` tuples (best-first) plus the batches, /// mirroring [`read_id_and_scores`] but with caller-supplied query options so the /// refine factor can be requested. async fn read_id_and_scores_with_options( @@ -1133,8 +1157,7 @@ async fn read_id_and_scores_with_options( .with_query_vector(query) .with_limit(limit) .with_options(options); - let batches = builder - .execute_read() + let batches = async { builder.execute().await?.new_read_builder().read().await } .await .expect("primary-key vector rerank read failed") .try_collect::>() @@ -1374,7 +1397,7 @@ async fn pk_vector_refine_factor_with_no_indexed_candidates_is_noop() { .unwrap(); // A positive refine factor is set, but with no indexed candidates the rerank is - // gated off: execute_read must not error, and returns the exact-fallback rows. + // gated off: result_read must not error, and returns the exact-fallback rows. let (ids, scores, batches) = read_id_and_scores_with_options(&table, query.to_vec(), k, refine_factor_option(2)).await; assert_eq!( @@ -1437,7 +1460,7 @@ async fn pk_vector_invalid_refine_factor_fails_loud_on_empty_table() { .with_query_vector(vec![0.0f32; DIM]) .with_limit(3) .with_options(HashMap::from([(refine_key.clone(), "abc".to_string())])); - let err = match builder.execute_read().await { + let err = match async { builder.execute().await?.new_read_builder().read().await }.await { Ok(_) => panic!("a non-integer refine factor must fail loud on an empty table"), Err(e) => e, }; @@ -1460,7 +1483,7 @@ async fn pk_vector_invalid_refine_factor_fails_loud_on_empty_table() { .with_query_vector(vec![0.0f32; DIM]) .with_limit(3) .with_options(HashMap::from([(refine_key.clone(), "0".to_string())])); - let err = match builder.execute_read().await { + let err = match async { builder.execute().await?.new_read_builder().read().await }.await { Ok(_) => panic!("a zero refine factor must fail loud on an empty table"), Err(e) => e, }; @@ -1482,7 +1505,7 @@ async fn pk_vector_invalid_refine_factor_fails_loud_on_empty_table() { .with_vector_column(VECTOR_COLUMN) .with_query_vector(vec![0.0f32; DIM]) .with_limit(3); - let err = match builder.execute_read().await { + let err = match async { builder.execute().await?.new_read_builder().read().await }.await { Ok(_) => panic!("a non-integer table refine factor must fail loud on an empty table"), Err(e) => e, }; diff --git a/crates/paimon/tests/pk_vector_batch_test.rs b/crates/paimon/tests/pk_vector_batch_test.rs index e2a21d210..c7fcb32d8 100644 --- a/crates/paimon/tests/pk_vector_batch_test.rs +++ b/crates/paimon/tests/pk_vector_batch_test.rs @@ -15,18 +15,9 @@ // specific language governing permissions and limitations // under the License. -//! End-to-end acceptance gate for BATCH primary-key vector search. -//! -//! Builds a complete, self-contained primary-key vector table entirely from Rust -//! (mirroring `pk_vector_baseline_test`), then reads it back through the public -//! batch surface `new_batch_vector_search_builder().execute_read()` and asserts: -//! - batch-of-one == the single-query `execute_read`; -//! - an N-query batch yields one stream per query, each matching the -//! corresponding independent single-query read (arity + order + independence); -//! - an empty snapshot yields N empty streams (arity preserved); -//! - the PK batch `execute()` (scored) fails loud, directing to `execute_read`; -//! - a shared residual filter reshapes every query's rows with no cross-query -//! bleed. +//! Batch PK search returns one snapshot-scoped result per query, in input order. +//! These fixtures compare positions and projected result reads against independent +//! single queries, including empty results and residual-filtered queries. use std::collections::HashMap; @@ -382,15 +373,14 @@ async fn drain_ids_and_scores(stream: ArrowRecordBatchStream) -> (Vec, Vec< (ids, scores) } -/// Single-query `execute_read` into `(id, score)` tuples. +/// Single-query `SearchResultReadBuilder::read` into `(id, score)` tuples. async fn single_read(table: &Table, query: Vec, limit: usize) -> (Vec, Vec) { let mut builder = table.new_vector_search_builder(); builder .with_vector_column(VECTOR_COLUMN) .with_query_vector(query) .with_limit(limit); - let stream = builder - .execute_read() + let stream = async { builder.execute().await?.new_read_builder().read().await } .await .expect("single-query read failed"); drain_ids_and_scores(stream).await @@ -419,15 +409,16 @@ async fn batch_of_one_equals_single_read() { let (single_ids, single_scores) = single_read(&table, query.clone(), 3).await; let mut builder = table.new_batch_vector_search_builder(); - let mut streams = builder + let mut results = builder .with_vector_column(VECTOR_COLUMN) .with_query_vectors(vec![query]) .with_limit(3) - .execute_read() + .execute() .await .expect("batch-of-one read failed"); - assert_eq!(streams.len(), 1, "batch-of-one yields exactly one stream"); - let (batch_ids, batch_scores) = drain_ids_and_scores(streams.remove(0)).await; + assert_eq!(results.len(), 1, "batch-of-one yields exactly one stream"); + let (batch_ids, batch_scores) = + drain_ids_and_scores(results.remove(0).new_read_builder().read().await.unwrap()).await; assert_eq!(batch_ids, single_ids); assert_eq!(batch_scores.len(), single_scores.len()); @@ -454,21 +445,22 @@ async fn n_query_batch_matches_n_independent_single_reads() { } let mut builder = table.new_batch_vector_search_builder(); - let streams = builder + let results = builder .with_vector_column(VECTOR_COLUMN) .with_query_vectors(queries.clone()) .with_limit(3) - .execute_read() + .execute() .await .expect("batch read failed"); assert_eq!( - streams.len(), + results.len(), queries.len(), "one stream per query, in input order" ); - for (i, stream) in streams.into_iter().enumerate() { - let (ids, scores) = drain_ids_and_scores(stream).await; + for (i, result) in results.into_iter().enumerate() { + let (ids, scores) = + drain_ids_and_scores(result.new_read_builder().read().await.unwrap()).await; let (want_ids, want_scores) = &expected[i]; assert_eq!(&ids, want_ids, "query {i} rows must match its single read"); assert_eq!(scores.len(), want_scores.len()); @@ -483,7 +475,7 @@ async fn n_query_batch_matches_n_independent_single_reads() { #[cfg(not(windows))] #[tokio::test] -async fn empty_snapshot_yields_n_empty_streams() { +async fn empty_snapshot_yields_n_empty_results() { let (_tmp, table) = build_empty_table().await; let queries = vec![ vec![1.0, 0.0, 0.0, 0.0], @@ -492,41 +484,67 @@ async fn empty_snapshot_yields_n_empty_streams() { ]; let mut builder = table.new_batch_vector_search_builder(); - let streams = builder + let results = builder .with_vector_column(VECTOR_COLUMN) .with_query_vectors(queries.clone()) .with_limit(3) - .execute_read() + .execute() .await .expect("empty-snapshot batch read failed"); assert_eq!( - streams.len(), + results.len(), queries.len(), "arity preserved: one empty stream per query" ); - for stream in streams { - let (ids, _scores) = drain_ids_and_scores(stream).await; + for result in results { + let (ids, _scores) = + drain_ids_and_scores(result.new_read_builder().read().await.unwrap()).await; assert!(ids.is_empty(), "no-hit query must yield an empty stream"); } } #[cfg(not(windows))] #[tokio::test] -async fn pk_batch_execute_scored_fails_loud() { +async fn pk_batch_execute_returns_positions_and_reads_without_replanning() { let vectors = fixture(); let (_tmp, table) = build_table(&vectors).await; - let err = table + let mut results = table .new_batch_vector_search_builder() .with_vector_column(VECTOR_COLUMN) .with_query_vectors(vec![vec![10.0, 0.0, 0.0, 0.0]]) .with_limit(3) .execute() .await - .expect_err("PK batch execute() must fail loud"); - assert!( - format!("{err:?}").contains("execute_read"), - "PK batch execute() should point at execute_read, got: {err:?}" + .unwrap(); + assert_eq!(results.len(), 1); + let result = results.remove(0); + assert_eq!(result.snapshot_id(), Some(1)); + assert!(result.row_ids().is_err()); + let positions = result.positions().unwrap(); + assert_eq!( + positions.iter().map(|p| p.row_position).collect::>(), + vec![5, 1, 3] ); + let expected_scores = positions.iter().map(|p| p.score).collect::>(); + + // Search has finished. Removing the manifest makes replanning impossible, + // while the retained file positions still allow a projected read. + let manager = table.snapshot_manager(); + let snapshot = manager.get_snapshot(1).await.unwrap(); + table + .file_io() + .delete_file(&manager.manifest_path(snapshot.index_manifest().unwrap())) + .await + .unwrap(); + let stream = result + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + .unwrap(); + let (ids, scores) = drain_ids_and_scores(stream).await; + assert_eq!(ids, vec![5, 1, 3]); + assert_eq!(scores, expected_scores); } /// A shared residual filter reshapes every query's rows (excluding low ids) with @@ -574,18 +592,19 @@ async fn shared_residual_filter_applies_per_query_without_bleed() { ); let mut builder = table.new_batch_vector_search_builder(); - let streams = builder + let results = builder .with_vector_column(VECTOR_COLUMN) .with_query_vectors(queries.clone()) .with_limit(3) .with_filter(build_filter()) - .execute_read() + .execute() .await .expect("residual batch read failed"); - assert_eq!(streams.len(), queries.len()); + assert_eq!(results.len(), queries.len()); - for (i, stream) in streams.into_iter().enumerate() { - let (ids, _scores) = drain_ids_and_scores(stream).await; + for (i, result) in results.into_iter().enumerate() { + let (ids, _scores) = + drain_ids_and_scores(result.new_read_builder().read().await.unwrap()).await; for &id in &ids { assert!( id >= threshold, @@ -596,14 +615,12 @@ async fn shared_residual_filter_applies_per_query_without_bleed() { } } -/// A table with no primary-key vector index is not a valid target for the batch -/// materialized read: it produces scored global row ids, not physical rows. The -/// batch `execute_read` must reject it up front rather than silently route it -/// through the primary-key materialization path. +/// A table without a PK-vector index uses DE results, including the shared +/// result-reading API for an empty snapshot. // Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. #[cfg(not(windows))] #[tokio::test] -async fn batch_execute_read_on_non_pk_vector_table_fails_loud() { +async fn batch_empty_de_result_can_be_read() { let tmp = tempfile::tempdir().expect("create temp dir"); let location = format!("file://{}", tmp.path().display()); let file_io = FileIOBuilder::new("file").build().unwrap(); @@ -638,16 +655,13 @@ async fn batch_execute_read_on_non_pk_vector_table_fails_loud() { .with_vector_column(VECTOR_COLUMN) .with_query_vectors(vec![vec![1.0, 0.0, 0.0, 0.0]]) .with_limit(3) - .execute_read() - .await; - let err = match result { - Ok(_) => panic!("batch execute_read on a non-PK-vector table must fail loud"), - Err(e) => e, - }; - assert!( - err.to_string().contains("primary-key vector path"), - "expected a message directing to the primary-key vector path, got: {err}" - ); + .execute() + .await + .unwrap(); + assert_eq!(result.len(), 1); + assert!(result[0].row_ids().unwrap().is_empty()); + let (ids, _) = drain_ids_and_scores(result[0].new_read_builder().read().await.unwrap()).await; + assert!(ids.is_empty()); } /// A malformed query (wrong dimension) must fail loud even when the plan is @@ -666,7 +680,7 @@ async fn empty_snapshot_still_rejects_malformed_query() { .with_vector_column(VECTOR_COLUMN) .with_query_vectors(queries) .with_limit(3) - .execute_read() + .execute() .await; let err = match result { Ok(_) => panic!("a wrong-dimension query must fail loud even on an empty snapshot"), @@ -693,7 +707,7 @@ async fn empty_array_snapshot_still_rejects_non_finite_query() { .with_vector_column(VECTOR_COLUMN) .with_query_vectors(queries) .with_limit(3) - .execute_read() + .execute() .await; let err = match result { Ok(_) => panic!("a NaN query must fail loud for ARRAY even on an empty snapshot"), @@ -719,7 +733,7 @@ async fn empty_snapshot_still_rejects_zero_limit() { .with_vector_column(VECTOR_COLUMN) .with_query_vectors(vec![vec![1.0, 0.0, 0.0, 0.0]]) .with_limit(0) - .execute_read() + .execute() .await; let err = match result { Ok(_) => panic!("a zero limit must fail loud even on an empty snapshot"), diff --git a/crates/paimon/tests/pk_vector_bucket_split_read_test.rs b/crates/paimon/tests/pk_vector_bucket_split_read_test.rs index d0dec3949..25edf9d95 100644 --- a/crates/paimon/tests/pk_vector_bucket_split_read_test.rs +++ b/crates/paimon/tests/pk_vector_bucket_split_read_test.rs @@ -183,20 +183,32 @@ fn batch_f32(batches: &[RecordBatch], column: &str) -> Vec { } async fn read_over_splits(table: &Table, splits: &[Vec], limit: usize) -> Vec { - let refs: Vec<&[u8]> = splits.iter().map(Vec::as_slice).collect(); + let splits = splits + .iter() + .map(|bytes| BucketVectorSearchSplit::deserialize(bytes)) + .collect::>>() + .unwrap(); let mut builder = table.new_vector_search_builder(); builder .with_vector_column(VECTOR_COLUMN) .with_query_vector(vec![0.0, 0.0]) - .with_limit(limit) - .with_projection(&["id"]); - builder - .execute_read_for_bucket_splits(&refs) - .await - .expect("bucket-split read over the Java fixture failed") - .try_collect::>() - .await - .expect("collecting read batches failed") + .with_limit(limit); + async { + let plan = builder.new_scan()?.plan_from_bucket_splits(splits)?; + builder + .new_read()? + .read(plan) + .await? + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + } + .await + .expect("bucket-split read over the Java fixture failed") + .try_collect::>() + .await + .expect("collecting read batches failed") } /// The read is driven entirely by the Java-planned split: no index manifest is @@ -236,15 +248,21 @@ async fn agrees_with_the_manifest_route() { builder .with_vector_column(VECTOR_COLUMN) .with_query_vector(vec![0.0, 0.0]) - .with_limit(3) - .with_projection(&["id"]); - let from_manifest = builder - .execute_read() - .await - .expect("manifest-route read failed") - .try_collect::>() - .await - .expect("collecting manifest-route batches failed"); + .with_limit(3); + let from_manifest = async { + builder + .execute() + .await? + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + } + .await + .expect("manifest-route read failed") + .try_collect::>() + .await + .expect("collecting manifest-route batches failed"); assert_eq!( batch_i32(&from_splits, "id"), @@ -279,12 +297,24 @@ async fn rejects_an_empty_split_list() { .with_vector_column(VECTOR_COLUMN) .with_query_vector(vec![0.0, 0.0]) .with_limit(3); - let error = match builder.execute_read_for_bucket_splits(&[]).await { + let error = match async { + let plan = builder.new_scan()?.plan_from_bucket_splits(Vec::new())?; + builder + .new_read()? + .read(plan) + .await? + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + } + .await + { Ok(_) => panic!("an empty split list must be rejected"), Err(e) => e, }; assert!( - error.to_string().contains("at least one split"), + error.to_string().contains("at least one bucket split"), "unexpected error: {error}" ); } @@ -293,10 +323,7 @@ async fn rejects_an_empty_split_list() { /// invalid data rather than as an internal fault. #[tokio::test] async fn rejects_corrupt_split_bytes() { - // A decoder assertion, not a read assertion. The entry point takes SERIALIZED - // bytes and decodes them at that boundary, so corrupt input is refused there - // and never reaches a search; driving the read as well would run the same - // decoder behind a query that cannot execute either way. + // Decode at the caller boundary before creating a typed plan or reader. let (_tmp, _table, splits) = open_bucket_split_fixture().await; let mut corrupt = splits[0].clone(); corrupt[0] ^= 0xFF; // break the PKVSPLIT magic @@ -365,3 +392,59 @@ fn with_row_range(bytes: &[u8], file: &str, to: i64) -> Vec { out.extend_from_slice(&to.to_be_bytes()); out } + +#[tokio::test] +async fn typed_bucket_plan_is_reusable_without_the_index_manifest() { + let (_tmp, table, bytes) = open_bucket_split_fixture().await; + let splits = bytes + .iter() + .map(|bytes| BucketVectorSearchSplit::deserialize(bytes)) + .collect::>>() + .unwrap(); + let mut builder = table.new_vector_search_builder(); + builder.with_vector_column(VECTOR_COLUMN); + let scan = builder.new_scan().unwrap(); + let plan = scan.plan_from_bucket_splits(splits).unwrap(); + let snapshot_id = plan.snapshot_id().unwrap(); + let manager = table.snapshot_manager(); + let snapshot = manager.get_snapshot(snapshot_id).await.unwrap(); + table + .file_io() + .delete_file(&manager.manifest_path(snapshot.index_manifest().unwrap())) + .await + .unwrap(); + assert!( + scan.plan().await.is_err(), + "replanning must require the removed manifest" + ); + builder.with_query_vector(vec![0.0, 0.0]).with_limit(1); + let read = builder.new_read().unwrap(); + drop(builder); + drop(scan); + let result = read.read(plan.clone()).await.unwrap(); + assert_eq!(result.positions().unwrap()[0].row_position, 0); + assert_eq!(result.snapshot_id(), Some(snapshot_id)); + let read = table + .new_batch_vector_search_builder() + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(vec![vec![0.0, 0.0], vec![2.0, 0.0]]) + .with_limit(1) + .new_read() + .unwrap(); + let results = read.read(plan).await.unwrap(); + assert_eq!(results.len(), 2); + for (result, row_position) in results.iter().zip([0, 2]) { + assert_eq!(result.positions().unwrap()[0].row_position, row_position); + assert_eq!(result.snapshot_id(), Some(snapshot_id)); + let batches = result + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batch_i32(&batches, "id"), vec![row_position as i32]); + } +} diff --git a/crates/paimon/tests/pk_vector_java_fixture_test.rs b/crates/paimon/tests/pk_vector_java_fixture_test.rs index f64584bbe..91cfcd9d7 100644 --- a/crates/paimon/tests/pk_vector_java_fixture_test.rs +++ b/crates/paimon/tests/pk_vector_java_fixture_test.rs @@ -30,25 +30,8 @@ //! Physical-position contract: Java never writes `first_row_id` on a primary-key //! table (row-tracking is forbidden for PK tables), so the committed data files //! below carry NO `first_row_id`. Row identity is therefore physical -//! `(file, position)`, not a global row id. Two consequences the assertions pin: -//! * `execute_scored()` — which reports global row ids — MUST fail on this -//! table, because a global row id cannot be recovered without `first_row_id`. -//! * `execute_read()` — which materializes rows by physical position — MUST -//! succeed and return the rows best-first. -//! -//! Provenance of `testdata/pkvector/pk_vector_ivf_flat` (opaque binary table -//! directory, regenerate rather than hand-edit): -//! * Source: Apache Paimon Java, module `paimon-vector`, commit `7234e4c34`. -//! * Generator: `PkVectorFixtureGenerator`. -//! * Command: `mvn -pl paimon-vector test -Dtest=PkVectorFixtureGenerator \ -//! -Dgen.pkvector.fixture=true -Drun.e2e.tests=true`. -//! * Config: primary key `id`, vector column `embedding`, `ivf-flat`, -//! `nlist = 1` (exact, deterministic single inverted list), `deduplicate` -//! merge engine, deletion-vectors enabled. -//! * Rows: `id == row position`, vectors `[0,0] [1,0] [2,0] [3,0] [4,0]`. -//! * Query `[0, 0]`, squared-L2 distances `[0, 1, 4, 9, 16]`; top-3 -> ids -//! `[0, 1, 2]`, distances `[0, 1, 4]`, scores `1/(1+d) = [1.0, 0.5, 0.2]`. -//! * Fixture tree checksum: `f6c21a447fa7be880713c3d1c27791e7dcb1db10`. +//! `(file, position)`. Search returns those positions with scores; a projected +//! result read must retrieve the same rows without inventing global row IDs. use std::path::Path; @@ -182,23 +165,27 @@ async fn reads_back_java_written_pk_vector_table() { "fixture top-3 ids must be [0, 1, 2]" ); - // execute_scored() reports global row ids. On a Java-written PK table the data - // files carry no `first_row_id`, so a global row id is unrecoverable and the - // scored path MUST fail loudly rather than fabricate ids. + // Java PK files need no first_row_id: the result retains physical positions. let scored = table .new_vector_search_builder() .with_vector_column(VECTOR_COLUMN) .with_query_vector(query.clone()) .with_limit(k) - .execute_scored() - .await; - assert!( - scored.is_err(), - "execute_scored() must fail on a primary-key vector table: global row ids \ - are unavailable when the data files carry no first_row_id" + .execute() + .await + .unwrap(); + assert!(scored.row_ids().is_err()); + assert_eq!( + scored + .positions() + .unwrap() + .iter() + .map(|p| p.row_position as i32) + .collect::>(), + expected_ids ); - // execute_read() materializes rows by physical position, so it MUST succeed + // the result reader materializes rows by physical position, so it MUST succeed // and emit the top-k best-first. The `id` column cross-checks the // position->id mapping (the fixture pins them equal) and `__paimon_search_score` // carries the metric score. @@ -206,15 +193,21 @@ async fn reads_back_java_written_pk_vector_table() { builder .with_vector_column(VECTOR_COLUMN) .with_query_vector(query) - .with_limit(k) - .with_projection(&["id"]); - let batches = builder - .execute_read() - .await - .expect("primary-key vector read over the Java fixture failed") - .try_collect::>() - .await - .expect("collecting read batches failed"); + .with_limit(k); + let batches = async { + builder + .execute() + .await? + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + } + .await + .expect("primary-key vector read over the Java fixture failed") + .try_collect::>() + .await + .expect("collecting read batches failed"); let ids = batch_i32(&batches, "id"); assert_eq!( diff --git a/docs/src/c-binding.md b/docs/src/c-binding.md index 1cd25ad94..8b2163364 100644 --- a/docs/src/c-binding.md +++ b/docs/src/c-binding.md @@ -332,6 +332,318 @@ and ranges. Combine predicates with `paimon_predicate_and`, A predicate that has not been consumed must be released with `paimon_predicate_free`. +## Vector Scan, Plan, and Read + +DE and primary-key vector searches use the same execution API: + +1. Configure a `paimon_vector_search_builder` with the column, query, limit, + options, predicate, and output projection. +2. Call `paimon_vector_search_builder_new_scan` and + `paimon_vector_search_builder_new_read` to create independent owned handles. + Creating a scan only requires the column; query validation happens when + creating the reader. +3. Call `paimon_vector_scan_plan` to resolve the source snapshot and search work. +4. Pass the reader and plan to `paimon_vector_read_read`. It returns the usual + Arrow record-batch reader with projected columns and `__paimon_search_score`. + +`paimon_vector_search_builder_execute_read` remains the convenience operation +for local planning and reading. A plan can also be reused with different query +vectors or limits. Its table, column, and pre-filter must match the reader. + +For Java-planned PK bucket work, decode each standalone +`BucketVectorSearchSplit.serialize` buffer using +`paimon_bucket_vector_search_split_deserialize`, then pass the decoded handles +to `paimon_vector_scan_plan_from_bucket_splits`. The returned common vector +plan is consumed by the same `paimon_vector_read_read` API. No table snapshot +or index manifest is read during this plan construction; supplied files, row +ranges, and snapshot IDs remain authoritative. Top-K is local to the supplied +buckets, so a distributed caller merges its per-worker results. + +The decoder accepts the versioned `PKVSPLIT` format. It does not accept Java +`ObjectOutputStream` envelopes or DE `IndexVectorSearchSplit` / +`RawVectorSearchSplit` object serialization. DE plans are currently obtained +through `paimon_vector_scan_plan`. + +| Owned handle | Release function | +|--------------|------------------| +| `paimon_vector_scan` | `paimon_vector_scan_free` | +| `paimon_vector_read` | `paimon_vector_read_free` | +| `paimon_vector_plan` | `paimon_vector_plan_free` | +| `paimon_bucket_vector_search_split` | `paimon_bucket_vector_search_split_free` | + +Decoded splits own their data, so input bytes can be released after decoding. +Plan construction copies the split metadata and leaves input handles intact, +including on failure. Free split handles after constructing the plan. Scans and +readers can outlive their builder; plans can outlive their scan. A read borrows +its plan, and the returned Arrow stream can outlive both the reader and plan. + +### Java-planned PK Bucket Splits + +Java plans the search once, and the caller sends each worker its assigned +`BucketVectorSearchSplit` buffers. The native worker follows this flow: + +```text +Java VectorScan.Plan + -> BucketVectorSearchSplit.serialize(DataOutputView), one buffer per split + -> Application transport + -> paimon_bucket_vector_search_split_deserialize, one handle per buffer + -> paimon_vector_scan_plan_from_bucket_splits + -> paimon_vector_read_read + -> paimon_record_batch_reader_next + -> Arrow consumer and global Top-K merge +``` + +`paimon_vector_search_builder_new_scan` creates the scan configuration; it does +not scan storage. For this path, construct the plan with +`paimon_vector_scan_plan_from_bucket_splits`. Calling `paimon_vector_scan_plan` +or `paimon_vector_search_builder_execute_read` would plan from the table again +and would not use the worker's assigned Java splits. + +#### Serialize on the Java Side + +Use the standalone serializer directly with `DataOutputViewStreamWrapper` over +a byte buffer. A Java object stream adds an envelope that the C decoder does +not accept. The following helper accepts a builder already configured with the +vector column and any planning predicates, for example one obtained from +`table.newVectorSearchBuilder().withVectorColumn("embedding")`: + +```java +import org.apache.paimon.io.DataOutputViewStreamWrapper; +import org.apache.paimon.table.source.BucketVectorSearchSplit; +import org.apache.paimon.table.source.VectorScan; +import org.apache.paimon.table.source.VectorSearchBuilder; +import org.apache.paimon.table.source.VectorSearchSplit; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +public final class VectorSplitSerializer { + public static List planAndSerialize(VectorSearchBuilder builder) + throws IOException { + VectorScan.Plan plan = builder.newVectorScan().scan(); + List buffers = new ArrayList<>(); + for (VectorSearchSplit split : plan.splits()) { + if (!(split instanceof BucketVectorSearchSplit)) { + throw new IllegalArgumentException("Expected a PK bucket vector split"); + } + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(bytes); + ((BucketVectorSearchSplit) split).serialize(out); + out.flush(); + buffers.add(bytes.toByteArray()); + } + return buffers; // Assign whole buffers to workers. + } +} +``` + +Preserve each buffer's length through transport. Decode one complete buffer at +a time; do not concatenate splits into a single decoder input. The buffer starts +with the eight bytes `PKVSPLIT`, followed by the big-endian format version +(currently `1`). Use matching Java and Rust format versions. + +Send the table location, branch and resolved Paimon `TableSchema` JSON alongside +the assigned buffers. The schema must retain its field IDs, primary keys and +table options, including vector index type, dimension, metric and deletion-vector +settings. This is Paimon schema JSON (`TableSchema.toString()`), not Arrow schema +JSON. Supply storage credentials/options separately when constructing the native +table; they are not merged into the table schema. A worker can use +`paimon_table_from_schema_json` without opening a catalog, or +`paimon_table_from_schema_json_with_file_io` with its own cache-enabled FileIO. +It must be able to access the data, deletion and index files named by the splits. + +For example, use the received metadata and worker-local storage options to +create the table (check `opened.error` before using `opened.table`): + +```c +paimon_result_get_table opened = paimon_table_from_schema_json( + table_path, table_schema_json, database, table_name, branch, + storage_options, storage_options_len); +``` + +Pass `opened.table` to the helper below and release it with `paimon_table_free` +after use. Pass `NULL, 0` for storage options when none are needed, and `NULL` +for `branch` only when the Java planner used the default `main` branch. + +The split buffers carry planned work, not the query vector, Top-K limit, +projection, query options or an executable scalar predicate. Send these query +parameters separately. If a scalar pre-filter is required, reconstruct it with +the `paimon_predicate_*` APIs and attach it to the native builder before creating +both scan and reader. Java file pruning and row ranges do not necessarily encode +the entire residual predicate. Applying that residual only after native Top-K +can discard winners without retrieving the next matching rows. PK data predicates +require deletion vectors enabled and merge-on-read disabled. + +#### Read the Assigned Splits through C + +This C11 helper searches `embedding`, projects `id`, and passes each batch to a +caller-provided callback. Adjust the column names to the table schema. The caller +provides a live table handle, a non-empty query of the configured dimension, a +positive `top_k`, and its assigned buffers as `paimon_byte_slice` values. + +The helper borrows the table and buffers and consumes the optional `filter`, +including on failure. The callback returns zero on success. It must either use +the batch synchronously or import/move its Arrow contents, marking the source +structures released according to the Arrow C Data Interface. It must not free +the Paimon batch container itself; the helper does that after the callback, +including when the callback fails. + +```c +#include +#include +#include + +#include "paimon.h" + +typedef int (*vector_batch_consumer)(void *context, paimon_arrow_batch batch); + +#define VECTOR_TRY(expression) \ + do { \ + error = (expression); \ + if (error != NULL) goto cleanup; \ + } while (0) + +int read_vector_splits( + const paimon_table *table, + const paimon_byte_slice *wire_splits, size_t split_count, + const float *query, size_t dimension, size_t top_k, + paimon_predicate *filter, + vector_batch_consumer consume_batch, void *context) { + int status = -1; + paimon_error *error = NULL; + paimon_bucket_vector_search_split **splits = NULL; + paimon_vector_search_builder *builder = NULL; + paimon_vector_scan *scan = NULL; + paimon_vector_plan *plan = NULL; + paimon_vector_read *read = NULL; + paimon_record_batch_reader *reader = NULL; + const char *projection[] = {"id", NULL}; + + if (table == NULL || consume_batch == NULL) goto cleanup; + if (split_count == 0) { + status = 0; // No assigned work; the plan API requires non-empty input. + goto cleanup; + } + if (wire_splits == NULL) goto cleanup; + splits = calloc(split_count, sizeof(*splits)); + if (splits == NULL) goto cleanup; + + for (size_t i = 0; i < split_count; ++i) { + paimon_result_bucket_vector_search_split decoded = + paimon_bucket_vector_search_split_deserialize( + wire_splits[i].data, wire_splits[i].len); + splits[i] = decoded.split; + VECTOR_TRY(decoded.error); + } + // All metadata is now owned by split handles; wire buffers can be released. + + paimon_result_vector_search_builder built = + paimon_table_new_vector_search_builder(table); + builder = built.builder; + VECTOR_TRY(built.error); + VECTOR_TRY(paimon_vector_search_builder_with_vector_column(builder, "embedding")); + VECTOR_TRY(paimon_vector_search_builder_with_query_vector(builder, query, dimension)); + VECTOR_TRY(paimon_vector_search_builder_with_limit(builder, top_k)); + VECTOR_TRY(paimon_vector_search_builder_with_projection(builder, projection)); + VECTOR_TRY(paimon_vector_search_builder_with_filter(builder, filter)); + filter = NULL; // Ownership transferred to the builder. + // Set paimon_vector_search_builder_with_options here if the query needs it. + + paimon_result_vector_scan scanned = paimon_vector_search_builder_new_scan(builder); + scan = scanned.scan; + VECTOR_TRY(scanned.error); + paimon_result_vector_read reading = paimon_vector_search_builder_new_read(builder); + read = reading.read; + VECTOR_TRY(reading.error); + + paimon_result_vector_plan planned = paimon_vector_scan_plan_from_bucket_splits( + scan, (const paimon_bucket_vector_search_split *const *)splits, split_count); + plan = planned.plan; + VECTOR_TRY(planned.error); + + // Plan construction copied the metadata and did not consume the handles. + for (size_t i = 0; i < split_count; ++i) { + paimon_bucket_vector_search_split_free(splits[i]); + splits[i] = NULL; + } + paimon_vector_scan_free(scan); + scan = NULL; + paimon_vector_search_builder_free(builder); + builder = NULL; + + paimon_result_record_batch_reader searched = paimon_vector_read_read(read, plan); + reader = searched.reader; + VECTOR_TRY(searched.error); + // The returned stream owns what it needs, independently of these handles. + paimon_vector_read_free(read); + read = NULL; + paimon_vector_plan_free(plan); + plan = NULL; + + for (;;) { + paimon_result_next_batch next = paimon_record_batch_reader_next(reader); + VECTOR_TRY(next.error); + if (next.batch.array == NULL && next.batch.schema == NULL) break; + int consumed = consume_batch(context, next.batch); + paimon_arrow_batch_free(next.batch); + if (consumed != 0) goto cleanup; + } + status = 0; + +cleanup: + if (error != NULL) { + fprintf(stderr, "Paimon error %d: ", error->code); + fwrite(error->message.data, 1, error->message.len, stderr); + fputc('\n', stderr); + paimon_error_free(error); + } + paimon_record_batch_reader_free(reader); + paimon_vector_read_free(read); + paimon_vector_plan_free(plan); + paimon_vector_scan_free(scan); + paimon_vector_search_builder_free(builder); + paimon_predicate_free(filter); + if (splits != NULL) { + for (size_t i = 0; i < split_count; ++i) { + paimon_bucket_vector_search_split_free(splits[i]); + } + free(splits); + } + return status; +} + +#undef VECTOR_TRY +``` + +Compile the helper as C and declare it with C linkage in the native worker. +When including the generated C header directly from C++, wrap the include +in `extern "C" { ... }` or generate a C++-compatible C header with +`cbindgen bindings/c --lang c --cpp-compat --output target/release/paimon.h`. +Replace the sample stderr reporting with the application's error handling as needed. + +Each call returns up to `top_k` rows across all splits supplied to that call, +in relevance order, with the requested user columns and a `FLOAT32` +`__paimon_search_score` column. Higher scores rank first, including for L2 (the +score is `1 / (1 + squared_distance)`, not the raw distance). The caller must merge +the results from its disjoint assignments and apply the final global Top-K. +Include the primary-key columns in the projection if the coordinator needs +them for row identity. This merge retains the configured search mode's ANN/exact +semantics; it does not make ANN search exact. + +All splits in a plan must come from one table, branch and snapshot, with at most +one split per `(partition, bucket)`. Mixed snapshot IDs and repeated buckets are +rejected. The caller is responsible for keeping the table/branch metadata paired +with the buffers and for avoiding duplicate assignments across workers. An empty +Java plan means no work; skip native plan construction. Supplied file-local row +ranges remain authoritative and are intersected with any native residual filter. + +For repeated queries over the same assignment, retain `paimon_vector_plan` and +create another reader with the new query/limit instead of decoding again. Readers +must use the same table, branch, vector column and pre-filter as the plan. +Changing a builder after `new_read` does not change an already-created reader. + ## Writing and Committing Writing uses a **write-then-commit** flow: From 024fabb29614699c1106ae91853b2a540d02b621 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 13 Sep 2026 22:23:29 +0800 Subject: [PATCH 2/3] refactor(table): align PK vector row ranges with Java --- .../src/table/pk_vector_orchestrator.rs | 46 ++-- crates/paimon/src/table/pk_vector_read.rs | 216 +++++------------- ..._tests.rs => residual_row_ranges_tests.rs} | 136 +++++++---- crates/paimon/src/table/pk_vector_scan.rs | 17 +- crates/paimon/src/vindex/pkvector/ann.rs | 201 +++++----------- crates/paimon/src/vindex/pkvector/bucket.rs | 83 +++---- crates/paimon/src/vindex/pkvector/mod.rs | 83 ++----- 7 files changed, 295 insertions(+), 487 deletions(-) rename crates/paimon/src/table/pk_vector_read/{residual_positions_tests.rs => residual_row_ranges_tests.rs} (81%) diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 0d8c68c38..5e37679b8 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -38,7 +38,7 @@ use crate::vindex::pkvector::bucket::{ }; use crate::vindex::pkvector::metric::{java_float_compare, VectorSearchMetric}; use crate::vindex::pkvector::result::PkVectorSearchResult; -use crate::vindex::pkvector::FileRowSelections; +use crate::vindex::pkvector::RowRangesByFile; fn data_invalid(message: impl Into) -> crate::Error { crate::Error::DataInvalid { @@ -356,13 +356,12 @@ impl PkVectorOrchestrator { /// and split so a caller can build a reader keyed to the specific split/file. /// `skip_exact_fallback` forwards to `bucket_search`. /// - /// `row_selections_by_split`, when present, carries one per-file row selection - /// per split (indexed parallel to `splits`), in the three states of - /// [`FileRowSelection`]: a file with NO entry is unrestricted, an empty entry - /// contributes no candidates, and a non-empty one limits which of its rows may. - /// A selection is either interval `Ranges` (from an engine's bucket split) or - /// `Positions` (from a residual data predicate). `None` restricts nothing at - /// all. The slice must have the same length as `splits`. + /// `row_ranges_by_split`, when present, carries one per-file row selection + /// per split (indexed parallel to `splits`), using [`RowRangesByFile`]: a file + /// with no entry is unrestricted, an empty list contributes no candidates, + /// and a non-empty list limits its rows. Both plans and residual predicates + /// use merged intervals. `None` restricts nothing; the slice must have the + /// same length as `splits`. /// /// This is the single-query wrapper over /// [`search_candidates_batch`](Self::search_candidates_batch): it searches the @@ -392,7 +391,7 @@ impl PkVectorOrchestrator { + Sync), search_options: &HashMap, skip_exact_fallback: bool, - row_selections_by_split: Option<&[FileRowSelections]>, + row_ranges_by_split: Option<&[RowRangesByFile]>, concurrency: usize, ) -> crate::Result { let mut results = self @@ -406,7 +405,7 @@ impl PkVectorOrchestrator { exact_file_search, search_options, skip_exact_fallback, - row_selections_by_split, + row_ranges_by_split, concurrency, ) .await?; @@ -424,7 +423,7 @@ impl PkVectorOrchestrator { /// into another's (independent per-query heaps). /// /// The row selections depend only on the filter and the plan, not the - /// query vector, so the SAME `row_selections_by_split` slice is shared across every + /// query vector, so the SAME `row_ranges_by_split` slice is shared across every /// query. Input-shape validation (positive limits, non-empty query, residual /// count) is applied per query / once as appropriate. /// @@ -458,7 +457,7 @@ impl PkVectorOrchestrator { + Sync), search_options: &HashMap, skip_exact_fallback: bool, - row_selections_by_split: Option<&[FileRowSelections]>, + row_ranges_by_split: Option<&[RowRangesByFile]>, concurrency: usize, ) -> crate::Result> { // Eager input-shape validation (Java checkArgument parity). @@ -476,7 +475,7 @@ impl PkVectorOrchestrator { return Err(data_invalid("vector search query must not be empty")); } } - if let Some(per_split) = row_selections_by_split { + if let Some(per_split) = row_ranges_by_split { if per_split.len() != splits.len() { return Err(data_invalid( "row selection map count does not match split count", @@ -529,8 +528,8 @@ impl PkVectorOrchestrator { ) }, ); - let row_selections = - row_selections_by_split.map(|per_split| &per_split[split_index]); + let row_ranges_by_file = + row_ranges_by_split.map(|per_split| &per_split[split_index]); let per_query = bucket_search_batch( ann_searcher, &split.ann_segments, @@ -543,7 +542,7 @@ impl PkVectorOrchestrator { limit, search_options, skip_exact_fallback, - row_selections, + row_ranges_by_file, concurrency, search_budget, ) @@ -1268,7 +1267,7 @@ mod e2e_tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { Ok(queries.iter().map(|_| self.hits.clone()).collect()) } @@ -1800,12 +1799,9 @@ mod e2e_tests { ); // Allow only positions 0 and 2 for "r.mosaic"; pos1 (the best hit) is // excluded by the residual. - let mut allowed = roaring::RoaringTreemap::new(); - allowed.insert(0); - allowed.insert(2); - let row_selections_by_split: Vec = vec![HashMap::from([( + let row_ranges_by_split: Vec = vec![HashMap::from([( "r.mosaic".to_string(), - crate::vindex::pkvector::FileRowSelection::Positions(allowed), + vec![RowRange::new(0, 0), RowRange::new(2, 2)], )])]; let opts = HashMap::new(); let result = PkVectorOrchestrator::new(make_reader(file_io, table_path)) @@ -1819,7 +1815,7 @@ mod e2e_tests { &factory, &opts, false, - Some(&row_selections_by_split), + Some(&row_ranges_by_split), 1, ) .await @@ -1860,7 +1856,7 @@ mod e2e_tests { }; let factory = unreachable_split_search(); // Two residual maps for a single split. - let row_selections_by_split: Vec = vec![HashMap::new(), HashMap::new()]; + let row_ranges_by_split: Vec = vec![HashMap::new(), HashMap::new()]; let opts = HashMap::new(); let err = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( @@ -1873,7 +1869,7 @@ mod e2e_tests { &factory, &opts, false, - Some(&row_selections_by_split), + Some(&row_ranges_by_split), 1, ) .await diff --git a/crates/paimon/src/table/pk_vector_read.rs b/crates/paimon/src/table/pk_vector_read.rs index 5869e9427..8a4ab2048 100644 --- a/crates/paimon/src/table/pk_vector_read.rs +++ b/crates/paimon/src/table/pk_vector_read.rs @@ -35,6 +35,7 @@ use crate::table::pk_vector_orchestrator::{ use crate::table::pk_vector_position_read::{PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN}; use crate::table::pk_vector_scan::PkVectorScanPlan; use crate::table::pk_vector_search_params::PkVectorSearchParams; +use crate::table::row_id_predicate::intersect_sorted_ranges; use crate::table::source::DataSplit; use crate::table::vector_read::Read; use crate::table::vector_search_common::{ @@ -46,12 +47,11 @@ use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; use crate::vindex::pkvector::ann::{AnnSegmentSource, PkVectorAnnSearcher, VindexAnnSearcher}; use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture}; use crate::vindex::pkvector::metric::VectorSearchMetric; -use crate::vindex::pkvector::{FileRowSelection, FileRowSelections}; +use crate::vindex::pkvector::RowRangesByFile; use crate::vindex::range_reader::{RangeReadLimiter, VindexFileReader}; use crate::vindex::reader::VindexVectorGlobalIndexReader; use arrow_array::{Array, Int64Array, RecordBatch}; use futures::{stream, TryStreamExt}; -use roaring::RoaringTreemap; use std::collections::{HashMap, HashSet}; use std::io::Cursor; use std::sync::Arc; @@ -183,41 +183,16 @@ pub(super) async fn materialize_positions( Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) } -/// Search an already-resolved plan across every query and return each query's raw -/// indexed and exact candidate lists, before any rerank or merge. -/// -/// Plan-dependent concurrency — the vindex segment count, batch-index parallelism -/// and the range-read bound — is derived here from the plan that is actually being -/// searched, so a narrowed plan can never be searched under limits computed for a -/// wider one. -/// Combine the two per-split row allow-lists a search can be handed: the physical -/// rows an engine-supplied plan restricts each file to, and the positions a residual -/// data predicate leaves behind. -/// -/// The two sides read a file's ABSENCE differently, and the merge has to respect -/// both readings: -/// -/// * The plan lists only what the engine's split narrowed, so an absent file is -/// unrestricted -- Java's `rowRangesByFile.get(file) == null`. -/// * The residual is exhaustive over the files a search can read from -/// (`residual_positions_by_file` registers every active file, empty when nothing -/// passed), so once a residual exists its silence about a file means "no rows". -/// -/// So: with no residual, a file the plan omits stays absent and unrestricted. With a -/// residual, a file it omits is excluded even if the plan restricted it, and a file -/// both describe keeps the intersection. Absent from BOTH is unrestricted, which is -/// what lets the ANN backend search unfiltered. -/// -/// The plan's ranges stay ranges. Expanding them into positions would be work sized -/// by row counts that arrived on the wire; where an intersection is genuinely needed -/// the residual positions — bounded by the rows its own read returned — are filtered -/// BY the ranges instead. When the residual was evaluated over those same ranges the -/// intersection cannot remove anything, and is kept as the invariant that says so. -fn intersect_row_allow_lists( - physical: Option<&[HashMap>]>, - residual: Option>>, +/// Intersect the plan's physical row ranges with the residual predicate's ranges. +/// The plan omits unrestricted files; the residual registers every active file, +/// including empty results. A file listed only by the plan must therefore stay +/// excluded when a residual exists. Files absent from both inputs stay absent. +/// Both inputs are sorted and merged, so intersection never expands large spans. +fn intersect_row_ranges_by_split( + physical: Option<&[RowRangesByFile]>, + residual: Option>, split_count: usize, -) -> crate::Result>> { +) -> crate::Result>> { if let Some(maps) = physical { if maps.len() != split_count { return Err(crate::Error::DataInvalid { @@ -241,72 +216,28 @@ fn intersect_row_allow_lists( } } match (physical, residual) { - (None, None) => Ok(None), - (None, Some(residual)) => Ok(Some( - residual - .into_iter() - .map(|per_file| { - per_file - .into_iter() - .map(|(file, positions)| (file, FileRowSelection::Positions(positions))) - .collect() - }) - .collect(), - )), - (Some(physical), None) => Ok(Some( + (None, residual) => Ok(residual), + (Some(physical), None) => Ok(Some(physical.to_vec())), + (Some(physical), Some(residual)) => Ok(Some( physical .iter() - .map(|per_file| { - per_file - .iter() - .map(|(file, ranges)| { - (file.clone(), FileRowSelection::Ranges(ranges.clone())) - }) - .collect() + .zip(residual) + .map(|(physical, mut residual)| { + for (file, ranges) in physical { + // The residual covers every active file. A missing entry + // must not restore rows excluded by that residual. + let allowed = residual.entry(file.clone()).or_default(); + *allowed = intersect_sorted_ranges(ranges, allowed); + } + residual }) .collect(), )), - (Some(physical), Some(residual)) => { - Ok(Some( - physical - .iter() - .zip(residual) - .map(|(physical, mut residual)| { - let mut merged: FileRowSelections = HashMap::new(); - for (file, ranges) in physical { - let range_selection = FileRowSelection::Ranges(ranges.clone()); - let selection = match residual.remove(file.as_str()) { - // Both restrict: keep the positions the ranges also - // allow. Filtering the positions (bounded by the read) - // by the ranges never expands the ranges. - Some(positions) => FileRowSelection::Positions( - positions - .iter() - .filter(|position| range_selection.contains(*position)) - .collect(), - ), - // The residual is exhaustive over the files the search - // can read from -- `residual_positions_by_file` - // registers every active file, empty when nothing - // passed. Its silence about a file therefore means "no - // rows", NOT "unrestricted", and must stay fail-closed - // here even though the plan has something to say. - None => FileRowSelection::Positions(RoaringTreemap::new()), - }; - merged.insert(file.clone(), selection); - } - // Whatever the residual restricted and the plan did not. - merged.extend(residual.into_iter().map(|(file, positions)| { - (file, FileRowSelection::Positions(positions)) - })); - merged - }) - .collect(), - )) - } } } +/// Search a resolved plan across all queries before reranking and merging. +/// Concurrency is derived from the actual plan, including external split subsets. #[allow(clippy::too_many_arguments)] async fn search_pk_raw_candidates_batch_with_plan( table: &Table, @@ -481,21 +412,11 @@ async fn search_pk_raw_candidates_batch_with_plan( field_name, scorer, loader, )); - // Residual (post-recall) filtering: for each candidate file, re-read its - // physical rows and keep the positions whose rows satisfy the filter. The - // per-split allow-list is threaded into the bucket search so the residual folds - // into recall (best-first order and Top-K are preserved). Built only when the - // filter has data (non-partition) conjuncts; a partition-only filter (or no - // filter) leaves `None`, which leaves the search unfiltered — partition - // pruning is already handled in planning. The residual depends only on the - // filter and the plan, not the query vector, so it is computed once here and - // shared across every query in the batch. The residual reader projects only - // the predicate columns and carries no pushdown; `residual_positions_by_file` - // recovers each surviving row's file-local physical position from its ordinal - // in the unfiltered scan (no `_ROW_ID`, no `first_row_id`). A file the - // allow-list leaves empty is skipped by the bucket search without opening an - // exact reader. - let residual_by_split: Option>> = match filter { + // Resolve data predicates before recall so both ANN and exact Top-K honor + // them. Partition-only predicates were already applied by the scan. The + // residual depends on the filter and plan, so its file-local ranges are + // shared by all queries. A file with an empty range list is skipped. + let residual_by_split: Option> = match filter { Some(filter) => { // The whole filter is pushed into scan planning (`PkVectorScan`), where // partition-only conjuncts already prune partitions/files. Re-applying @@ -537,7 +458,7 @@ async fn search_pk_raw_candidates_batch_with_plan( .as_ref() .and_then(|per_split| per_split.get(index)); per_split.push( - residual_positions_by_file( + residual_row_ranges_by_file( &residual_reader, &split.data_split, &split.active_files, @@ -552,11 +473,8 @@ async fn search_pk_raw_candidates_batch_with_plan( } None => None, }; - // Fold the plan's own positional restriction into the same allow-list. A plan - // built from engine-supplied bucket splits carries the physical positions each - // file is limited to; a plan read from the index manifest carries none. Both - // sides list what is permitted, so combining them is an intersection. - let row_selections_by_split = intersect_row_allow_lists( + // Preserve the external plan's ranges as well as the residual restriction. + let row_ranges_by_split = intersect_row_ranges_by_split( plan.physical_row_ranges_by_split.as_deref(), residual_by_split, plan.splits.len(), @@ -634,7 +552,7 @@ async fn search_pk_raw_candidates_batch_with_plan( &factory, &search_options, skip_exact_fallback, - row_selections_by_split.as_deref(), + row_ranges_by_split.as_deref(), concurrency, ) .await?; @@ -713,47 +631,30 @@ async fn search_pk_candidates_batch_with_plan( Ok(per_query_candidates) } -/// Compute, per data file in `split`, the set of file-LOCAL physical row -/// positions whose rows satisfy the residual predicate. Mirrors the -/// row-collecting half of Java `PrimaryKeyVectorRead`'s `executeFilter`: the -/// predicate is NOT pushed down (a pushed filter would drop rows before their -/// position could be recovered). Instead `reader` projects only the residual -/// columns and carries no pushdown predicate, the residual is evaluated here at the -/// Arrow level, and each surviving row's file-local 0-based position is recovered -/// from the selection the read was limited to. This needs no `_ROW_ID` and no -/// `first_row_id` — real primary-key tables never write one. +/// Read the plan's allowed physical rows and return merged ranges matching the +/// residual, as in Java `PrimaryKeyVectorRead.residualRowRanges`. /// -/// `allowed_rows` is the plan's per-file physical selection, keyed by data-file -/// name, with the plan's three states: a file it does not list is unrestricted and -/// the whole file is scanned; an empty range list excludes the file, which is -/// registered empty without a read; a non-empty list is scanned over exactly those -/// ranges, because an engine-supplied bucket split can restrict a huge file to a -/// handful of ranges and reading all of it to discard the rest would defeat the -/// split. +/// The reader projects predicate columns without pushing the predicate down, so +/// each emitted row can be mapped back to its physical position. This uses neither +/// `_ROW_ID` nor `first_row_id`. Ascending matches are coalesced as they arrive. /// -/// Every *active* data file in the split gets an entry in the RESULT, possibly -/// empty, and that exhaustiveness is load-bearing. The search kernel reads a file's -/// absence from its selections as "unrestricted", so an active file missing here -/// would reach the search with no predicate applied at all -- the residual would be -/// silently dropped for it. (The merge below reads a residual's silence about a -/// file the PLAN listed as exclusion, so only a file both omit falls through, which -/// is exactly the case this exhaustiveness rules out.) Non-active files (e.g. -/// level-0 files the bucket search excludes) are skipped entirely: they are never -/// searched, so re-reading them would be wasted IO. +/// Missing `allowed_rows` entries permit the entire file; empty ranges skip it +/// without a read. Every active file gets an output entry, including empty results, +/// so a rejected file cannot become unrestricted in the search kernel. Inactive +/// files are not searched and need no residual read. /// -/// `reader` must be predicate-free and project the residual columns; -/// `residual.file_fields` are the fields the residual leaf indices point into -/// (resolved by name against each emitted batch). -async fn residual_positions_by_file( +/// `reader` must be predicate-free; `residual.file_fields` resolves predicate +/// indices against each emitted batch by name. +async fn residual_row_ranges_by_file( reader: &DataFileReader, split: &DataSplit, active_files: &[BucketActiveFile], residual: &FilePredicates, - allowed_rows: Option<&HashMap>>, -) -> crate::Result> { + allowed_rows: Option<&RowRangesByFile>, +) -> crate::Result { let scan_fields = reader.read_type().to_vec(); let active_names: HashSet<&str> = active_files.iter().map(|f| f.file_name.as_str()).collect(); - let mut out: HashMap = HashMap::new(); + let mut out: RowRangesByFile = HashMap::new(); for file_meta in split.data_files() { // Only files the bucket search actually recalls from need residual // positions; skip everything else to avoid a wasted read. @@ -786,18 +687,18 @@ async fn residual_positions_by_file( }; // Register the file up front so a file whose rows all fail the residual // still appears in the map (empty set). - let positions = out.entry(file_meta.file_name.clone()).or_default(); + let ranges = out.entry(file_meta.file_name.clone()).or_default(); // Rows arrive in ascending physical order, and the read emitted exactly what // was selected (no pushdown predicate, no deletion vector), so walking the // selection in step with the rows recovers each row's file-local position. - let mut selected: Box + Send> = match &selection { + let mut selected: Box + Send> = match &selection { Some(ranges) => Box::new( ranges .clone() .into_iter() - .flat_map(|range| (range.from() as u64)..=(range.to() as u64)), + .flat_map(|range| range.from()..=range.to()), ), - None => Box::new(0..file_meta.row_count.max(0) as u64), + None => Box::new(0..file_meta.row_count.max(0)), }; while let Some(batch) = stream.try_next().await? { let num_rows = batch.num_rows(); @@ -823,7 +724,14 @@ async fn residual_positions_by_file( None => true, }; if keep { - positions.insert(position); + // Reads return ascending physical positions, so coalesce + // consecutive matches directly into Java's range form. + match ranges.last_mut() { + Some(last) if last.to().checked_add(1) == Some(position) => { + *last = RowRange::new(last.from(), position); + } + _ => ranges.push(RowRange::new(position, position)), + } } } } @@ -1023,4 +931,4 @@ async fn rerank_indexed_positional( mod tests; #[cfg(test)] -mod residual_positions_tests; +mod residual_row_ranges_tests; diff --git a/crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs b/crates/paimon/src/table/pk_vector_read/residual_row_ranges_tests.rs similarity index 81% rename from crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs rename to crates/paimon/src/table/pk_vector_read/residual_row_ranges_tests.rs index 7dc65762b..f50ef9d0b 100644 --- a/crates/paimon/src/table/pk_vector_read/residual_positions_tests.rs +++ b/crates/paimon/src/table/pk_vector_read/residual_row_ranges_tests.rs @@ -181,8 +181,48 @@ fn residual_id_gt(threshold: i32) -> FilePredicates { } } -fn sorted(t: &roaring::RoaringTreemap) -> Vec { - t.iter().collect() +fn sorted(ranges: &[RowRange]) -> Vec { + ranges + .iter() + .flat_map(|r| (r.from() as u64)..=(r.to() as u64)) + .collect() +} + +#[tokio::test] +async fn residual_ranges_coalesce_matches_without_bridging_gaps() { + let (reader, split, active) = build_reader_and_split( + "memory:/residual_range_gaps", + &[("part-0.mosaic", vec![7, 7, 0, 7, 7, 0, 7], 100)], + ) + .await; + let ranges = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(0), None) + .await + .unwrap(); + assert_eq!( + ranges["part-0.mosaic"], + vec![ + RowRange::new(0, 1), + RowRange::new(3, 4), + RowRange::new(6, 6) + ] + ); + + let allowed = HashMap::from([( + "part-0.mosaic".to_string(), + vec![RowRange::new(1, 3), RowRange::new(5, 6)], + )]); + let ranges = + residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(0), Some(&allowed)) + .await + .unwrap(); + assert_eq!( + ranges["part-0.mosaic"], + vec![ + RowRange::new(1, 1), + RowRange::new(3, 3), + RowRange::new(6, 6) + ] + ); } #[tokio::test] @@ -193,7 +233,7 @@ async fn test_residual_selects_matching_positions() { &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)], ) .await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), None) + let map = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(2), None) .await .unwrap(); assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]); @@ -215,7 +255,7 @@ async fn test_residual_only_evaluates_the_rows_the_plan_allows() { .await; let allowed = HashMap::from([("part-0.mosaic".to_string(), vec![RowRange::new(3, 4)])]); let map = - residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), Some(&allowed)) + residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(2), Some(&allowed)) .await .unwrap(); assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]); @@ -233,7 +273,7 @@ async fn test_residual_does_not_read_a_file_the_plan_excludes() { .await; let excluded = HashMap::from([("part-0.mosaic".to_string(), Vec::new())]); - let map = residual_positions_by_file( + let map = residual_row_ranges_by_file( &reader, &split, &active, @@ -246,7 +286,7 @@ async fn test_residual_does_not_read_a_file_the_plan_excludes() { assert!(sorted(&map["part-0.mosaic"]).is_empty()); let unrestricted = HashMap::new(); - let map = residual_positions_by_file( + let map = residual_row_ranges_by_file( &reader, &split, &active, @@ -263,7 +303,7 @@ async fn test_residual_matches_none_yields_empty_entry() { // id > 100 matches nothing; the file still gets a (present, empty) entry. let (reader, split, active) = build_reader_and_split("memory:/rpf_none", &[("part-0.mosaic", vec![1, 2, 3], 0)]).await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(100), None) + let map = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(100), None) .await .unwrap(); assert!(map.contains_key("part-0.mosaic")); @@ -274,7 +314,7 @@ async fn test_residual_matches_none_yields_empty_entry() { async fn test_residual_matches_all_yields_full_set() { let (reader, split, active) = build_reader_and_split("memory:/rpf_all", &[("part-0.mosaic", vec![1, 2, 3], 0)]).await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(0), None) + let map = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(0), None) .await .unwrap(); assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); @@ -292,7 +332,7 @@ async fn test_residual_positions_are_file_local_across_files() { ], ) .await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(3), None) + let map = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(3), None) .await .unwrap(); assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]); @@ -335,7 +375,7 @@ async fn test_non_active_files_are_skipped() { .with_data_files(metas) .build() .unwrap(); - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(2), None) + let map = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(2), None) .await .unwrap(); assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]); @@ -351,7 +391,7 @@ async fn test_missing_first_row_id_recovers_local_positions() { // recovered from each row's ordinal in the scan, so the residual still // works: ids [1,2,3] with id > 0 -> all match -> local positions [0,1,2]. let (reader, split, active) = build_reader_and_split_no_first_row_id().await; - let map = residual_positions_by_file(&reader, &split, &active, &residual_id_gt(0), None) + let map = residual_row_ranges_by_file(&reader, &split, &active, &residual_id_gt(0), None) .await .expect("missing first_row_id must not fail the residual read"); assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]); @@ -397,22 +437,13 @@ async fn build_reader_and_split_no_first_row_id( // ---- combining the plan's positional restriction with the residual ---- -fn allow_list(entries: &[(&str, &[u64])]) -> HashMap { - entries - .iter() - .map(|(file, positions)| ((*file).to_string(), positions.iter().copied().collect())) - .collect() -} - -/// The plan side carries ranges, so its fixtures are built from the positions -/// each file allows and coalesced the way the planner normalizes them. -fn range_allow_list(entries: &[(&str, &[u64])]) -> HashMap> { +fn range_allow_list(entries: &[(&str, &[u64])]) -> RowRangesByFile { entries .iter() .map(|(file, positions)| { let ranges = positions .iter() - .map(|p| RowRange::new(*p as i64, *p as i64)) + .map(|&p| RowRange::new(p as i64, p as i64)) .collect(); ((*file).to_string(), merge_row_ranges(ranges)) }) @@ -421,11 +452,10 @@ fn range_allow_list(entries: &[(&str, &[u64])]) -> HashMap /// The positions a merged selection allows, expanded for readable assertions. /// Test-only: the production path never expands a range. -fn listed(map: &FileRowSelections, file: &str) -> Vec { +fn listed(map: &RowRangesByFile, file: &str) -> Vec { match map.get(file) { None => Vec::new(), - Some(FileRowSelection::Positions(positions)) => positions.iter().collect(), - Some(FileRowSelection::Ranges(ranges)) => ranges + Some(ranges) => ranges .iter() .flat_map(|range| (range.from() as u64)..=(range.to() as u64)) .collect(), @@ -434,25 +464,41 @@ fn listed(map: &FileRowSelections, file: &str) -> Vec { #[test] fn no_restriction_on_either_side_stays_unrestricted() { - assert!(intersect_row_allow_lists(None, None, 1).unwrap().is_none()); + assert!(intersect_row_ranges_by_split(None, None, 1) + .unwrap() + .is_none()); +} + +#[test] +fn range_intersection_keeps_large_spans_compact_and_inclusive() { + let physical = vec![HashMap::from([( + "f".to_string(), + vec![RowRange::new(1, i64::MAX)], + )])]; + let residual = vec![HashMap::from([( + "f".to_string(), + vec![RowRange::new(0, 1), RowRange::new(i64::MAX - 1, i64::MAX)], + )])]; + let combined = intersect_row_ranges_by_split(Some(&physical), Some(residual), 1) + .unwrap() + .unwrap(); + assert_eq!( + combined[0]["f"], + vec![RowRange::new(1, 1), RowRange::new(i64::MAX - 1, i64::MAX)] + ); } #[test] fn one_side_alone_passes_through() { let physical = vec![range_allow_list(&[("d0", &[1, 2])])]; - let only_physical = intersect_row_allow_lists(Some(&physical), None, 1) + let only_physical = intersect_row_ranges_by_split(Some(&physical), None, 1) .unwrap() .expect("a plan restriction survives on its own"); assert_eq!(listed(&only_physical[0], "d0"), vec![1, 2]); - // Still intervals. Expanding them here is the unbounded step the plan side - // must never take, and the positions above cannot tell the two apart. - assert!( - matches!(only_physical[0]["d0"], FileRowSelection::Ranges(_)), - "the plan's ranges must reach the search as ranges" - ); + assert_eq!(only_physical[0]["d0"], vec![RowRange::new(1, 2)]); - let residual = vec![allow_list(&[("d0", &[3])])]; - let only_residual = intersect_row_allow_lists(None, Some(residual), 1) + let residual = vec![range_allow_list(&[("d0", &[3])])]; + let only_residual = intersect_row_ranges_by_split(None, Some(residual), 1) .unwrap() .expect("a residual survives on its own"); assert_eq!(listed(&only_residual[0], "d0"), vec![3]); @@ -466,13 +512,13 @@ fn both_sides_intersect_and_the_residual_stays_fail_closed() { // must not resurrect the file, and neither may its absence make it // unrestricted. let physical = vec![range_allow_list(&[("d0", &[1, 2, 3]), ("d1", &[0, 1])])]; - let residual = vec![allow_list(&[("d0", &[2, 3, 4])])]; - let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) + let residual = vec![range_allow_list(&[("d0", &[2, 3, 4])])]; + let combined = intersect_row_ranges_by_split(Some(&physical), Some(residual), 1) .unwrap() .expect("both sides restrict"); assert_eq!(listed(&combined[0], "d0"), vec![2, 3]); assert!( - combined[0]["d1"].is_excluded(), + combined[0]["d1"].is_empty(), "a file the residual omits must stay excluded" ); } @@ -483,13 +529,13 @@ fn a_file_neither_side_restricts_stays_absent() { // entry for a file no one narrowed, or the ANN backend takes the filtered // path for a query that filters nothing. let physical = vec![range_allow_list(&[("d0", &[1])])]; - let combined = intersect_row_allow_lists(Some(&physical), None, 1) + let combined = intersect_row_ranges_by_split(Some(&physical), None, 1) .unwrap() .expect("the plan restricts d0"); assert!(!combined[0].contains_key("d1")); - let residual = vec![allow_list(&[("d0", &[1])])]; - let combined = intersect_row_allow_lists(Some(&physical), Some(residual), 1) + let residual = vec![range_allow_list(&[("d0", &[1])])]; + let combined = intersect_row_ranges_by_split(Some(&physical), Some(residual), 1) .unwrap() .expect("both restrict d0"); assert!(!combined[0].contains_key("d1")); @@ -502,7 +548,7 @@ fn a_plan_that_restricts_nothing_produces_an_empty_selection_map() { // and that must survive the merge as an empty map (which the ANN layer reads // as "nothing to mask"), not become a per-file all-permitting mask. let physical = vec![HashMap::new()]; - let combined = intersect_row_allow_lists(Some(&physical), None, 1) + let combined = intersect_row_ranges_by_split(Some(&physical), None, 1) .unwrap() .expect("a split-driven plan is always Some"); assert!(combined[0].is_empty()); @@ -523,13 +569,13 @@ fn take_only_result_rejects_bad_batch_arity() { #[test] fn rejects_allow_lists_that_do_not_cover_every_split() { let physical = vec![range_allow_list(&[("d0", &[1])])]; - let error = intersect_row_allow_lists(Some(&physical), None, 2) + let error = intersect_row_ranges_by_split(Some(&physical), None, 2) .map(|_| ()) .expect_err("an allow-list per split is what makes the index meaningful"); assert!(error.to_string().contains("for 2 splits"), "{error}"); - let residual = vec![allow_list(&[("d0", &[1])])]; - let error = intersect_row_allow_lists(Some(&physical), Some(residual), 2) + let residual = vec![range_allow_list(&[("d0", &[1])])]; + let error = intersect_row_ranges_by_split(Some(&physical), Some(residual), 2) .map(|_| ()) .expect_err("the residual must cover every split too"); assert!(error.to_string().contains("for 2 splits"), "{error}"); diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index e9fae268e..9b845af7d 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -37,6 +37,7 @@ use crate::table::source::{merge_row_ranges, DataSplit, DataSplitBuilder, Deleti use crate::table::vector_scan::Scan; use crate::table::Table; use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; +use crate::vindex::pkvector::RowRangesByFile; /// A payload whose bucket-local path is resolved in planning Phase C, once the /// owning bucket's data split (and directory) is known. @@ -229,7 +230,7 @@ pub(crate) struct PkVectorScanPlan { // carry row ranges the engine's own planner already resolved -- possibly an // empty map, when that planner narrowed nothing. `None` for a plan read from // this table's index manifest, which places no positional restriction at all. - pub physical_row_ranges_by_split: Option>>>, + pub physical_row_ranges_by_split: Option>, } pub(crate) struct PkVectorScan { @@ -1417,18 +1418,10 @@ mod tests { let plan = plan_from_bucket_splits(&index_type, field_id, None, "/tbl", false, vec![split]) .unwrap(); - let selections: crate::vindex::pkvector::FileRowSelections = plan + let row_ranges_by_file = plan .physical_row_ranges_by_split .expect("split-driven plan") - .remove(0) - .into_iter() - .map(|(file, ranges)| { - ( - file, - crate::vindex::pkvector::FileRowSelection::Ranges(ranges), - ) - }) - .collect(); + .remove(0); let active: HashSet = source_meta .source_files() @@ -1440,7 +1433,7 @@ mod tests { source_meta.source_files(), &active, &HashMap::new(), - Some(&selections), + Some(&row_ranges_by_file), ) .unwrap() .is_none(), diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index a2baa220d..39de1a03a 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -25,7 +25,7 @@ use super::bucket::BucketAnnSegment; use super::data_invalid; use super::metric::{java_float_compare, VectorSearchMetric}; use super::result::PkVectorSearchResult; -use super::{FileRowSelection, FileRowSelections}; +use super::{contains_row_position, RowRangesByFile}; use crate::deletion_vector::DeletionVector; use crate::spec::{ PrimaryKeyIndexSourceFile as PkVectorSourceFile, @@ -77,9 +77,9 @@ fn charge_live_rows(remaining: &mut u64, rows: u64) -> crate::Result<()> { Ok(()) } -/// `row_selections` restricts each source file to the rows a pre-filter allows, +/// `row_ranges_by_file` restricts each source file to the rows a pre-filter allows, /// keyed by data-file name. A file with **no entry is unrestricted**, an empty -/// entry excludes it, and a non-empty one limits it — see [`FileRowSelection`]. +/// entry excludes it, and a non-empty one limits it — see [`RowRangesByFile`]. /// Mirrors Java `rowRangesByFile`. /// /// Returns `None` when nothing is restricted, every source file is active, AND no @@ -96,7 +96,7 @@ pub(crate) fn build_live_row_ids( source_files: &[PkVectorSourceFile], active_source_files: &HashSet, deletion_vectors: &HashMap>, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result> { let all_active = source_files .iter() @@ -109,7 +109,7 @@ pub(crate) fn build_live_row_ids( // more segments unfiltered, but it also changes which backend entry point they // take (`search` vs `search_with_filter`), and those can differ in recall. Not // worth diverging for. - let nothing_restricted = row_selections.is_none_or(FileRowSelections::is_empty); + let nothing_restricted = row_ranges_by_file.is_none_or(RowRangesByFile::is_empty); if nothing_restricted && all_active && !has_relevant_dv { return Ok(None); } @@ -126,7 +126,8 @@ pub(crate) fn build_live_row_ids( .ok_or_else(|| data_invalid("vector source row counts overflow u64"))?; let active = active_source_files.contains(source_file.file_name()); if active && row_count > 0 { - match row_selections.and_then(|selections| selections.get(source_file.file_name())) { + match row_ranges_by_file.and_then(|selections| selections.get(source_file.file_name())) + { // Unrestricted: the whole active file range is live. This is the // no-entry case Java spells as `rowRanges == null`. None => { @@ -137,7 +138,7 @@ pub(crate) fn build_live_row_ids( // these bounds ride in on an engine-supplied split, so walking them // would be unbounded work driven by untrusted numbers. Mirrors Java // `live.addRange(range.addOffset(fileOffset))`. - Some(FileRowSelection::Ranges(ranges)) => { + Some(ranges) => { for range in ranges { // Java checks each range against the SOURCE file's row count. // On the bucket-split route that count came off the wire as @@ -162,32 +163,6 @@ pub(crate) fn build_live_row_ids( live.insert_range((file_offset + from)..=(file_offset + to)); } } - // Restricted to positions a residual predicate left behind. Bounded - // by the rows that read actually returned, so walking them is safe. - Some(FileRowSelection::Positions(allowed)) => { - // `len` plus a maximum of `row_count - 1` can only describe the - // full set; inserting it as one range subsumes the per-position - // bound check below. - if allowed.len() == row_count && allowed.max() == Some(row_count - 1) { - charge_live_rows(&mut budget, row_count)?; - live.insert_range(file_offset..end); - } else { - charge_live_rows(&mut budget, allowed.len())?; - for position in allowed.iter() { - if position >= row_count { - return Err(data_invalid(format!( - "residual position {position} is out of range for source file {} ({} rows)", - source_file.file_name(), - row_count - ))); - } - let global = file_offset.checked_add(position).ok_or_else(|| { - data_invalid("vector residual position overflows u64") - })?; - live.insert(global); - } - } - } } } if active { @@ -228,7 +203,7 @@ pub(crate) fn map_ann_results( source_meta: &PkVectorSourceMeta, active_source_files: &HashSet, deletion_vectors: &HashMap>, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, metric: VectorSearchMetric, ) -> crate::Result> { let mut results = Vec::with_capacity(scored.len()); @@ -252,9 +227,9 @@ pub(crate) fn map_ann_results( } // A file with no entry is unrestricted, so only an entry can reject. if let Some(selection) = - row_selections.and_then(|selections| selections.get(&data_file_name)) + row_ranges_by_file.and_then(|selections| selections.get(&data_file_name)) { - if !selection.contains(pos) { + if !contains_row_position(selection, row_position) { return Err(data_invalid(format!( "ANN segment returned row position {row_position} in {data_file_name} outside the row selection for that file" ))); @@ -314,7 +289,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>>; #[allow(clippy::too_many_arguments)] @@ -328,7 +303,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { match segment_source { AnnSegmentSource::Buffered(bytes) => self.search_batch( @@ -340,7 +315,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files, deletion_vectors, search_options, - row_selections, + row_ranges_by_file, ), AnnSegmentSource::Vindex(_) => Err(data_invalid( "ANN searcher does not support a range-backed segment source", @@ -362,7 +337,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result> { let mut results = self.search_batch( segment, @@ -373,7 +348,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files, deletion_vectors, search_options, - row_selections, + row_ranges_by_file, )?; if results.len() != 1 { return Err(data_invalid(format!( @@ -395,7 +370,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result> { let mut results = self.search_batch_source( segment, @@ -406,7 +381,7 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { active_source_files, deletion_vectors, search_options, - row_selections, + row_ranges_by_file, )?; if results.len() != 1 { return Err(data_invalid(format!( @@ -538,7 +513,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { self.search_batch_source( segment, @@ -549,7 +524,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { active_source_files, deletion_vectors, search_options, - row_selections, + row_ranges_by_file, ) } @@ -563,7 +538,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { if limit == 0 { return Err(data_invalid("vector search limit must be positive")); @@ -579,7 +554,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { source_files, active_source_files, deletion_vectors, - row_selections, + row_ranges_by_file, )?; let mut searches = Vec::with_capacity(queries.len()); for query in queries { @@ -608,7 +583,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { &segment.source_meta, active_source_files, deletion_vectors, - row_selections, + row_ranges_by_file, metric, )? } @@ -623,6 +598,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { #[cfg(test)] mod tests { use super::*; + use crate::table::RowRange; use roaring::RoaringBitmap; /// A trivial loader returning empty bytes — the synthetic scorers below ignore @@ -653,26 +629,22 @@ mod tests { Arc::new(DeletionVector::from_bitmap(bitmap)) } - /// A residual selection: the physical positions of one file that passed a data - /// predicate. - fn positions(at: &[u64]) -> FileRowSelection { - let mut t = roaring::RoaringTreemap::new(); - for &p in at { - t.insert(p); - } - FileRowSelection::Positions(t) - } - - /// A pre-filter selection in the interval form an engine's split carries. - fn ranges(bounds: &[(i64, i64)]) -> FileRowSelection { - FileRowSelection::Ranges( - bounds - .iter() - .map(|(from, to)| crate::table::RowRange::new(*from, *to)) + /// Build the merged ranges a residual predicate hands to the search. + fn positions(at: &[u64]) -> Vec { + crate::table::merge_row_ranges( + at.iter() + .map(|&p| RowRange::new(p as i64, p as i64)) .collect(), ) } + fn ranges(bounds: &[(i64, i64)]) -> Vec { + bounds + .iter() + .map(|&(from, to)| RowRange::new(from, to)) + .collect() + } + fn active_set(names: &[&str]) -> HashSet { names.iter().map(|n| (*n).to_string()).collect() } @@ -979,14 +951,6 @@ mod tests { assert!(results.is_empty()); } - fn treemap(positions: &[u64]) -> roaring::RoaringTreemap { - let mut t = roaring::RoaringTreemap::new(); - for &p in positions { - t.insert(p); - } - t - } - #[test] fn test_build_live_row_ids_residual_intersects_with_active_and_dv() { // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). Both active. @@ -1002,10 +966,7 @@ mod tests { let mut dvs = HashMap::new(); dvs.insert("f0".to_string(), dv(&[1])); let mut residual = HashMap::new(); - residual.insert( - "f0".to_string(), - FileRowSelection::Positions(treemap(&[0, 1])), - ); + residual.insert("f0".to_string(), positions(&[0, 1])); let live = build_live_row_ids(&files, &active_set(&["f0", "f1"]), &dvs, Some(&residual)) .unwrap() .unwrap(); @@ -1023,14 +984,8 @@ mod tests { ]; let active = active_set(&["f0", "f1"]); let mut residual = HashMap::new(); - residual.insert( - "f0".to_string(), - FileRowSelection::Positions(treemap(&[0, 1, 2])), - ); - residual.insert( - "f1".to_string(), - FileRowSelection::Positions(treemap(&[0, 1])), - ); + residual.insert("f0".to_string(), positions(&[0, 1, 2])); + residual.insert("f1".to_string(), positions(&[0, 1])); let spelled_out = build_live_row_ids(&files, &active, &HashMap::new(), Some(&residual)) .unwrap() @@ -1049,10 +1004,7 @@ mod tests { let mut dvs = HashMap::new(); dvs.insert("f0".to_string(), dv(&[1])); let mut residual = HashMap::new(); - residual.insert( - "f0".to_string(), - FileRowSelection::Positions(treemap(&[0, 1, 2])), - ); + residual.insert("f0".to_string(), positions(&[0, 1, 2])); let live = build_live_row_ids(&files, &active_set(&["f0"]), &dvs, Some(&residual)) .unwrap() .unwrap(); @@ -1100,8 +1052,8 @@ mod tests { PkVectorSourceFile::new("f1".into(), 2).unwrap(), ]; let mut residual = HashMap::new(); - residual.insert("f0".to_string(), FileRowSelection::Positions(treemap(&[2]))); - residual.insert("f1".to_string(), FileRowSelection::Positions(treemap(&[1]))); + residual.insert("f0".to_string(), positions(&[2])); + residual.insert("f1".to_string(), positions(&[1])); let live = build_live_row_ids( &files, &active_set(&["f0", "f1"]), @@ -1119,10 +1071,7 @@ mod tests { // present, a mask is always required. let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()]; let mut residual = HashMap::new(); - residual.insert( - "f0".to_string(), - FileRowSelection::Positions(treemap(&[0, 2])), - ); + residual.insert("f0".to_string(), positions(&[0, 2])); let live = build_live_row_ids( &files, &active_set(&["f0"]), @@ -1140,10 +1089,7 @@ mod tests { // naming position 3 is out of range and must fail loud, not be skipped. let files = source_meta(&[("f0", 3)]); let mut residual = HashMap::new(); - residual.insert( - "f0".to_string(), - FileRowSelection::Positions(treemap(&[0, 3])), - ); + residual.insert("f0".to_string(), positions(&[0, 3])); let err = build_live_row_ids( files.source_files(), &active_set(&["f0"]), @@ -1182,7 +1128,7 @@ mod tests { // (e.g. an ANN reader that ignored include_row_ids) must fail loud. let meta = source_meta(&[("f0", 3)]); let mut residual = HashMap::new(); - residual.insert("f0".to_string(), FileRowSelection::Positions(treemap(&[0]))); + residual.insert("f0".to_string(), positions(&[0])); let err = map_ann_results( &[(1u64, 0.5)], &meta, @@ -1215,10 +1161,7 @@ mod tests { ); let segment = BucketAnnSegment::for_test(source_meta(&[("f0", 3)])); let mut residual = HashMap::new(); - residual.insert( - "f0".to_string(), - FileRowSelection::Positions(treemap(&[0, 2])), - ); + residual.insert("f0".to_string(), positions(&[0, 2])); searcher .search( &segment, @@ -1255,7 +1198,7 @@ mod tests { ), ); let segment = BucketAnnSegment::for_test(source_meta(&[("f0", 3)])); - let no_prefilter: FileRowSelections = HashMap::new(); + let no_prefilter: RowRangesByFile = HashMap::new(); searcher .search( &segment, @@ -1395,16 +1338,9 @@ mod tests { ); } - /// A treemap holding `0..=to`, built as one run so the test itself stays cheap. - fn positions_through(to: u64) -> roaring::RoaringTreemap { - let mut t = roaring::RoaringTreemap::new(); - t.insert_range(0..=to); - t - } - #[test] fn an_oversized_range_selection_is_charged() { - // The `Ranges` charge site, distinct from the unrestricted one: the file is + // The range charge site, distinct from the unrestricted one: the file is // restricted, so it never reaches the whole-file insert. let files = vec![PkVectorSourceFile::new("f0".into(), i32::MAX as i64 + 1).unwrap()]; let mut selections = HashMap::new(); @@ -1421,46 +1357,21 @@ mod tests { } #[test] - fn an_oversized_whole_file_position_set_is_charged() { - // The `Positions` whole-file shortcut: `len` equals the row count and the - // maximum is the last row, so it inserts as one range. - let rows = i32::MAX as u64 + 1; - let files = vec![PkVectorSourceFile::new("f0".into(), rows as i64).unwrap()]; - let mut selections = HashMap::new(); - selections.insert( - "f0".to_string(), - FileRowSelection::Positions(positions_through(rows - 1)), - ); - let error = build_live_row_ids( - &files, - &active_set(&["f0"]), - &HashMap::new(), - Some(&selections), - ) - .map(|_| ()) - .expect_err("a whole-file position set this large cannot be filtered"); - assert!(error.to_string().contains("more than"), "{error}"); - } - - #[test] - fn an_oversized_sparse_position_set_is_charged() { - // The per-position `Positions` path: the set is large but is NOT the whole - // file, so the shortcut above does not apply and the loop would walk it. - let rows = i32::MAX as u64 + 5; - let files = vec![PkVectorSourceFile::new("f0".into(), rows as i64).unwrap()]; - let mut selections = HashMap::new(); - selections.insert( + fn the_live_row_budget_is_shared_across_ranges() { + // Each range fits individually; their combined size exceeds the mask + // budget. Keep the first range tiny so the rejected test allocates little. + let files = vec![PkVectorSourceFile::new("f0".into(), i32::MAX as i64 + 2).unwrap()]; + let selections = HashMap::from([( "f0".to_string(), - FileRowSelection::Positions(positions_through(i32::MAX as u64)), - ); + ranges(&[(0, 0), (2, i32::MAX as i64 + 1)]), + )]); let error = build_live_row_ids( &files, &active_set(&["f0"]), &HashMap::new(), Some(&selections), ) - .map(|_| ()) - .expect_err("a position set this large cannot be filtered"); + .unwrap_err(); assert!(error.to_string().contains("more than"), "{error}"); } diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index a724544f0..62f5a1acc 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -26,9 +26,10 @@ use super::ann::PkVectorAnnSearcher; use super::data_invalid; use super::metric::{java_float_compare, VectorSearchMetric}; use super::result::PkVectorSearchResult; -use super::{FileRowSelection, FileRowSelections}; +use super::{contains_row_position, RowRangesByFile}; use crate::deletion_vector::DeletionVector; use crate::spec::PrimaryKeyIndexSourceMeta as PkVectorSourceMeta; +use crate::table::RowRange; use crate::vindex::executor::{ acquire_process_global_search_permit, drain_indexed_jobs, execute_global_index, }; @@ -164,7 +165,7 @@ fn validate_per_query_len( /// closure borrows the allow-list for its lifetime. fn position_excluder( dv: Option>, - selection: Option<&FileRowSelection>, + selection: Option<&[RowRange]>, ) -> impl Fn(i64) -> bool + Sync + '_ { move |position: i64| -> bool { let dv_deleted = match &dv { @@ -180,10 +181,7 @@ fn position_excluder( // No entry: the file is unrestricted, so the row is allowed. None => false, // Restricted: exclude positions the selection does not list. - Some(selection) => match u64::try_from(position) { - Ok(p) => !selection.contains(p), - Err(_) => true, - }, + Some(ranges) => !contains_row_position(ranges, position), } } } @@ -355,11 +353,11 @@ enum BucketLeaf { /// `ann_searcher` may be `None` only when there are no ANN segments; segments /// present with `None` is an error. /// -/// `row_selections` is the pre-filter allow-list keyed by data-file name: a file +/// `row_ranges_by_file` is the pre-filter allow-list keyed by data-file name: a file /// with **no entry is unrestricted**, an empty entry excludes it (the file is /// skipped without a read), and a non-empty one limits which of its rows may /// produce candidates. `None` restricts nothing at all. Mirrors Java -/// `rowRangesByFile`; see [`FileRowSelection`]. +/// `rowRangesByFile`; see [`RowRangesByFile`]. #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] pub(crate) async fn bucket_search( @@ -382,7 +380,7 @@ pub(crate) async fn bucket_search( exact_limit: usize, search_options: &HashMap, skip_exact_fallback: bool, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, concurrency: usize, search_budget: Option, ) -> crate::Result { @@ -517,7 +515,7 @@ pub(crate) async fn bucket_search( let ann_shared = searcher.as_ref().map(|searcher| { ( searcher.clone(), - row_selections.map(|selections| Arc::new(selections.clone())), + row_ranges_by_file.map(|selections| Arc::new(selections.clone())), Arc::new(active_source_files.clone()), Arc::new(deletion_vectors.clone()), Arc::new(search_options.clone()), @@ -538,13 +536,16 @@ pub(crate) async fn bucket_search( // No entry means unrestricted; an empty one excludes the file, which is // skipped without a read. Mirrors Java // `if (rowRanges != null && rowRanges.isEmpty()) continue;`. - let selection: Option<&FileRowSelection> = - match row_selections.and_then(|selections| selections.get(&file.file_name)) { - Some(selection) if selection.is_excluded() => continue, + let selection: Option<&Vec> = + match row_ranges_by_file.and_then(|selections| selections.get(&file.file_name)) { + Some(selection) if selection.is_empty() => continue, other => other, }; let dv = deletion_vectors.get(&file.file_name).cloned(); - exact_tasks.push((file, Box::new(position_excluder(dv, selection)))); + exact_tasks.push(( + file, + Box::new(position_excluder(dv, selection.map(Vec::as_slice))), + )); } } @@ -661,7 +662,7 @@ pub(crate) async fn bucket_search_batch( exact_limit: usize, search_options: &HashMap, skip_exact_fallback: bool, - row_selections: Option<&FileRowSelections>, + row_ranges_by_file: Option<&RowRangesByFile>, concurrency: usize, search_budget: Option, ) -> crate::Result> { @@ -684,7 +685,7 @@ pub(crate) async fn bucket_search_batch( exact_limit, search_options, skip_exact_fallback, - row_selections, + row_ranges_by_file, concurrency, search_budget, ) @@ -809,7 +810,7 @@ pub(crate) async fn bucket_search_batch( Arc::new(queries.iter().map(|q| q.to_vec()).collect()); ( searcher.clone(), - row_selections.map(|selections| Arc::new(selections.clone())), + row_ranges_by_file.map(|selections| Arc::new(selections.clone())), Arc::new(active_source_files.clone()), Arc::new(deletion_vectors.clone()), Arc::new(search_options.clone()), @@ -828,13 +829,16 @@ pub(crate) async fn bucket_search_batch( // No entry means unrestricted; an empty one excludes the file, which is // skipped without a read. Mirrors Java // `if (rowRanges != null && rowRanges.isEmpty()) continue;`. - let selection: Option<&FileRowSelection> = - match row_selections.and_then(|selections| selections.get(&file.file_name)) { - Some(selection) if selection.is_excluded() => continue, + let selection: Option<&Vec> = + match row_ranges_by_file.and_then(|selections| selections.get(&file.file_name)) { + Some(selection) if selection.is_empty() => continue, other => other, }; let dv = deletion_vectors.get(&file.file_name).cloned(); - exact_tasks.push((file, Box::new(position_excluder(dv, selection)))); + exact_tasks.push(( + file, + Box::new(position_excluder(dv, selection.map(Vec::as_slice))), + )); } } @@ -1080,7 +1084,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { Ok(queries.iter().map(|_| self.result.clone()).collect()) } @@ -1743,17 +1747,14 @@ mod tests { assert!(covered.is_empty()); } - fn treemap(positions: &[u64]) -> roaring::RoaringTreemap { - let mut t = roaring::RoaringTreemap::new(); - for &p in positions { - t.insert(p); - } - t - } - /// A residual selection over one file's physical positions. - fn selected(positions: &[u64]) -> FileRowSelection { - FileRowSelection::Positions(treemap(positions)) + fn selected(positions: &[u64]) -> Vec { + crate::table::merge_row_ranges( + positions + .iter() + .map(|&p| RowRange::new(p as i64, p as i64)) + .collect(), + ) } #[tokio::test] @@ -1766,7 +1767,7 @@ mod tests { Some(vec![2.0, 0.0]), Some(vec![3.0, 0.0]), ]); - let mut residual: FileRowSelections = HashMap::new(); + let mut residual: RowRangesByFile = HashMap::new(); residual.insert("data-1".into(), selected(&[0, 2])); let out = bucket_search( None, @@ -1837,7 +1838,7 @@ mod tests { }) }, ); - let mut residual: FileRowSelections = HashMap::new(); + let mut residual: RowRangesByFile = HashMap::new(); residual.insert("data-1".into(), selected(&[0, 1])); let out = bucket_search( None, @@ -1887,7 +1888,7 @@ mod tests { }) }, ); - let mut residual: FileRowSelections = HashMap::new(); + let mut residual: RowRangesByFile = HashMap::new(); residual.insert("data-1".into(), selected(&[])); let out = bucket_search( None, @@ -1925,7 +1926,7 @@ mod tests { let mut bm = RoaringBitmap::new(); bm.insert(0); // pos0 deleted dvs.insert("data-1".into(), Arc::new(DeletionVector::from_bitmap(bm))); - let mut residual: FileRowSelections = HashMap::new(); + let mut residual: RowRangesByFile = HashMap::new(); residual.insert("data-1".into(), selected(&[0, 1, 2])); let out = bucket_search( None, @@ -2560,7 +2561,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { use std::sync::atomic::Ordering::SeqCst; let current = self.inflight.fetch_add(1, SeqCst) + 1; @@ -2856,7 +2857,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { panic!("scorer panic to exercise JoinError mapping"); } @@ -2921,7 +2922,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { let file = segment.source_meta.source_files()[0] .file_name() @@ -3099,7 +3100,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { // Runs on the blocking pool. Announce arrival, then wait (bounded) for the // exact leaf. Both arriving proves overlap; a timeout means no overlap. @@ -3224,7 +3225,7 @@ mod tests { _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, - _row_selections: Option<&FileRowSelections>, + _row_ranges_by_file: Option<&RowRangesByFile>, ) -> crate::Result>> { // The bytes handed to the scorer must be exactly this segment's loaded // bytes (its path), proving load→score threads the right payload. diff --git a/crates/paimon/src/vindex/pkvector/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs index e383a8a69..fcc6e8a64 100644 --- a/crates/paimon/src/vindex/pkvector/mod.rs +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -36,71 +36,24 @@ pub(crate) fn data_invalid(message: impl Into) -> crate::Error { } } -/// Which physical rows of one data file a bucket search may read. -/// -/// Mirrors Java `rowRangesByFile` (`PkVectorAnnSegmentSearcher.liveRowPositions`, -/// `PrimaryKeyVectorBucketSearch.search`), which is a three-state per file and -/// spells the third state as the absence of a map entry: -/// -/// * **absent from [`FileRowSelections`]** — unrestricted, every row is readable. -/// Java records an entry only for a file its own pre-filter narrowed, so a -/// split that narrowed nothing carries an empty map and restricts nothing. -/// * [`Ranges`](Self::Ranges)/[`Positions`](Self::Positions) **empty** — excluded, -/// no row of the file is readable (Java's empty `List`). -/// * non-empty — restricted to what it lists. -/// -/// The two non-absent variants differ only in where the restriction came from, -/// which decides its shape. A plan built from an engine's bucket split carries -/// interval [`Ranges`] straight off the wire and must never be expanded into -/// positions: the row counts bounding those intervals are untrusted, so -/// materializing one row per allowed position is unbounded work. A residual data -/// predicate produces [`Positions`], whose size is bounded by the rows its own -/// read actually returned. -#[derive(Debug, Clone, PartialEq)] -pub(crate) enum FileRowSelection { - /// Inclusive physical row ranges, sorted and non-overlapping. Empty excludes - /// the file. - Ranges(Vec), - /// Physical row positions. Empty excludes the file. - Positions(roaring::RoaringTreemap), -} - -/// Per-data-file row selections for one split. A file with no entry is -/// unrestricted; see [`FileRowSelection`]. -pub(crate) type FileRowSelections = std::collections::HashMap; +/// Per-file physical row ranges, matching Java `rowRangesByFile`. +/// Missing entries impose no restriction; an empty list excludes the file. +/// Non-empty lists contain sorted, merged, inclusive ranges. Engine-planned +/// ranges stay compact, and residual predicates produce the same representation. +pub(crate) type RowRangesByFile = std::collections::HashMap>; -impl FileRowSelection { - /// Whether the selection permits no row at all, which is how Java's empty - /// `List` reads: the file is skipped rather than searched. - pub(crate) fn is_excluded(&self) -> bool { - match self { - Self::Ranges(ranges) => ranges.is_empty(), - Self::Positions(positions) => positions.is_empty(), - } - } - - /// Whether `position` is permitted. Ranges are binary-searched rather than - /// expanded, mirroring Java `PkVectorAnnSegmentSearcher.contains`. - pub(crate) fn contains(&self, position: u64) -> bool { - match self { - Self::Ranges(ranges) => { - let position = match i64::try_from(position) { - Ok(position) => position, - Err(_) => return false, - }; - ranges - .binary_search_by(|range| { - if position < range.from() { - std::cmp::Ordering::Greater - } else if position > range.to() { - std::cmp::Ordering::Less - } else { - std::cmp::Ordering::Equal - } - }) - .is_ok() +/// Test a physical position against sorted, non-overlapping ranges without +/// expanding them, as in Java's PK vector readers. +pub(crate) fn contains_row_position(ranges: &[crate::table::RowRange], position: i64) -> bool { + ranges + .binary_search_by(|range| { + if position < range.from() { + std::cmp::Ordering::Greater + } else if position > range.to() { + std::cmp::Ordering::Less + } else { + std::cmp::Ordering::Equal } - Self::Positions(positions) => positions.contains(position), - } - } + }) + .is_ok() } From be23031241c98448fc55acbc610997f439bbe219 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 13 Sep 2026 22:40:39 +0800 Subject: [PATCH 3/3] fix(table): resolve vector search clippy failures --- .../src/table/batch_vector_search_builder.rs | 5 ++--- crates/paimon/src/table/de_vector_read.rs | 2 +- crates/paimon/src/table/vector_read.rs | 15 +++++---------- crates/paimon/src/table/vector_scan.rs | 16 ++++++++-------- crates/paimon/src/table/vector_search_builder.rs | 6 +++--- crates/paimon/tests/pk_vector_baseline_test.rs | 4 ++-- 6 files changed, 21 insertions(+), 27 deletions(-) diff --git a/crates/paimon/src/table/batch_vector_search_builder.rs b/crates/paimon/src/table/batch_vector_search_builder.rs index 4e1f9d9b5..f80a7b248 100644 --- a/crates/paimon/src/table/batch_vector_search_builder.rs +++ b/crates/paimon/src/table/batch_vector_search_builder.rs @@ -120,7 +120,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { /// Create an owned batch reader; result i belongs to input query i. pub fn new_read(&self) -> crate::Result { let column = self.column()?; - PlanContext::new( + let context = PlanContext::new( self.table, column, self.filter.as_ref(), @@ -145,8 +145,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { limit, &self.options, self.filter.as_ref(), - self.include_row_ids.as_ref(), - self.prepared_filter.as_ref(), + context, ) } diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index 328962ae0..e033f79fb 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -225,7 +225,7 @@ impl Read for DeVectorRead { } evaluate_batch_vector_search( VectorSearchEvaluation { - table: Some(&pinned_table), + table: Some(pinned_table), file_io: pinned_table.file_io(), table_path: pinned_table.location(), table_options: pinned_table.schema().options(), diff --git a/crates/paimon/src/table/vector_read.rs b/crates/paimon/src/table/vector_read.rs index a1c166c3d..be2c55f0d 100644 --- a/crates/paimon/src/table/vector_read.rs +++ b/crates/paimon/src/table/vector_read.rs @@ -19,17 +19,14 @@ use crate::spec::{CoreOptions, Predicate}; use crate::table::de_vector_read::DeVectorRead; -use crate::table::de_vector_scan::PreparedVectorSearchFilter; use crate::table::pk_vector_read::PkVectorRead; use crate::table::pk_vector_search_params::PkVectorSearchParams; use crate::table::vector_scan::{PlanContext, VectorScanPlan, VectorScanWork}; use crate::table::vector_search_common::{take_only_result, targets_primary_key_column}; use crate::table::Table; use crate::vector_search::SearchResult; -use roaring::RoaringTreemap; use std::collections::HashMap; use std::future::Future; -use std::sync::Arc; /// Execute a resolved plan without resolving another snapshot or manifest. /// @@ -64,7 +61,7 @@ pub struct BatchVectorRead { enum VectorReadKind { DataEvolution(DeVectorRead), - PrimaryKey(PkVectorRead), + PrimaryKey(Box), } impl BatchVectorRead { @@ -75,18 +72,16 @@ impl BatchVectorRead { limit: usize, options: &HashMap, filter: Option<&Predicate>, - include_row_ids: Option<&Arc>, - prepared: Option<&PreparedVectorSearchFilter>, + context: PlanContext, ) -> crate::Result { - let context = PlanContext::new(table, column, filter, include_row_ids, prepared)?; let core = CoreOptions::new(table.schema().options()); let reader = if targets_primary_key_column(&core, column) { let pk_col = core.primary_key_vector_index_column()?; let params = PkVectorSearchParams::resolve(table, options, filter, &pk_col, queries, limit)?; - VectorReadKind::PrimaryKey(PkVectorRead::new( + VectorReadKind::PrimaryKey(Box::new(PkVectorRead::new( table, options, filter, &pk_col, queries, limit, params, - )) + ))) } else { VectorReadKind::DataEvolution(DeVectorRead::new(column, queries, limit, options)?) }; @@ -103,7 +98,7 @@ impl BatchVectorRead { } match (&self.reader, plan.work) { (VectorReadKind::DataEvolution(reader), VectorScanWork::DataEvolution(plan)) => { - reader.read(plan).await + reader.read(*plan).await } (VectorReadKind::PrimaryKey(reader), VectorScanWork::PrimaryKey(plan)) => { reader.read(plan).await diff --git a/crates/paimon/src/table/vector_scan.rs b/crates/paimon/src/table/vector_scan.rs index a1190e004..396833522 100644 --- a/crates/paimon/src/table/vector_scan.rs +++ b/crates/paimon/src/table/vector_scan.rs @@ -48,7 +48,7 @@ pub struct VectorScanPlan { #[derive(Clone)] pub(super) enum VectorScanWork { - DataEvolution(DeVectorScanPlan), + DataEvolution(Box), PrimaryKey(PkVectorScanPlan), } @@ -127,8 +127,8 @@ pub struct VectorScan { } enum VectorScanKind { - DataEvolution(DeVectorScan), - PrimaryKey(PkVectorScan), + DataEvolution(Box), + PrimaryKey(Box), } impl VectorScan { @@ -150,19 +150,19 @@ impl VectorScan { source: None, } })?; - VectorScanKind::PrimaryKey(PkVectorScan::new( + VectorScanKind::PrimaryKey(Box::new(PkVectorScan::new( table, field_id, core.primary_key_vector_index_type(&column)?, filter.cloned(), - )) + ))) } else { - VectorScanKind::DataEvolution(DeVectorScan::new( + VectorScanKind::DataEvolution(Box::new(DeVectorScan::new( table, filter, include_row_ids, prepared, - )) + ))) }; Ok(Self { context, scan }) } @@ -170,7 +170,7 @@ impl VectorScan { pub async fn plan(&self) -> crate::Result { let work = match &self.scan { VectorScanKind::DataEvolution(scan) => { - VectorScanWork::DataEvolution(scan.plan().await?) + VectorScanWork::DataEvolution(Box::new(scan.plan().await?)) } VectorScanKind::PrimaryKey(scan) => VectorScanWork::PrimaryKey(scan.plan().await?), }; diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index f0c742f99..176d75f9c 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -19,7 +19,7 @@ use crate::spec::{CoreOptions, Predicate}; use crate::table::vector_read::{BatchVectorRead, VectorRead}; -use crate::table::vector_scan::VectorScan; +use crate::table::vector_scan::{PlanContext, VectorScan}; use crate::table::Table; use crate::vector_search::SearchResult; use std::collections::HashMap; @@ -98,6 +98,7 @@ impl<'a> VectorSearchBuilder<'a> { /// Create an owned reader; query errors are reported before planning. pub fn new_read(&self) -> crate::Result { let (column, query, limit) = self.query()?; + let context = PlanContext::new(self.table, column, self.filter.as_ref(), None, None)?; Ok(VectorRead { batch: BatchVectorRead::new( self.table, @@ -106,8 +107,7 @@ impl<'a> VectorSearchBuilder<'a> { limit, &self.options, self.filter.as_ref(), - None, - None, + context, )?, }) } diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs b/crates/paimon/tests/pk_vector_baseline_test.rs index 667cc7bcc..c11bb0c7f 100644 --- a/crates/paimon/tests/pk_vector_baseline_test.rs +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -951,7 +951,7 @@ async fn pk_vector_residual_filter_excludes_non_matching_rows() { .positions() .unwrap() .iter() - .map(|p| p.row_position as i64) + .map(|p| p.row_position) .collect::>(), unfiltered_ids .iter() @@ -973,7 +973,7 @@ async fn pk_vector_residual_filter_excludes_non_matching_rows() { .positions() .unwrap() .iter() - .map(|p| p.row_position as i64) + .map(|p| p.row_position) .collect::>(), expected_ids.iter().map(|id| *id as i64).collect::>() );