This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new cec2d4c1 perf(vindex): execute compatible queries through native batch
search (#646)
cec2d4c1 is described below
commit cec2d4c104590c25a20e16d6e479ea2bac371dec
Author: Junrui Lee <[email protected]>
AuthorDate: Tue Aug 4 21:38:31 2026 +0800
perf(vindex): execute compatible queries through native batch search (#646)
---
crates/paimon/src/table/vector_search_builder.rs | 196 +++++--
crates/paimon/src/vindex/range_reader.rs | 78 ++-
crates/paimon/src/vindex/reader.rs | 715 +++++++++++++++++++++--
3 files changed, 880 insertions(+), 109 deletions(-)
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index d882a551..b2d67597 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -65,6 +65,7 @@ use arrow_select::interleave::interleave_record_batch;
use futures::{stream, TryStreamExt};
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::{BinaryHeap, HashMap, HashSet};
@@ -98,11 +99,30 @@ impl VectorIndexBackend {
}
}
-fn current_tokio_runtime_handle() -> crate::Result<tokio::runtime::Handle> {
- 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)),
+async fn execute_vindex_searches<S: SeekRead + 'static>(
+ io_meta: GlobalIndexIOMeta,
+ options: HashMap<String, String>,
+ vector_searches: Vec<VectorSearch>,
+ source: S,
+ file_name: String,
+ shard_concurrency: usize,
+) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
+ 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(panic_context, move || {
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options)
+ .with_batch_shard_concurrency(shard_concurrency);
+ 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
}
pub struct VectorSearchBuilder<'a> {
@@ -839,7 +859,8 @@ async fn plan_and_search_pk_candidates_batch(
reader.visit_batch_vector_search(searches, |_|
Ok(Cursor::new(data)))
}
VectorIndexBackend::Vindex => {
- let mut reader =
VindexVectorGlobalIndexReader::new(io_meta, options.clone());
+ let mut reader =
VindexVectorGlobalIndexReader::new(io_meta, options.clone())
+ .with_batch_shard_concurrency(concurrency);
reader.visit_batch_vector_search(searches, |_|
Ok(Cursor::new(data)))
}
}
@@ -1424,10 +1445,21 @@ async fn evaluate_batch_vector_search(
let mut merged = vec![SearchResult::empty(); vector_searches.len()];
if !vector_entries.is_empty() {
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 range_read_permits =
Arc::new(tokio::sync::Semaphore::new(concurrency));
let futures: Vec<_> = vector_entries
.into_iter()
.map(|entry| {
+ let range_read_permits = Arc::clone(&range_read_permits);
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");
@@ -1482,64 +1514,62 @@ async fn evaluate_batch_vector_search(
.await?
}
VectorIndexBackend::Vindex => {
- if vector_searches.len() > 1 {
- let data = input.read().await.map_err(|e| {
- crate::Error::DataInvalid {
- message: format!(
- "Failed to read vindex index file
'{}': {}",
- file_name, e
- ),
- source: None,
- }
- })?;
- execute_global_index(
- "vindex global-index batch search task
failed",
- move || {
- let mut reader =
VindexVectorGlobalIndexReader::new(
- io_meta, options,
- );
-
reader.visit_batch_vector_search(&vector_searches, |_| {
- Ok(Cursor::new(data))
- })
- },
- )
- .await?
- } else {
- 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,
- }
- })?;
- let source = VindexFileReader::new(
- Arc::new(file_reader),
- current_tokio_runtime_handle()?,
- file_size,
- file_name.clone(),
- );
- execute_global_index(
- "vindex global-index search task failed",
- move || {
- let mut reader =
VindexVectorGlobalIndexReader::new(
- io_meta, options,
- );
- 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?
+ match tokio::runtime::Handle::try_current() {
+ Ok(runtime) => {
+ 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,
+ }
+ })?;
+ let source =
VindexFileReader::new_with_permits(
+ Arc::new(file_reader),
+ runtime,
+ range_read_permits,
+ file_size,
+ file_name.clone(),
+ );
+ execute_vindex_searches(
+ io_meta,
+ options,
+ vector_searches,
+ source,
+ file_name,
+ concurrency,
+ )
+ .await?
+ }
+ Err(_) if query_count > 1 => {
+ let data = input.read().await.map_err(|e| {
+ crate::Error::DataInvalid {
+ message: format!(
+ "Failed to read vindex index
file '{}': {}",
+ file_name, e
+ ),
+ source: None,
+ }
+ })?;
+ execute_vindex_searches(
+ io_meta,
+ options,
+ vector_searches,
+ Cursor::new(data),
+ file_name,
+ concurrency,
+ )
+ .await?
+ }
+ Err(error) => {
+ return Err(crate::Error::UnexpectedError {
+ message:
+ "Vector index range reader
requires a Tokio runtime"
+ .to_string(),
+ source: Some(Box::new(error)),
+ });
+ }
}
}
};
@@ -3627,6 +3657,46 @@ mod tests {
});
}
+ #[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();
diff --git a/crates/paimon/src/vindex/range_reader.rs
b/crates/paimon/src/vindex/range_reader.rs
index 1e66472c..07431749 100644
--- a/crates/paimon/src/vindex/range_reader.rs
+++ b/crates/paimon/src/vindex/range_reader.rs
@@ -67,16 +67,33 @@ pub(crate) struct VindexFileReader {
}
impl VindexFileReader {
+ #[cfg(test)]
pub(crate) fn new(
reader: Arc<dyn FileRead>,
runtime: tokio::runtime::Handle,
file_size: u64,
path: String,
+ ) -> Self {
+ Self::new_with_permits(
+ reader,
+ runtime,
+ Arc::new(tokio::sync::Semaphore::new(RANGE_READ_CONCURRENCY)),
+ file_size,
+ path,
+ )
+ }
+
+ pub(crate) fn new_with_permits(
+ reader: Arc<dyn FileRead>,
+ runtime: tokio::runtime::Handle,
+ permits: Arc<tokio::sync::Semaphore>,
+ file_size: u64,
+ path: String,
) -> Self {
Self {
reader,
runtime,
- permits:
Arc::new(tokio::sync::Semaphore::new(RANGE_READ_CONCURRENCY)),
+ permits,
file_size,
path,
scalar_cache: None,
@@ -288,6 +305,7 @@ mod tests {
use super::*;
use crate::io::FileIO;
use async_trait::async_trait;
+ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
use std::time::Duration;
@@ -324,6 +342,12 @@ mod tests {
runtime_id: Mutex<Option<tokio::runtime::Id>>,
}
+ struct ConcurrencyTrackingRead {
+ data: Bytes,
+ active: AtomicUsize,
+ max_active: AtomicUsize,
+ }
+
#[async_trait]
impl FileRead for RuntimeTrackingRead {
async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
@@ -332,6 +356,17 @@ mod tests {
}
}
+ #[async_trait]
+ impl FileRead for ConcurrencyTrackingRead {
+ async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
+ let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
+ self.max_active.fetch_max(active, Ordering::SeqCst);
+ tokio::time::sleep(Duration::from_millis(25)).await;
+ self.active.fetch_sub(1, Ordering::SeqCst);
+ Ok(self.data.slice(range.start as usize..range.end as usize))
+ }
+ }
+
#[async_trait]
impl FileRead for TrackingRead {
async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
@@ -557,6 +592,47 @@ mod tests {
assert_eq!(cloned.read_capabilities(),
SeekReadCapabilities::default());
}
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn shared_permits_bound_reads_across_independent_readers() {
+ let data = Bytes::from(vec![8u8; 1024]);
+ let tracking = Arc::new(ConcurrencyTrackingRead {
+ data: data.clone(),
+ active: AtomicUsize::new(0),
+ max_active: AtomicUsize::new(0),
+ });
+ let permits = Arc::new(tokio::sync::Semaphore::new(1));
+ let make_reader = |path: &str| {
+ let source: Arc<dyn FileRead> = tracking.clone();
+ VindexFileReader::new_with_permits(
+ source,
+ tokio::runtime::Handle::current(),
+ Arc::clone(&permits),
+ data.len() as u64,
+ path.to_string(),
+ )
+ };
+ let mut first_reader = make_reader("first.index");
+ let mut second_reader = make_reader("second.index");
+ assert!(Arc::ptr_eq(&first_reader.permits, &second_reader.permits));
+
+ let first = tokio::task::spawn_blocking(move || {
+ let mut output = [0u8; 128];
+ first_reader
+ .pread(&mut [ReadRequest::new(0, &mut output)])
+ .unwrap();
+ });
+ let second = tokio::task::spawn_blocking(move || {
+ let mut output = [0u8; 128];
+ second_reader
+ .pread(&mut [ReadRequest::new(128, &mut output)])
+ .unwrap();
+ });
+ first.await.unwrap();
+ second.await.unwrap();
+
+ assert_eq!(tracking.max_active.load(Ordering::SeqCst), 1);
+ }
+
#[test]
fn local_fs_read_completes_with_one_host_blocking_thread() {
let temp_dir = tempfile::tempdir().unwrap();
diff --git a/crates/paimon/src/vindex/reader.rs
b/crates/paimon/src/vindex/reader.rs
index ecc39a57..3dc9dda1 100644
--- a/crates/paimon/src/vindex/reader.rs
+++ b/crates/paimon/src/vindex/reader.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use crate::spec::CoreOptions;
use crate::vector_search::{GlobalIndexIOMeta, VectorSearch};
use paimon_vindex_core::distance::MetricType;
use paimon_vindex_core::index::{
@@ -27,6 +28,7 @@ use std::io;
const DEFAULT_NPROBE: usize = 16;
const NPROBE_PARAMETER: &str = "ivf.nprobe";
+const NATIVE_BATCH_OPERATION_WORKING_SET_BYTES: usize = 64 * 1024 * 1024;
trait ErasedSeekRead: Send {
fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) ->
io::Result<()>;
@@ -76,6 +78,7 @@ impl SeekRead for VindexInput {
pub struct VindexVectorGlobalIndexReader {
io_meta: GlobalIndexIOMeta,
options: HashMap<String, String>,
+ batch_shard_concurrency: Option<usize>,
reader: Option<VIndexReader<VindexInput>>,
metadata: Option<VectorIndexMetadata>,
}
@@ -85,17 +88,23 @@ impl VindexVectorGlobalIndexReader {
Self {
io_meta,
options,
+ batch_shard_concurrency: None,
reader: None,
metadata: None,
}
}
+ pub(crate) fn with_batch_shard_concurrency(mut self, concurrency: usize)
-> Self {
+ self.batch_shard_concurrency = Some(concurrency.max(1));
+ self
+ }
+
pub fn visit_vector_search<S: SeekRead + 'static>(
&mut self,
vector_search: &VectorSearch,
stream_fn: impl FnOnce(&str) -> crate::Result<S>,
) -> crate::Result<Option<HashMap<u64, f32>>> {
- self.ensure_loaded(stream_fn)?;
+ self.ensure_loaded(stream_fn, |_| Ok(()))?;
self.search(vector_search)
}
@@ -104,11 +113,56 @@ impl VindexVectorGlobalIndexReader {
vector_searches: &[VectorSearch],
stream_fn: impl FnOnce(&str) -> crate::Result<S>,
) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
- self.ensure_loaded(stream_fn)?;
- vector_searches
- .iter()
- .map(|vector_search| self.search(vector_search))
- .collect()
+ self.ensure_loaded(stream_fn, |_| Ok(()))?;
+ self.search_batch(vector_searches)
+ }
+
+ #[cfg(test)]
+ pub(crate) fn load<S: SeekRead + 'static>(
+ &mut self,
+ stream_fn: impl FnOnce(&str) -> crate::Result<S>,
+ ) -> crate::Result<()> {
+ self.ensure_loaded(stream_fn, |_| Ok(()))
+ }
+
+ pub(crate) fn metadata(&self) -> crate::Result<&VectorIndexMetadata> {
+ self.metadata
+ .as_ref()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "vindex metadata not initialized".to_string(),
+ source: None,
+ })
+ }
+
+ pub(crate) fn search_batch(
+ &mut self,
+ vector_searches: &[VectorSearch],
+ ) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
+ let shard_concurrency = match self.batch_shard_concurrency {
+ Some(concurrency) => concurrency,
+ None => CoreOptions::new(&self.options).global_index_thread_num()?,
+ };
+ let reader = self
+ .reader
+ .as_mut()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "vindex reader not initialized".to_string(),
+ source: None,
+ })?;
+ let metadata = self
+ .metadata
+ .as_ref()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "vindex metadata not initialized".to_string(),
+ source: None,
+ })?;
+ search_batch_vindex(
+ reader,
+ metadata,
+ &self.options,
+ vector_searches,
+ shard_concurrency,
+ )
}
fn search(&mut self, vector_search: &VectorSearch) ->
crate::Result<Option<HashMap<u64, f32>>> {
@@ -130,12 +184,38 @@ impl VindexVectorGlobalIndexReader {
search_vindex(reader, metadata, &self.options, vector_search)
}
- fn ensure_loaded<S: SeekRead + 'static>(
+ fn ensure_loaded<S, F>(
&mut self,
stream_fn: impl FnOnce(&str) -> crate::Result<S>,
- ) -> crate::Result<()> {
+ validate: F,
+ ) -> crate::Result<()>
+ where
+ S: SeekRead + 'static,
+ F: FnOnce(&VectorIndexMetadata) -> crate::Result<()>,
+ {
+ self.ensure_loaded_with_optimizer(stream_fn, validate, |reader| {
+ reader
+ .optimize_for_search()
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("Failed to optimize paimon-vindex-core
reader: {}", e),
+ source: Some(Box::new(e)),
+ })
+ })
+ }
+
+ fn ensure_loaded_with_optimizer<S, F, O>(
+ &mut self,
+ stream_fn: impl FnOnce(&str) -> crate::Result<S>,
+ validate: F,
+ optimize: O,
+ ) -> crate::Result<()>
+ where
+ S: SeekRead + 'static,
+ F: FnOnce(&VectorIndexMetadata) -> crate::Result<()>,
+ O: FnOnce(&mut VIndexReader<VindexInput>) -> crate::Result<()>,
+ {
if self.reader.is_some() {
- return Ok(());
+ return validate(self.metadata()?);
}
let source = stream_fn(&self.io_meta.file_path)?;
@@ -146,12 +226,8 @@ impl VindexVectorGlobalIndexReader {
}
})?;
let metadata = reader.metadata();
- reader
- .optimize_for_search()
- .map_err(|e| crate::Error::DataInvalid {
- message: format!("Failed to optimize paimon-vindex-core
reader: {}", e),
- source: Some(Box::new(e)),
- })?;
+ validate(&metadata)?;
+ optimize(&mut reader)?;
self.reader = Some(reader);
self.metadata = Some(metadata);
@@ -165,12 +241,35 @@ fn search_vindex(
options: &HashMap<String, String>,
vector_search: &VectorSearch,
) -> crate::Result<Option<HashMap<u64, f32>>> {
- let expected_dim = metadata.dimension;
- if vector_search.vector.len() != expected_dim {
+ let Some(prepared) = prepare_search(metadata, options, vector_search)?
else {
+ return Ok(None);
+ };
+ let (labels, distances) = execute_scalar_search(reader, vector_search,
&prepared)?;
+ let id_to_scores = collect_results(&labels, &distances, prepared.top_k,
metadata.metric);
+ if id_to_scores.is_empty() {
+ return Ok(None);
+ }
+
+ Ok(Some(id_to_scores))
+}
+
+#[derive(Clone, PartialEq, Eq)]
+struct PreparedSearch {
+ top_k: usize,
+ nprobe: usize,
+ filter_bytes: Option<Vec<u8>>,
+}
+
+fn prepare_search(
+ metadata: &VectorIndexMetadata,
+ options: &HashMap<String, String>,
+ vector_search: &VectorSearch,
+) -> crate::Result<Option<PreparedSearch>> {
+ if vector_search.vector.len() != metadata.dimension {
return Err(crate::Error::DataInvalid {
message: format!(
"Query vector dimension mismatch: index expects {}, but got
{}",
- expected_dim,
+ metadata.dimension,
vector_search.vector.len()
),
source: None,
@@ -178,48 +277,195 @@ fn search_vindex(
}
let count = usize::try_from(metadata.total_vectors).unwrap_or(0);
- let effective_k = std::cmp::min(vector_search.limit, count);
- if effective_k == 0 {
+ let mut top_k = vector_search.limit.min(count);
+ if top_k == 0 {
return Ok(None);
}
-
let nprobe = int_parameter(options, NPROBE_PARAMETER, DEFAULT_NPROBE)?;
- let params = VectorSearchParams::new(effective_k, nprobe);
- let (labels, distances) = if let Some(include_ids) =
&vector_search.include_row_ids {
+ let filter_bytes = if let Some(include_ids) =
&vector_search.include_row_ids {
if include_ids.is_empty() {
return Ok(None);
}
- let ek = std::cmp::min(effective_k, include_ids.len() as usize);
- let params = VectorSearchParams::new(params.top_k.min(ek), nprobe);
- let mut filter_bytes = Vec::new();
+ top_k = top_k.min(include_ids.len() as usize);
+ let mut bytes = Vec::new();
include_ids
- .serialize_into(&mut filter_bytes)
+ .serialize_into(&mut bytes)
.map_err(|e| crate::Error::DataInvalid {
message: format!("Failed to serialize vector search row-id
filter: {}", e),
source: Some(Box::new(e)),
})?;
- reader
- .search_with_roaring_filter(&vector_search.vector, params,
&filter_bytes)
- .map_err(|e| crate::Error::DataInvalid {
- message: format!("paimon-vindex-core filtered search failed:
{}", e),
- source: Some(Box::new(e)),
- })?
+ Some(bytes)
} else {
- reader
- .search(&vector_search.vector, params)
+ None
+ };
+
+ Ok(Some(PreparedSearch {
+ top_k,
+ nprobe,
+ filter_bytes,
+ }))
+}
+
+fn execute_scalar_search(
+ reader: &mut VIndexReader<impl SeekRead>,
+ vector_search: &VectorSearch,
+ prepared: &PreparedSearch,
+) -> crate::Result<(Vec<i64>, Vec<f32>)> {
+ let params = VectorSearchParams::new(prepared.top_k, prepared.nprobe);
+ match &prepared.filter_bytes {
+ Some(filter) => reader
+ .search_with_roaring_filter(&vector_search.vector, params, filter)
.map_err(|e| crate::Error::DataInvalid {
- message: format!("paimon-vindex-core search failed: {}", e),
+ message: format!("paimon-vindex-core filtered search failed:
{}", e),
source: Some(Box::new(e)),
- })?
- };
+ }),
+ None => {
+ reader
+ .search(&vector_search.vector, params)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core search failed: {}",
e),
+ source: Some(Box::new(e)),
+ })
+ }
+ }
+}
- let id_to_scores = collect_results(&labels, &distances, effective_k,
metadata.metric);
- if id_to_scores.is_empty() {
- return Ok(None);
+fn search_batch_vindex(
+ reader: &mut VIndexReader<impl SeekRead>,
+ metadata: &VectorIndexMetadata,
+ options: &HashMap<String, String>,
+ vector_searches: &[VectorSearch],
+ shard_concurrency: usize,
+) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
+ let mut results: Vec<Option<HashMap<u64, f32>>> =
+ (0..vector_searches.len()).map(|_| None).collect();
+ let mut groups: Vec<(PreparedSearch, Vec<usize>)> = Vec::new();
+
+ for (index, search) in vector_searches.iter().enumerate() {
+ let Some(prepared) = prepare_search(metadata, options, search)? else {
+ continue;
+ };
+ if let Some((_, indices)) = groups.iter_mut().find(|(key, _)| key ==
&prepared) {
+ indices.push(index);
+ } else {
+ groups.push((prepared, vec![index]));
+ }
}
- Ok(Some(id_to_scores))
+ for (prepared, indices) in groups {
+ let chunk_size = native_batch_chunk_size(metadata, &prepared,
shard_concurrency);
+ for indices in indices.chunks(chunk_size) {
+ if indices.len() == 1 {
+ let index = indices[0];
+ let (labels, distances) =
+ execute_scalar_search(reader, &vector_searches[index],
&prepared)?;
+ let map = collect_results(&labels, &distances, prepared.top_k,
metadata.metric);
+ if !map.is_empty() {
+ results[index] = Some(map);
+ }
+ continue;
+ }
+
+ let mut queries = Vec::with_capacity(indices.len() *
metadata.dimension);
+ for &index in indices {
+ queries.extend_from_slice(&vector_searches[index].vector);
+ }
+ let params = VectorSearchParams::new(prepared.top_k,
prepared.nprobe);
+ let (labels, distances) = match &prepared.filter_bytes {
+ Some(filter) => reader
+ .search_batch_with_roaring_filter(&queries, indices.len(),
params, filter)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core filtered batch
search failed: {}", e),
+ source: Some(Box::new(e)),
+ })?,
+ None => reader
+ .search_batch(&queries, indices.len(), params)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core batch search
failed: {}", e),
+ source: Some(Box::new(e)),
+ })?,
+ };
+ let expected = indices.len() * prepared.top_k;
+ if labels.len() != expected || distances.len() != expected {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "paimon-vindex-core batch search returned
labels/distances of length {}/{}, expected {expected}",
+ labels.len(),
+ distances.len()
+ ),
+ source: None,
+ });
+ }
+ for (query_index, &result_index) in indices.iter().enumerate() {
+ let start = query_index * prepared.top_k;
+ let end = start + prepared.top_k;
+ let map = collect_results(
+ &labels[start..end],
+ &distances[start..end],
+ prepared.top_k,
+ metadata.metric,
+ );
+ if !map.is_empty() {
+ results[result_index] = Some(map);
+ }
+ }
+ }
+ }
+
+ Ok(results)
+}
+
+fn native_batch_chunk_size(
+ metadata: &VectorIndexMetadata,
+ prepared: &PreparedSearch,
+ shard_concurrency: usize,
+) -> usize {
+ let per_shard_budget = NATIVE_BATCH_OPERATION_WORKING_SET_BYTES
+ .checked_div(shard_concurrency.max(1))
+ .unwrap_or(0);
+ let filter_bytes = prepared
+ .filter_bytes
+ .as_ref()
+ .map_or(0, |filter| filter.len().saturating_mul(2));
+ let query_budget = per_shard_budget.saturating_sub(filter_bytes);
+ query_budget
+ .checked_div(native_batch_query_working_set_bytes(metadata, prepared))
+ .unwrap_or(0)
+ .max(1)
+}
+
+fn native_batch_query_working_set_bytes(
+ metadata: &VectorIndexMetadata,
+ prepared: &PreparedSearch,
+) -> usize {
+ let query_vectors = metadata
+ .dimension
+ .saturating_mul(std::mem::size_of::<f32>() * 2);
+ let centroid_products =
metadata.nlist.saturating_mul(std::mem::size_of::<f32>());
+ let probe_results = prepared
+ .nprobe
+ .min(metadata.nlist)
+ .saturating_mul(std::mem::size_of::<usize>() +
std::mem::size_of::<f32>());
+ let top_k_results = prepared.top_k.saturating_mul(
+ std::mem::size_of::<i64>() + std::mem::size_of::<f32>() +
std::mem::size_of::<(f32, i64)>(),
+ );
+ let pq_tables = match (metadata.pq_m, metadata.pq_bits) {
+ (Some(m), Some(bits)) => 1usize
+ .checked_shl(bits as u32)
+ .unwrap_or(usize::MAX)
+ .saturating_mul(m)
+ .saturating_mul(std::mem::size_of::<f32>()),
+ _ => 0,
+ };
+
+ query_vectors
+ .saturating_add(centroid_products)
+ .saturating_add(probe_results)
+ .saturating_add(top_k_results)
+ .saturating_add(pq_tables)
+ .saturating_add(256)
+ .max(1)
}
fn collect_results(
@@ -305,6 +551,8 @@ mod tests {
use bytes::Bytes;
use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer,
VectorIndexWriter};
use paimon_vindex_core::io::{PosWriter, SeekReadCapabilities};
+ use std::cell::Cell;
+ use std::io::Cursor;
use std::ops::Range;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
@@ -385,17 +633,65 @@ mod tests {
Bytes::from(output)
}
- fn query() -> VectorSearch {
+ fn query_for_cluster(cluster: usize, limit: usize) -> VectorSearch {
+ let cluster = cluster as f32 * 100.0;
VectorSearch::new(
(0..TEST_DIMENSION)
- .map(|dimension| dimension as f32 * 0.01)
+ .map(|dimension| cluster + dimension as f32 * 0.01)
.collect(),
- 10,
+ limit,
"embedding".to_string(),
)
.unwrap()
}
+ fn query() -> VectorSearch {
+ query_for_cluster(0, 10)
+ }
+
+ async fn tracked_batch_search(
+ index: Bytes,
+ query_count: usize,
+ ) -> (Vec<Option<HashMap<u64, f32>>>, usize) {
+ tracked_batch_search_with_options(
+ index,
+ query_count,
+ HashMap::from([(NPROBE_PARAMETER.to_string(), "1".to_string())]),
+ 32,
+ )
+ .await
+ }
+
+ async fn tracked_batch_search_with_options(
+ index: Bytes,
+ query_count: usize,
+ options: HashMap<String, String>,
+ shard_concurrency: usize,
+ ) -> (Vec<Option<HashMap<u64, f32>>>, usize) {
+ let tracking = TrackingIndexRead::new(index.clone());
+ let source: Arc<dyn FileRead> = tracking.clone();
+ let runtime = tokio::runtime::Handle::current();
+ let results = tokio::task::spawn_blocking(move || {
+ let source = VindexFileReader::new(
+ source,
+ runtime,
+ index.len() as u64,
+ "batch.index".to_string(),
+ );
+ let io_meta =
+ GlobalIndexIOMeta::new("batch.index".to_string(), index.len()
as u64, Vec::new());
+ let searches = vec![query(); query_count];
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
options)
+ .with_batch_shard_concurrency(shard_concurrency);
+ reader
+ .visit_batch_vector_search(&searches, |_| Ok(source))
+ .unwrap()
+ })
+ .await
+ .unwrap();
+ (results, tracking.bytes_read.load(Ordering::SeqCst))
+ }
+
#[test]
fn test_convert_distance_to_score() {
assert_eq!(convert_distance_to_score(0.0, MetricType::L2), 1.0);
@@ -434,6 +730,43 @@ mod tests {
assert!(!result.contains_key(&3));
}
+ #[test]
+ fn native_batch_chunk_size_tracks_working_set_inputs() {
+ let base_metadata = VectorIndexMetadata {
+ index_type: paimon_vindex_core::index::IndexType::IvfFlat,
+ dimension: 128,
+ nlist: 256,
+ metric: MetricType::L2,
+ total_vectors: 8192,
+ pq_m: None,
+ pq_bits: None,
+ rq_bits: None,
+ diskann: None,
+ };
+ let base_prepared = PreparedSearch {
+ top_k: 10,
+ nprobe: 16,
+ filter_bytes: None,
+ };
+ let base = native_batch_chunk_size(&base_metadata, &base_prepared, 32);
+
+ let mut larger_index = base_metadata.clone();
+ larger_index.dimension *= 2;
+ larger_index.nlist *= 2;
+ assert!(native_batch_chunk_size(&larger_index, &base_prepared, 32) <
base);
+
+ let mut larger_top_k = base_prepared.clone();
+ larger_top_k.top_k *= 4;
+ assert!(native_batch_chunk_size(&base_metadata, &larger_top_k, 32) <
base);
+
+ let mut pq_metadata = base_metadata.clone();
+ pq_metadata.pq_m = Some(64);
+ pq_metadata.pq_bits = Some(8);
+ assert!(native_batch_chunk_size(&pq_metadata, &base_prepared, 32) <
base);
+
+ assert!(native_batch_chunk_size(&base_metadata, &base_prepared, 64) <
base);
+ }
+
#[test]
fn test_int_parameter() {
let mut options = HashMap::new();
@@ -500,4 +833,296 @@ mod tests {
"range search unexpectedly read the entire index"
);
}
+
+ #[test]
+ fn mixed_batch_matches_scalar_searches_and_preserves_order() {
+ let index = build_ivf_flat_index();
+ let options = HashMap::from([(NPROBE_PARAMETER.to_string(),
"1".to_string())]);
+ let mut include_row_ids = roaring::RoaringTreemap::new();
+ include_row_ids.insert(0);
+ include_row_ids.insert(16);
+ include_row_ids.insert(32);
+ let searches = vec![
+ query_for_cluster(0, 4),
+ query_for_cluster(5, 1),
+ query_for_cluster(15, 4),
+ query_for_cluster(0, 10).with_include_row_ids(include_row_ids),
+ ];
+
+ let scalar_meta =
+ GlobalIndexIOMeta::new("scalar.index".to_string(), index.len() as
u64, Vec::new());
+ let mut scalar_reader =
VindexVectorGlobalIndexReader::new(scalar_meta, options.clone());
+ scalar_reader
+ .load(|_| Ok(Cursor::new(index.clone())))
+ .unwrap();
+ let expected: Vec<_> = searches
+ .iter()
+ .map(|search| scalar_reader.search(search).unwrap())
+ .collect();
+
+ let batch_meta =
+ GlobalIndexIOMeta::new("batch.index".to_string(), index.len() as
u64, Vec::new());
+ let mut batch_reader = VindexVectorGlobalIndexReader::new(batch_meta,
options);
+ let actual = batch_reader
+ .visit_batch_vector_search(&searches, |_| Ok(Cursor::new(index)))
+ .unwrap();
+
+ assert_eq!(actual, expected);
+ assert_ne!(
+ actual[0], actual[2],
+ "interleaved batch groups lost query order"
+ );
+ assert_eq!(actual[1].as_ref().map(HashMap::len), Some(1));
+ assert_eq!(actual[3].as_ref().map(HashMap::len), Some(3));
+ }
+
+ #[test]
+ fn metadata_validation_runs_before_optimization() {
+ let index = build_ivf_flat_index();
+ let io_meta = GlobalIndexIOMeta::new(
+ "validated.index".to_string(),
+ index.len() as u64,
+ Vec::new(),
+ );
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
HashMap::new());
+ let optimized = Cell::new(false);
+
+ let error = reader
+ .ensure_loaded_with_optimizer(
+ |_| Ok(Cursor::new(index)),
+ |metadata| {
+ assert_eq!(metadata.dimension, TEST_DIMENSION);
+ Err(crate::Error::DataInvalid {
+ message: "rejected test metric".to_string(),
+ source: None,
+ })
+ },
+ |_| {
+ optimized.set(true);
+ Ok(())
+ },
+ )
+ .unwrap_err();
+
+ assert!(error.to_string().contains("rejected test metric"));
+ assert!(!optimized.get());
+ assert!(reader.metadata().is_err());
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn batch_search_reuses_probed_lists_and_avoids_full_file_read() {
+ let index = build_ivf_flat_index();
+ let batch_index = index.clone();
+ let options = HashMap::from([(NPROBE_PARAMETER.to_string(),
"1".to_string())]);
+ let search = query();
+
+ let scalar_tracking = TrackingIndexRead::new(index.clone());
+ let scalar_source: Arc<dyn FileRead> = scalar_tracking.clone();
+ let scalar_options = options.clone();
+ let scalar_search = search.clone();
+ let scalar_runtime = tokio::runtime::Handle::current();
+ let scalar_results = tokio::task::spawn_blocking(move || {
+ let source = VindexFileReader::new(
+ scalar_source,
+ scalar_runtime,
+ index.len() as u64,
+ "scalar.index".to_string(),
+ );
+ let io_meta =
+ GlobalIndexIOMeta::new("scalar.index".to_string(), index.len()
as u64, Vec::new());
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
scalar_options);
+ reader.load(|_| Ok(source)).unwrap();
+ vec![
+ reader.search(&scalar_search).unwrap(),
+ reader.search(&scalar_search).unwrap(),
+ ]
+ })
+ .await
+ .unwrap();
+
+ let batch_tracking = TrackingIndexRead::new(batch_index.clone());
+ let batch_source: Arc<dyn FileRead> = batch_tracking.clone();
+ let batch_options = options;
+ let batch_search = search.clone();
+ let batch_runtime = tokio::runtime::Handle::current();
+ let batch_results = tokio::task::spawn_blocking(move || {
+ let source = VindexFileReader::new(
+ batch_source,
+ batch_runtime,
+ batch_index.len() as u64,
+ "batch.index".to_string(),
+ );
+ let io_meta = GlobalIndexIOMeta::new(
+ "batch.index".to_string(),
+ batch_index.len() as u64,
+ Vec::new(),
+ );
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
batch_options);
+ reader
+ .visit_batch_vector_search(&[batch_search.clone(),
batch_search], |_| Ok(source))
+ .unwrap()
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(batch_results, scalar_results);
+ let scalar_bytes = scalar_tracking.bytes_read.load(Ordering::SeqCst);
+ let batch_bytes = batch_tracking.bytes_read.load(Ordering::SeqCst);
+ assert!(
+ batch_bytes < scalar_bytes,
+ "batch should read a shared probed list once: batch={batch_bytes},
scalar={scalar_bytes}"
+ );
+ assert!(
+ batch_bytes < batch_tracking.data.len() / 2,
+ "nprobe=1 should read substantially less than the full index:
read={batch_bytes}, file={} ",
+ batch_tracking.data.len()
+ );
+ assert!(batch_tracking
+ .ranges()
+ .iter()
+ .all(|range| { range.start != 0 || range.end !=
batch_tracking.data.len() as u64 }));
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn filtered_batch_matches_scalar_searches_and_reuses_probed_lists() {
+ let index = build_ivf_flat_index();
+ let batch_index = index.clone();
+ let options = HashMap::from([(NPROBE_PARAMETER.to_string(),
"1".to_string())]);
+ let mut include_row_ids = roaring::RoaringTreemap::new();
+ for row_id in (0..256).step_by(16) {
+ include_row_ids.insert(row_id);
+ }
+ let searches = vec![
+ query().with_include_row_ids(include_row_ids.clone()),
+ query().with_include_row_ids(include_row_ids),
+ ];
+
+ let scalar_tracking = TrackingIndexRead::new(index.clone());
+ let scalar_source: Arc<dyn FileRead> = scalar_tracking.clone();
+ let scalar_options = options.clone();
+ let scalar_searches = searches.clone();
+ let scalar_runtime = tokio::runtime::Handle::current();
+ let scalar_results = tokio::task::spawn_blocking(move || {
+ let source = VindexFileReader::new(
+ scalar_source,
+ scalar_runtime,
+ index.len() as u64,
+ "scalar-filtered.index".to_string(),
+ );
+ let io_meta = GlobalIndexIOMeta::new(
+ "scalar-filtered.index".to_string(),
+ index.len() as u64,
+ Vec::new(),
+ );
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
scalar_options);
+ reader.load(|_| Ok(source)).unwrap();
+ scalar_searches
+ .iter()
+ .map(|search| reader.search(search).unwrap())
+ .collect::<Vec<_>>()
+ })
+ .await
+ .unwrap();
+
+ let batch_tracking = TrackingIndexRead::new(batch_index.clone());
+ let batch_source: Arc<dyn FileRead> = batch_tracking.clone();
+ let batch_runtime = tokio::runtime::Handle::current();
+ let batch_results = tokio::task::spawn_blocking(move || {
+ let source = VindexFileReader::new(
+ batch_source,
+ batch_runtime,
+ batch_index.len() as u64,
+ "batch-filtered.index".to_string(),
+ );
+ let io_meta = GlobalIndexIOMeta::new(
+ "batch-filtered.index".to_string(),
+ batch_index.len() as u64,
+ Vec::new(),
+ );
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
options);
+ reader
+ .visit_batch_vector_search(&searches, |_| Ok(source))
+ .unwrap()
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(batch_results, scalar_results);
+ let scalar_bytes = scalar_tracking.bytes_read.load(Ordering::SeqCst);
+ let batch_bytes = batch_tracking.bytes_read.load(Ordering::SeqCst);
+ assert!(
+ batch_bytes < scalar_bytes,
+ "filtered batch should read a shared probed list once:
batch={batch_bytes}, scalar={scalar_bytes}"
+ );
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn large_homogeneous_batch_reuses_lists_across_previous_boundary() {
+ let previous_batch_boundary = 16;
+ let index = build_ivf_flat_index();
+ let (within_limit_results, within_limit_bytes) =
+ tracked_batch_search(index.clone(), previous_batch_boundary).await;
+ let (over_limit_results, over_limit_bytes) =
+ tracked_batch_search(index, previous_batch_boundary + 1).await;
+
+ assert_eq!(within_limit_results.len(), previous_batch_boundary);
+ assert_eq!(over_limit_results.len(), previous_batch_boundary + 1);
+ assert!(within_limit_results
+ .iter()
+ .all(|result| result == &within_limit_results[0]));
+ assert!(over_limit_results
+ .iter()
+ .all(|result| result == &over_limit_results[0]));
+ assert_eq!(over_limit_results[0], within_limit_results[0]);
+ assert_eq!(
+ over_limit_bytes, within_limit_bytes,
+ "compatible queries should reuse the same probed list across the
former 16-query boundary"
+ );
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn homogeneous_batch_chunks_at_working_set_boundary() {
+ let shard_concurrency = 4096;
+ let metadata = VectorIndexMetadata {
+ index_type: paimon_vindex_core::index::IndexType::IvfFlat,
+ dimension: TEST_DIMENSION,
+ nlist: 16,
+ metric: MetricType::L2,
+ total_vectors: 8192,
+ pq_m: None,
+ pq_bits: None,
+ rq_bits: None,
+ diskann: None,
+ };
+ let options = HashMap::from([
+ (NPROBE_PARAMETER.to_string(), "1".to_string()),
+ ("global-index.thread-num".to_string(), "1".to_string()),
+ ]);
+ let prepared = prepare_search(&metadata, &options, &query())
+ .unwrap()
+ .unwrap();
+ let chunk_size = native_batch_chunk_size(&metadata, &prepared,
shard_concurrency);
+ assert!(chunk_size > 16);
+
+ let index = build_ivf_flat_index();
+ let (within_results, within_bytes) = tracked_batch_search_with_options(
+ index.clone(),
+ chunk_size,
+ options.clone(),
+ shard_concurrency,
+ )
+ .await;
+ let (over_results, over_bytes) =
+ tracked_batch_search_with_options(index, chunk_size + 1, options,
shard_concurrency)
+ .await;
+
+ assert_eq!(within_results.len(), chunk_size);
+ assert_eq!(over_results.len(), chunk_size + 1);
+ assert!(within_results
+ .iter()
+ .all(|result| result == &within_results[0]));
+ assert!(over_results.iter().all(|result| result == &over_results[0]));
+ assert_eq!(over_results[0], within_results[0]);
+ assert!(over_bytes > within_bytes);
+ }
}
