Skip to content

perf(vindex): execute compatible queries through native batch search #646

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JingsongLi merged 3 commits into from
Aug 4, 2026
+880 −109
Merged

perf(vindex): execute compatible queries through native batch search #646

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
414853f
perf(vindex): execute compatible queries through native batch search
JunRuiLee Aug 1, 2026
8dcb20d
fix(vindex): bound native batch resource usage
JunRuiLee Aug 4, 2026
49545a9
fix(vindex): limit test constructor to tests
JunRuiLee Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
196 changes: 133 additions & 63 deletions crates/paimon/src/table/vector_search_builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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)))
}
}
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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)),
});
}
}
}
};
Expand Down Expand Up @@ -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();
Expand Down
78 changes: 77 additions & 1 deletion crates/paimon/src/vindex/range_reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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> {
Expand All @@ -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> {
Expand Down Expand Up @@ -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();
Expand Down
Loading
Loading