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 8cd11661 perf(vindex): read search indexes through positional ranges
(#645)
8cd11661 is described below
commit 8cd1166110504f9a882d779cb1bf7643e95de54b
Author: Junrui Lee <[email protected]>
AuthorDate: Mon Aug 3 16:39:15 2026 +0800
perf(vindex): read search indexes through positional ranges (#645)
---
Cargo.lock | 1 +
benchmarks/tpcds/DEPENDENCIES.rust.tsv | 1 +
bindings/c/DEPENDENCIES.rust.tsv | 1 +
bindings/go/DEPENDENCIES.rust.tsv | 1 +
crates/integration_tests/DEPENDENCIES.rust.tsv | 1 +
crates/paimon-rest-server/DEPENDENCIES.rust.tsv | 1 +
crates/paimon/Cargo.toml | 1 +
crates/paimon/src/spec/core_options.rs | 7 +-
crates/paimon/src/table/vector_search_builder.rs | 155 +++++-
crates/paimon/src/vindex/executor.rs | 605 +++++++++++++++++++++
crates/paimon/src/vindex/mod.rs | 2 +
crates/paimon/src/vindex/range_reader.rs | 644 +++++++++++++++++++++++
crates/paimon/src/vindex/reader.rs | 231 +++++++-
13 files changed, 1602 insertions(+), 49 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 45b953b4..87cdeca7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4535,6 +4535,7 @@ dependencies = [
"bytes",
"chrono",
"crc32fast",
+ "crossbeam-channel",
"futures",
"hex",
"hmac 0.12.1",
diff --git a/benchmarks/tpcds/DEPENDENCIES.rust.tsv
b/benchmarks/tpcds/DEPENDENCIES.rust.tsv
index da26dc28..4b76a75b 100644
--- a/benchmarks/tpcds/DEPENDENCIES.rust.tsv
+++ b/benchmarks/tpcds/DEPENDENCIES.rust.tsv
@@ -81,6 +81,7 @@ [email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
diff --git a/bindings/c/DEPENDENCIES.rust.tsv b/bindings/c/DEPENDENCIES.rust.tsv
index f196c1cf..244fb110 100644
--- a/bindings/c/DEPENDENCIES.rust.tsv
+++ b/bindings/c/DEPENDENCIES.rust.tsv
@@ -60,6 +60,7 @@ [email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
diff --git a/bindings/go/DEPENDENCIES.rust.tsv
b/bindings/go/DEPENDENCIES.rust.tsv
index f196c1cf..244fb110 100644
--- a/bindings/go/DEPENDENCIES.rust.tsv
+++ b/bindings/go/DEPENDENCIES.rust.tsv
@@ -60,6 +60,7 @@ [email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
diff --git a/crates/integration_tests/DEPENDENCIES.rust.tsv
b/crates/integration_tests/DEPENDENCIES.rust.tsv
index fffee985..07114e66 100644
--- a/crates/integration_tests/DEPENDENCIES.rust.tsv
+++ b/crates/integration_tests/DEPENDENCIES.rust.tsv
@@ -60,6 +60,7 @@ [email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
diff --git a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv
b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv
index fbfc944b..e068e1d7 100644
--- a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv
+++ b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv
@@ -63,6 +63,7 @@ [email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
[email protected] X
X
diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml
index b8db7fa9..22fb72d1 100644
--- a/crates/paimon/Cargo.toml
+++ b/crates/paimon/Cargo.toml
@@ -97,6 +97,7 @@ arrow-schema = { workspace = true }
arrow-select = { workspace = true }
arrow-string = { workspace = true }
futures = "0.3"
+crossbeam-channel = "0.5"
tokio-util = { workspace = true, features = ["compat"] }
parquet = { workspace = true, features = ["async", "zstd", "lz4", "snap"] }
orc-rust = "0.8.0"
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index e26c6bc3..81d85ed4 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -647,9 +647,10 @@ impl<'a> CoreOptions<'a> {
/// Maximum number of concurrent tasks for global-index I/O, mirroring Java
/// `CoreOptions.GLOBAL_INDEX_THREAD_NUM` (key `global-index.thread-num`,
/// default 32). Used as the per-operation fan-out limit for sorted BTree
and
- /// bitmap shard reads and for primary-key vector search. A value of `1`
- /// reproduces strict sequential execution. A non-positive value is a
- /// misconfiguration and fails loud rather than being silently clamped.
+ /// bitmap shard reads, global-index vector search, and primary-key vector
+ /// search. A value of `1` reproduces strict sequential execution. A
+ /// non-positive value is a misconfiguration and fails loud rather than
being
+ /// silently clamped.
pub fn global_index_thread_num(&self) -> crate::Result<usize> {
let value = self
.parse_i64_option(GLOBAL_INDEX_THREAD_NUM_OPTION)?
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index 82183f38..281f77bb 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -50,10 +50,14 @@ 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::{
+ drain_indexed_jobs, ensure_global_index_executor_capacity,
execute_global_index,
+};
use crate::vindex::pkvector::ann::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::range_reader::VindexFileReader;
use crate::vindex::reader::VindexVectorGlobalIndexReader;
use crate::vindex::{is_vindex_index_type, VindexVectorIndexOptions};
use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array,
ListArray, RecordBatch};
@@ -94,6 +98,13 @@ 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)),
+ })
+}
+
pub struct VectorSearchBuilder<'a> {
table: &'a Table,
vector_column: Option<String>,
@@ -1412,6 +1423,8 @@ 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()?;
+ ensure_global_index_executor_capacity(concurrency);
let futures: Vec<_> = vector_entries
.into_iter()
.map(|entry| {
@@ -1440,39 +1453,102 @@ async fn evaluate_batch_vector_search(
let input = evaluation.file_io.new_input(&path);
async move {
let input = input?;
- let bytes = input.read().await.map_err(|e|
crate::Error::DataInvalid {
- message: format!(
- "Failed to read {} index file '{}': {}",
- backend.error_name(),
- file_name,
- e
- ),
- source: None,
- })?;
-
+ let query_count = vector_searches.len();
let io_meta =
GlobalIndexIOMeta::new(file_name.clone(), file_size,
index_meta_bytes);
- let data = bytes.to_vec();
let results = match backend {
VectorIndexBackend::Lumina => {
- let mut reader =
LuminaVectorGlobalIndexReader::new(io_meta, options);
- reader.visit_batch_vector_search(&vector_searches,
|_| {
- Ok(Cursor::new(data))
- })?
+ 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,
+ }
+ })?;
+ execute_global_index(
+ "Lumina global-index batch search task failed",
+ move || {
+ let mut reader =
+
LuminaVectorGlobalIndexReader::new(io_meta, options);
+
reader.visit_batch_vector_search(&vector_searches, |_| {
+ Ok(Cursor::new(data))
+ })
+ },
+ )
+ .await?
}
VectorIndexBackend::Vindex => {
- let mut reader =
VindexVectorGlobalIndexReader::new(io_meta, options);
- reader.visit_batch_vector_search(&vector_searches,
|_| {
- Ok(Cursor::new(data))
- })?
+ 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?
+ }
}
};
- if results.len() != vector_searches.len() {
+ if results.len() != query_count {
return Err(crate::Error::DataInvalid {
message: format!(
"Batch vector search backend returned {}
results for {} query vectors",
results.len(),
- vector_searches.len()
+ query_count
),
source: None,
});
@@ -1492,7 +1568,7 @@ async fn evaluate_batch_vector_search(
})
.collect();
- let results = futures::future::try_join_all(futures).await?;
+ let results = drain_indexed_jobs(futures.into_iter(),
concurrency).await?;
for per_entry in &results {
for (query_index, result) in per_entry.iter().enumerate() {
merged[query_index] = merged[query_index].or(result);
@@ -3514,6 +3590,41 @@ mod tests {
.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:?}"
+ );
+ });
}
#[tokio::test]
diff --git a/crates/paimon/src/vindex/executor.rs
b/crates/paimon/src/vindex/executor.rs
new file mode 100644
index 00000000..d8947102
--- /dev/null
+++ b/crates/paimon/src/vindex/executor.rs
@@ -0,0 +1,605 @@
+// 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 crossbeam_channel::{Receiver, RecvTimeoutError, Sender};
+use futures::stream::{self, StreamExt};
+use std::any::Any;
+use std::panic::{catch_unwind, AssertUnwindSafe};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex, OnceLock};
+use std::time::Duration;
+use tokio::sync::oneshot;
+
+type Job = Box<dyn FnOnce() + Send + 'static>;
+
+const DEFAULT_IO_BOUND_WORKERS: usize = 32;
+const IO_BOUND_WORKERS_PER_CPU: usize = 4;
+const WORKER_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
+
+struct ExecutorState {
+ receiver: Receiver<Job>,
+ max_workers: AtomicUsize,
+ minimum_workers: usize,
+ physical_worker_limit: usize,
+ worker_idle_timeout: Duration,
+ worker_count: AtomicUsize,
+ outstanding_jobs: AtomicUsize,
+ next_worker_id: AtomicUsize,
+}
+
+struct GlobalIndexExecutor {
+ sender: Sender<Job>,
+ state: Arc<ExecutorState>,
+}
+
+#[derive(Default)]
+struct TaskCompletionState {
+ started: bool,
+ cancelled: bool,
+}
+
+struct TaskCompletion {
+ state: Mutex<TaskCompletionState>,
+}
+
+impl TaskCompletion {
+ fn new() -> Self {
+ Self {
+ state: Mutex::new(TaskCompletionState::default()),
+ }
+ }
+
+ fn try_start(&self) -> bool {
+ let mut state = self.state.lock().unwrap_or_else(|error|
error.into_inner());
+ if state.cancelled {
+ return false;
+ }
+ state.started = true;
+ true
+ }
+
+ fn cancel_if_queued(&self) {
+ let mut state = self.state.lock().unwrap_or_else(|error|
error.into_inner());
+ if !state.started {
+ state.cancelled = true;
+ }
+ }
+}
+
+struct TaskCompletionGuard(Arc<TaskCompletion>);
+
+impl Drop for TaskCompletionGuard {
+ fn drop(&mut self) {
+ self.0.cancel_if_queued();
+ }
+}
+
+impl GlobalIndexExecutor {
+ fn new(max_workers: usize) -> Self {
+ Self::new_with_limits(
+ max_workers,
+ max_physical_worker_count(),
+ WORKER_IDLE_TIMEOUT,
+ )
+ }
+
+ fn new_with_limits(
+ max_workers: usize,
+ physical_worker_limit: usize,
+ worker_idle_timeout: Duration,
+ ) -> Self {
+ let (sender, receiver) = crossbeam_channel::unbounded();
+ let physical_worker_limit = physical_worker_limit.max(1);
+ let max_workers = max_workers.max(1).min(physical_worker_limit);
+ Self {
+ sender,
+ state: Arc::new(ExecutorState {
+ receiver,
+ max_workers: AtomicUsize::new(max_workers),
+ minimum_workers: max_workers,
+ physical_worker_limit,
+ worker_idle_timeout,
+ worker_count: AtomicUsize::new(0),
+ outstanding_jobs: AtomicUsize::new(0),
+ next_worker_id: AtomicUsize::new(0),
+ }),
+ }
+ }
+
+ fn ensure_capacity(&self, max_workers: usize) {
+ let max_workers =
max_workers.max(1).min(self.state.physical_worker_limit);
+ self.state
+ .max_workers
+ .fetch_max(max_workers, Ordering::SeqCst);
+ // Another query may already have queued work, so raising the
high-watermark
+ // can immediately make more workers useful even before this caller
submits.
+ let _ = self.ensure_workers_for_load();
+ }
+
+ fn submit(&self, job: Job) -> crate::Result<()> {
+ self.state.outstanding_jobs.fetch_add(1, Ordering::SeqCst);
+ if let Err(error) = self.ensure_workers_for_load() {
+ self.state.outstanding_jobs.fetch_sub(1, Ordering::SeqCst);
+ return Err(error);
+ }
+ if self.sender.send(job).is_err() {
+ self.state.outstanding_jobs.fetch_sub(1, Ordering::SeqCst);
+ return Err(crate::Error::UnexpectedError {
+ message: "global-index executor stopped before accepting a
task".to_string(),
+ source: None,
+ });
+ }
+ Ok(())
+ }
+
+ fn ensure_workers_for_load(&self) -> crate::Result<()> {
+ ensure_workers_for_load(&self.state)
+ }
+}
+
+fn ensure_workers_for_load(state: &Arc<ExecutorState>) -> crate::Result<()> {
+ loop {
+ let current = state.worker_count.load(Ordering::SeqCst);
+ let max_workers = state.max_workers.load(Ordering::SeqCst);
+ let outstanding = state.outstanding_jobs.load(Ordering::SeqCst);
+ let target = outstanding.min(max_workers);
+ if current >= target {
+ return Ok(());
+ }
+ if state
+ .worker_count
+ .compare_exchange(current, current + 1, Ordering::SeqCst,
Ordering::SeqCst)
+ .is_err()
+ {
+ continue;
+ }
+
+ let worker_id = state.next_worker_id.fetch_add(1, Ordering::Relaxed);
+ let worker_state = Arc::clone(state);
+ if let Err(error) = std::thread::Builder::new()
+ .name(format!("paimon-global-index-{worker_id}"))
+ .spawn(move || run_worker(worker_state))
+ {
+ state.worker_count.fetch_sub(1, Ordering::SeqCst);
+ return Err(crate::Error::UnexpectedError {
+ message: format!("failed to start global-index executor
thread: {error}"),
+ source: Some(Box::new(error)),
+ });
+ }
+ }
+}
+
+fn run_worker(state: Arc<ExecutorState>) {
+ loop {
+ match state.receiver.recv_timeout(state.worker_idle_timeout) {
+ Ok(job) => {
+ let _ = catch_unwind(AssertUnwindSafe(job));
+ state.outstanding_jobs.fetch_sub(1, Ordering::SeqCst);
+ }
+ Err(RecvTimeoutError::Timeout) => {
+ if retire_idle_worker(&state) {
+ // A submit racing with retirement may have observed the
old worker
+ // count. Re-check after decrementing so queued work can
replace us.
+ let _ = ensure_workers_for_load(&state);
+ return;
+ }
+ }
+ Err(RecvTimeoutError::Disconnected) => {
+ state.worker_count.fetch_sub(1, Ordering::SeqCst);
+ return;
+ }
+ }
+ }
+}
+
+fn retire_idle_worker(state: &ExecutorState) -> bool {
+ loop {
+ let current = state.worker_count.load(Ordering::SeqCst);
+ if current <= state.minimum_workers {
+ return false;
+ }
+ if state
+ .worker_count
+ .compare_exchange(current, current - 1, Ordering::SeqCst,
Ordering::SeqCst)
+ .is_ok()
+ {
+ return true;
+ }
+ }
+}
+
+fn default_worker_count() -> usize {
+ std::thread::available_parallelism()
+ .map(|parallelism| parallelism.get())
+ .unwrap_or(1)
+}
+
+fn max_physical_worker_count() -> usize {
+ default_worker_count()
+ .saturating_mul(IO_BOUND_WORKERS_PER_CPU)
+ .max(DEFAULT_IO_BOUND_WORKERS)
+}
+
+/// Shared global-index executor for synchronous search and range-I/O waits. It
+/// starts at the machine's available parallelism and grows lazily, but keeps
the
+/// configured logical fan-out separate from a physical cap of four workers per
+/// CPU (and at least the default 32 I/O workers). Growth workers expire after
the
+/// same one-minute idle interval used by Java's `GlobalIndexReadThreadPool`.
+/// Smaller query limits are enforced by each query's bounded job scheduler.
+fn global_executor() -> &'static GlobalIndexExecutor {
+ static EXECUTOR: OnceLock<GlobalIndexExecutor> = OnceLock::new();
+ EXECUTOR.get_or_init(|| GlobalIndexExecutor::new(default_worker_count()))
+}
+
+pub(crate) fn ensure_global_index_executor_capacity(max_workers: usize) {
+ global_executor().ensure_capacity(max_workers);
+}
+
+/// Runs bounded global-index jobs to completion, restores submission order,
and
+/// returns the first error by submission index. Dedicated executor jobs
cannot be
+/// interrupted after they start, so short-circuiting would only hide
background
+/// work and make the surfaced error depend on completion timing.
+pub(crate) async fn drain_indexed_jobs<T, F>(
+ jobs: impl Iterator<Item = F>,
+ concurrency: usize,
+) -> crate::Result<Vec<T>>
+where
+ F: std::future::Future<Output = crate::Result<T>>,
+{
+ let indexed = jobs
+ .enumerate()
+ .map(|(index, job)| async move { (index, job.await) });
+ let mut collected: Vec<(usize, crate::Result<T>)> = stream::iter(indexed)
+ .buffer_unordered(concurrency.max(1))
+ .collect()
+ .await;
+ collected.sort_by_key(|(index, _)| *index);
+ collected.into_iter().map(|(_, result)| result).collect()
+}
+
+pub(crate) async fn execute_global_index<T, F>(
+ panic_context: &'static str,
+ task: F,
+) -> crate::Result<T>
+where
+ T: Send + 'static,
+ F: FnOnce() -> crate::Result<T> + Send + 'static,
+{
+ execute_on(global_executor(), panic_context, task).await
+}
+
+async fn execute_on<T, F>(
+ executor: &GlobalIndexExecutor,
+ panic_context: &'static str,
+ task: F,
+) -> crate::Result<T>
+where
+ T: Send + 'static,
+ F: FnOnce() -> crate::Result<T> + Send + 'static,
+{
+ let (sender, receiver) = oneshot::channel();
+ let completion = Arc::new(TaskCompletion::new());
+ let task_completion = Arc::clone(&completion);
+ executor.submit(Box::new(move || {
+ if !task_completion.try_start() {
+ return;
+ }
+ let outcome = catch_unwind(AssertUnwindSafe(task));
+ let _ = sender.send(outcome);
+ }))?;
+ // Cancellation marks a queued job so the worker skips it. Started work is
+ // detached like Tokio's spawn_blocking tasks; Drop must not block a
runtime
+ // thread because the search may still depend on async range I/O from it.
The
+ // caller runtime must remain driven until that I/O completes; dropping the
+ // runtime cancels the spawned read and releases the worker.
+ let _completion_guard = TaskCompletionGuard(completion);
+
+ let outcome = receiver
+ .await
+ .map_err(|error| crate::Error::UnexpectedError {
+ message: "global-index executor dropped a task result".to_string(),
+ source: Some(Box::new(error)),
+ })?;
+ match outcome {
+ Ok(result) => result,
+ Err(payload) => Err(crate::Error::UnexpectedError {
+ message: format!("{panic_context}: {}",
panic_message(payload.as_ref())),
+ source: None,
+ }),
+ }
+}
+
+fn panic_message(payload: &(dyn Any + Send)) -> &str {
+ if let Some(message) = payload.downcast_ref::<&str>() {
+ message
+ } else if let Some(message) = payload.downcast_ref::<String>() {
+ message.as_str()
+ } else {
+ "non-string panic payload"
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::atomic::{AtomicBool, AtomicUsize};
+ use std::time::Duration;
+
+ #[tokio::test]
+ async fn executor_respects_worker_limit() {
+ let executor = GlobalIndexExecutor::new(2);
+ let in_flight = Arc::new(AtomicUsize::new(0));
+ let peak = Arc::new(AtomicUsize::new(0));
+ let tasks = (0..8).map(|_| {
+ let in_flight = Arc::clone(&in_flight);
+ let peak = Arc::clone(&peak);
+ execute_on(&executor, "test global-index task failed", move || {
+ let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
+ peak.fetch_max(now, Ordering::SeqCst);
+ std::thread::sleep(Duration::from_millis(20));
+ in_flight.fetch_sub(1, Ordering::SeqCst);
+ Ok(())
+ })
+ });
+
+ futures::future::try_join_all(tasks).await.unwrap();
+ assert!(peak.load(Ordering::SeqCst) <= 2);
+ }
+
+ #[tokio::test]
+ async fn executor_grows_to_requested_capacity() {
+ let executor = Arc::new(GlobalIndexExecutor::new(1));
+ executor.ensure_capacity(3);
+ let (started_sender, mut started_receiver) =
tokio::sync::mpsc::unbounded_channel();
+ let (release_sender, release_receiver) = crossbeam_channel::bounded(3);
+ let executor_for_tasks = Arc::clone(&executor);
+ let tasks = (0..3).map(move |_| {
+ let executor = Arc::clone(&executor_for_tasks);
+ let started_sender = started_sender.clone();
+ let release_receiver = release_receiver.clone();
+ async move {
+ execute_on(&executor, "test global-index task failed", move ||
{
+ started_sender.send(()).unwrap();
+ release_receiver.recv().unwrap();
+ Ok(())
+ })
+ .await
+ }
+ });
+ let join = tokio::spawn(async move {
futures::future::try_join_all(tasks).await });
+
+ for _ in 0..3 {
+ tokio::time::timeout(Duration::from_secs(1),
started_receiver.recv())
+ .await
+ .expect("executor did not grow to requested capacity")
+ .expect("executor stopped before all tasks started");
+ }
+ for _ in 0..3 {
+ release_sender.send(()).unwrap();
+ }
+ join.await.unwrap().unwrap();
+ assert_eq!(executor.state.worker_count.load(Ordering::SeqCst), 3);
+ }
+
+ #[test]
+ fn executor_caps_requested_physical_capacity() {
+ let executor = GlobalIndexExecutor::new_with_limits(1, 3,
Duration::from_secs(60));
+ executor.ensure_capacity(usize::MAX);
+
+ assert_eq!(executor.state.max_workers.load(Ordering::SeqCst), 3);
+ }
+
+ #[tokio::test]
+ async fn executor_reclaims_idle_growth() {
+ let executor = Arc::new(GlobalIndexExecutor::new_with_limits(
+ 1,
+ 3,
+ Duration::from_millis(20),
+ ));
+ executor.ensure_capacity(3);
+ let (started_sender, mut started_receiver) =
tokio::sync::mpsc::unbounded_channel();
+ let (release_sender, release_receiver) = crossbeam_channel::bounded(3);
+ let executor_for_tasks = Arc::clone(&executor);
+ let tasks = (0..3).map(move |_| {
+ let executor = Arc::clone(&executor_for_tasks);
+ let started_sender = started_sender.clone();
+ let release_receiver = release_receiver.clone();
+ async move {
+ execute_on(&executor, "test global-index task failed", move ||
{
+ started_sender.send(()).unwrap();
+ release_receiver.recv().unwrap();
+ Ok(())
+ })
+ .await
+ }
+ });
+ let join = tokio::spawn(async move {
futures::future::try_join_all(tasks).await });
+
+ for _ in 0..3 {
+ tokio::time::timeout(Duration::from_secs(1),
started_receiver.recv())
+ .await
+ .expect("executor did not grow before idle reclamation")
+ .expect("executor stopped before all tasks started");
+ }
+ for _ in 0..3 {
+ release_sender.send(()).unwrap();
+ }
+ join.await.unwrap().unwrap();
+
+ tokio::time::timeout(Duration::from_secs(1), async {
+ while executor.state.worker_count.load(Ordering::SeqCst) != 1 {
+ tokio::time::sleep(Duration::from_millis(5)).await;
+ }
+ })
+ .await
+ .expect("executor did not reclaim idle growth workers");
+
+ execute_on(&executor, "test global-index task failed", || Ok(()))
+ .await
+ .unwrap();
+ }
+
+ #[tokio::test]
+ async fn drain_indexed_jobs_bounds_concurrency_and_restores_order() {
+ let in_flight = Arc::new(AtomicUsize::new(0));
+ let peak = Arc::new(AtomicUsize::new(0));
+ let jobs = (0..8).map(|index| {
+ let in_flight = Arc::clone(&in_flight);
+ let peak = Arc::clone(&peak);
+ async move {
+ let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
+ peak.fetch_max(now, Ordering::SeqCst);
+ tokio::time::sleep(Duration::from_millis((8 - index) as
u64)).await;
+ in_flight.fetch_sub(1, Ordering::SeqCst);
+ Ok(index)
+ }
+ });
+
+ let results = drain_indexed_jobs(jobs, 2).await.unwrap();
+ assert_eq!(results, (0..8).collect::<Vec<_>>());
+ assert_eq!(peak.load(Ordering::SeqCst), 2);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn cancelling_waiter_does_not_stop_submitted_task() {
+ let (started_sender, started_receiver) = oneshot::channel();
+ let (release_sender, release_receiver) = std::sync::mpsc::channel();
+ let completed = Arc::new(AtomicBool::new(false));
+ let completed_for_task = Arc::clone(&completed);
+ let join = tokio::spawn(async move {
+ execute_global_index("test global-index task failed", move || {
+ let _ = started_sender.send(());
+ release_receiver.recv().unwrap();
+ completed_for_task.store(true, Ordering::SeqCst);
+ Ok(())
+ })
+ .await
+ });
+ started_receiver.await.unwrap();
+
+ let releaser = std::thread::spawn(move || {
+ std::thread::sleep(Duration::from_millis(20));
+ release_sender.send(()).unwrap();
+ });
+ join.abort();
+ assert!(join.await.unwrap_err().is_cancelled());
+ releaser.join().unwrap();
+ assert!(completed.load(Ordering::SeqCst));
+ }
+
+ #[test]
+ fn cancelling_started_search_keeps_current_thread_runtime_live() {
+ let runtime = tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ .unwrap();
+ runtime.block_on(async {
+ let handle = tokio::runtime::Handle::current();
+ let (started_sender, started_receiver) = oneshot::channel();
+ let completed = Arc::new(AtomicBool::new(false));
+ let completed_for_task = Arc::clone(&completed);
+ let join = tokio::spawn(async move {
+ execute_global_index("test global-index task failed", move || {
+ let _ = started_sender.send(());
+ let (sender, receiver) = std::sync::mpsc::sync_channel(1);
+ handle.spawn(async move {
+ tokio::task::yield_now().await;
+ let _ = sender.send(());
+ });
+ receiver.recv().unwrap();
+ completed_for_task.store(true, Ordering::SeqCst);
+ Ok(())
+ })
+ .await
+ });
+ started_receiver.await.unwrap();
+
+ join.abort();
+ let error = tokio::time::timeout(Duration::from_secs(1), join)
+ .await
+ .expect("cancelling a search blocked the current-thread
runtime")
+ .unwrap_err();
+ assert!(error.is_cancelled());
+ tokio::time::timeout(Duration::from_secs(1), async {
+ while !completed.load(Ordering::SeqCst) {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .expect("detached search did not finish after caller
cancellation");
+ });
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn cancelling_queued_task_returns_without_waiting_for_queue() {
+ let executor = Arc::new(GlobalIndexExecutor::new(1));
+ let (started_sender, started_receiver) = oneshot::channel();
+ let (release_sender, release_receiver) = std::sync::mpsc::channel();
+ let first_executor = Arc::clone(&executor);
+ let first = tokio::spawn(async move {
+ execute_on(&first_executor, "first task failed", move || {
+ let _ = started_sender.send(());
+ release_receiver.recv().unwrap();
+ Ok(())
+ })
+ .await
+ });
+ started_receiver.await.unwrap();
+
+ let second_ran = Arc::new(AtomicBool::new(false));
+ let second_ran_in_task = Arc::clone(&second_ran);
+ let second_executor = Arc::clone(&executor);
+ let mut second = tokio::spawn(async move {
+ execute_on(&second_executor, "second task failed", move || {
+ second_ran_in_task.store(true, Ordering::SeqCst);
+ Ok(())
+ })
+ .await
+ });
+ tokio::time::timeout(Duration::from_secs(1), async {
+ while executor.state.outstanding_jobs.load(Ordering::SeqCst) < 2 {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .expect("second task was not queued");
+
+ second.abort();
+ let cancelled_without_waiting =
tokio::time::timeout(Duration::from_secs(1), &mut second)
+ .await
+ .is_ok();
+
+ release_sender.send(()).unwrap();
+ first.await.unwrap().unwrap();
+ if !cancelled_without_waiting {
+ assert!(second.await.unwrap_err().is_cancelled());
+ }
+ tokio::time::timeout(Duration::from_secs(1), async {
+ while executor.state.outstanding_jobs.load(Ordering::SeqCst) != 0 {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .expect("cancelled task was not drained from the executor queue");
+
+ assert!(
+ cancelled_without_waiting,
+ "cancelling a queued task waited for the task ahead of it"
+ );
+ assert!(!second_ran.load(Ordering::SeqCst));
+ }
+}
diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs
index f6e35ab5..c6db3adf 100644
--- a/crates/paimon/src/vindex/mod.rs
+++ b/crates/paimon/src/vindex/mod.rs
@@ -15,6 +15,8 @@
// specific language governing permissions and limitations
// under the License.
+pub(crate) mod executor;
+pub(crate) mod range_reader;
pub mod reader;
pub mod pkvector;
diff --git a/crates/paimon/src/vindex/range_reader.rs
b/crates/paimon/src/vindex/range_reader.rs
new file mode 100644
index 00000000..1e66472c
--- /dev/null
+++ b/crates/paimon/src/vindex/range_reader.rs
@@ -0,0 +1,644 @@
+// 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::io::FileRead;
+use bytes::Bytes;
+use futures::future::try_join_all;
+use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities};
+use std::io;
+use std::ops::Range;
+use std::sync::{mpsc, Arc};
+
+const SCALAR_READ_MAX: usize = 64;
+const SCALAR_READ_AHEAD: u64 = 64 * 1024;
+const RANGE_COALESCE_GAP: u64 = 16 * 1024;
+const RANGE_READ_CONCURRENCY: usize = 32;
+
+struct CachedRange {
+ start: u64,
+ data: Bytes,
+}
+
+impl CachedRange {
+ fn end(&self) -> u64 {
+ self.start + self.data.len() as u64
+ }
+
+ fn contains(&self, range: &Range<u64>) -> bool {
+ self.start <= range.start && range.end <= self.end()
+ }
+}
+
+struct RequestedRange {
+ request_index: usize,
+ range: Range<u64>,
+}
+
+struct MergedRange {
+ range: Range<u64>,
+ request_indices: Vec<usize>,
+ requested_bytes: u64,
+}
+
+/// Bridges vindex-core's synchronous positional reads to Paimon's asynchronous
+/// range reader. This type is consumed from a blocking search task; it
captures
+/// the surrounding Tokio runtime so remote storage reads still run
asynchronously.
+pub(crate) struct VindexFileReader {
+ reader: Arc<dyn FileRead>,
+ runtime: tokio::runtime::Handle,
+ permits: Arc<tokio::sync::Semaphore>,
+ file_size: u64,
+ path: String,
+ scalar_cache: Option<CachedRange>,
+}
+
+impl VindexFileReader {
+ pub(crate) fn new(
+ reader: Arc<dyn FileRead>,
+ runtime: tokio::runtime::Handle,
+ file_size: u64,
+ path: String,
+ ) -> Self {
+ Self {
+ reader,
+ runtime,
+ permits:
Arc::new(tokio::sync::Semaphore::new(RANGE_READ_CONCURRENCY)),
+ file_size,
+ path,
+ scalar_cache: None,
+ }
+ }
+
+ fn validate_range(&self, pos: u64, len: usize) -> io::Result<Range<u64>> {
+ let end = pos.checked_add(len as u64).ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidInput,
+ format!("vindex read range overflows for '{}'", self.path),
+ )
+ })?;
+ if end > self.file_size {
+ return Err(io::Error::new(
+ io::ErrorKind::UnexpectedEof,
+ format!(
+ "vindex read range {pos}..{end} exceeds file '{}' size {}",
+ self.path, self.file_size
+ ),
+ ));
+ }
+ Ok(pos..end)
+ }
+
+ fn read_one(&mut self, pos: u64, buf: &mut [u8]) -> io::Result<()> {
+ if buf.is_empty() {
+ return Ok(());
+ }
+ let range = self.validate_range(pos, buf.len())?;
+ if buf.len() <= SCALAR_READ_MAX {
+ if let Some(cache) = &self.scalar_cache {
+ if cache.contains(&range) {
+ let start = (range.start - cache.start) as usize;
+ buf.copy_from_slice(&cache.data[start..start + buf.len()]);
+ return Ok(());
+ }
+ }
+
+ let read_end = range
+ .start
+ .saturating_add(SCALAR_READ_AHEAD)
+ .max(range.end)
+ .min(self.file_size);
+ let data = self.fetch_exact(range.start..read_end)?;
+ buf.copy_from_slice(&data[..buf.len()]);
+ self.scalar_cache = Some(CachedRange {
+ start: range.start,
+ data,
+ });
+ return Ok(());
+ }
+
+ let data = self.fetch_exact(range)?;
+ buf.copy_from_slice(&data);
+ Ok(())
+ }
+
+ fn fetch_exact(&self, range: Range<u64>) -> io::Result<Bytes> {
+ let mut results =
self.fetch_range_batch(std::slice::from_ref(&range))?;
+ Ok(results.pop().expect("one requested range"))
+ }
+
+ fn fetch_range_batch(&self, ranges: &[Range<u64>]) ->
io::Result<Vec<Bytes>> {
+ debug_assert!(ranges.len() <= RANGE_READ_CONCURRENCY);
+ let reader = Arc::clone(&self.reader);
+ let permits = Arc::clone(&self.permits);
+ let path = self.path.clone();
+ let requested = ranges.to_vec();
+ let (sender, receiver) = mpsc::sync_channel(1);
+ self.runtime.spawn(async move {
+ let fetched = try_join_all(requested.iter().cloned().map(|range| {
+ let reader = Arc::clone(&reader);
+ let permits = Arc::clone(&permits);
+ let path = path.clone();
+ async move {
+ let _permit = permits.acquire_owned().await.map_err(|_| {
+ io::Error::other("vindex range read concurrency
limiter closed")
+ })?;
+ let expected = (range.end - range.start) as usize;
+ let data =
reader.read(range.clone()).await.map_err(|error| {
+ io::Error::other(format!(
+ "failed to read vindex file '{path}' range {}..{}:
{error}",
+ range.start, range.end
+ ))
+ })?;
+ if data.len() != expected {
+ return Err(io::Error::new(
+ io::ErrorKind::UnexpectedEof,
+ format!(
+ "short read for vindex file '{path}' range
{}..{}: expected {expected} bytes, got {}",
+ range.start,
+ range.end,
+ data.len()
+ ),
+ ));
+ }
+ Ok(data)
+ }
+ }))
+ .await;
+ let _ = sender.send(fetched);
+ });
+ receiver.recv().map_err(|_| {
+ io::Error::other(format!(
+ "vindex range read task for '{}' was cancelled",
+ self.path
+ ))
+ })?
+ }
+
+ fn read_many(&self, requests: &mut [ReadRequest<'_>]) -> io::Result<()> {
+ let mut requested = Vec::with_capacity(requests.len());
+ for (request_index, request) in requests.iter().enumerate() {
+ if request.buf.is_empty() {
+ continue;
+ }
+ requested.push(RequestedRange {
+ request_index,
+ range: self.validate_range(request.pos, request.buf.len())?,
+ });
+ }
+ if requested.is_empty() {
+ return Ok(());
+ }
+
+ requested.sort_by_key(|request| request.range.start);
+ let mut merged = Vec::<MergedRange>::new();
+ for request in requested {
+ if let Some(last) = merged.last_mut() {
+ let merged_end = last.range.end.max(request.range.end);
+ let merged_bytes = merged_end - last.range.start;
+ let requested_bytes = last
+ .requested_bytes
+ .saturating_add(request.range.end - request.range.start);
+ let within_gap =
+ request.range.start <=
last.range.end.saturating_add(RANGE_COALESCE_GAP);
+ let bounded_amplification = merged_bytes <=
requested_bytes.saturating_mul(2);
+ if within_gap && bounded_amplification {
+ last.range.end = last.range.end.max(request.range.end);
+ last.request_indices.push(request.request_index);
+ last.requested_bytes = requested_bytes;
+ continue;
+ }
+ }
+ let requested_bytes = request.range.end - request.range.start;
+ merged.push(MergedRange {
+ range: request.range,
+ request_indices: vec![request.request_index],
+ requested_bytes,
+ });
+ }
+
+ for batch in merged.chunks(RANGE_READ_CONCURRENCY) {
+ let ranges: Vec<_> = batch.iter().map(|merged|
merged.range.clone()).collect();
+ let fetched = self.fetch_range_batch(&ranges)?;
+ for (merged_range, data) in batch.iter().zip(fetched) {
+ for &request_index in &merged_range.request_indices {
+ let request = &mut requests[request_index];
+ let start = (request.pos - merged_range.range.start) as
usize;
+ request
+ .buf
+ .copy_from_slice(&data[start..start +
request.buf.len()]);
+ }
+ }
+ }
+ Ok(())
+ }
+}
+
+impl SeekRead for VindexFileReader {
+ fn pread(&mut self, requests: &mut [ReadRequest<'_>]) -> io::Result<()> {
+ let non_empty = requests
+ .iter()
+ .filter(|request| !request.buf.is_empty())
+ .count();
+ if non_empty == 1 {
+ let request = requests
+ .iter_mut()
+ .find(|request| !request.buf.is_empty())
+ .expect("one non-empty request");
+ self.read_one(request.pos, request.buf)
+ } else {
+ self.read_many(requests)
+ }
+ }
+
+ fn try_clone_reader(&self) -> io::Result<Option<Self>> {
+ Ok(Some(Self {
+ reader: Arc::clone(&self.reader),
+ runtime: self.runtime.clone(),
+ permits: Arc::clone(&self.permits),
+ file_size: self.file_size,
+ path: self.path.clone(),
+ scalar_cache: None,
+ }))
+ }
+
+ fn read_capabilities(&self) -> SeekReadCapabilities {
+ // This adapter accepts any number of ranges and splits them
internally.
+ // The efficient window size depends on the underlying FileRead
backend,
+ // so leave both storage-specific hints unspecified.
+ SeekReadCapabilities::default()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::io::FileIO;
+ use async_trait::async_trait;
+ use std::sync::Mutex;
+ use std::time::Duration;
+
+ struct TrackingRead {
+ data: Bytes,
+ ranges: Mutex<Vec<Range<u64>>>,
+ short_read: bool,
+ }
+
+ impl TrackingRead {
+ fn new(data: Bytes) -> Arc<Self> {
+ Arc::new(Self {
+ data,
+ ranges: Mutex::new(Vec::new()),
+ short_read: false,
+ })
+ }
+
+ fn short(data: Bytes) -> Arc<Self> {
+ Arc::new(Self {
+ data,
+ ranges: Mutex::new(Vec::new()),
+ short_read: true,
+ })
+ }
+
+ fn ranges(&self) -> Vec<Range<u64>> {
+ self.ranges.lock().unwrap().clone()
+ }
+ }
+
+ struct RuntimeTrackingRead {
+ data: Bytes,
+ runtime_id: Mutex<Option<tokio::runtime::Id>>,
+ }
+
+ #[async_trait]
+ impl FileRead for RuntimeTrackingRead {
+ async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
+ *self.runtime_id.lock().unwrap() =
Some(tokio::runtime::Handle::current().id());
+ 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> {
+ self.ranges.lock().unwrap().push(range.clone());
+ let mut end = range.end as usize;
+ if self.short_read && end > range.start as usize {
+ end -= 1;
+ }
+ Ok(self.data.slice(range.start as usize..end))
+ }
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn scalar_reads_reuse_bounded_read_ahead() {
+ let data = Bytes::from((0..200_000).map(|value| value as
u8).collect::<Vec<_>>());
+ let tracking = TrackingRead::new(data.clone());
+ let source: Arc<dyn FileRead> = tracking.clone();
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+
+ tokio::task::spawn_blocking(move || {
+ let mut first = [0u8; 8];
+ reader
+ .pread(&mut [ReadRequest::new(32, &mut first)])
+ .unwrap();
+ assert_eq!(&first, &data[32..40]);
+
+ let mut second = [0u8; 4];
+ reader
+ .pread(&mut [ReadRequest::new(128, &mut second)])
+ .unwrap();
+ assert_eq!(&second, &data[128..132]);
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(tracking.ranges(), vec![32..32 + SCALAR_READ_AHEAD]);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn multi_range_reads_merge_only_nearby_requests() {
+ let data = Bytes::from((0..100_000).map(|value| value as
u8).collect::<Vec<_>>());
+ let tracking = TrackingRead::new(data.clone());
+ let source: Arc<dyn FileRead> = tracking.clone();
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+
+ tokio::task::spawn_blocking(move || {
+ let mut first = [0u8; 4];
+ let mut second = [0u8; 4];
+ let mut third = [0u8; 4];
+ reader
+ .pread(&mut [
+ ReadRequest::new(0, &mut first),
+ ReadRequest::new(8, &mut second),
+ ReadRequest::new(20_000, &mut third),
+ ])
+ .unwrap();
+ assert_eq!(&first, &data[0..4]);
+ assert_eq!(&second, &data[8..12]);
+ assert_eq!(&third, &data[20_000..20_004]);
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(tracking.ranges(), vec![0..12, 20_000..20_004]);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn multi_range_reads_reject_sparse_amplification() {
+ let data = Bytes::from(vec![7u8; 100_000]);
+ let tracking = TrackingRead::new(data.clone());
+ let source: Arc<dyn FileRead> = tracking.clone();
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+
+ tokio::task::spawn_blocking(move || {
+ let mut first = [0u8; 4];
+ let mut second = [0u8; 4];
+ reader
+ .pread(&mut [
+ ReadRequest::new(0, &mut first),
+ ReadRequest::new(16_000, &mut second),
+ ])
+ .unwrap();
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(tracking.ranges(), vec![0..4, 16_000..16_004]);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn short_range_read_fails_loudly() {
+ let data = Bytes::from(vec![1u8; 1024]);
+ let tracking = TrackingRead::short(data.clone());
+ let source: Arc<dyn FileRead> = tracking;
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+
+ let error = tokio::task::spawn_blocking(move || {
+ let mut output = [0u8; 128];
+ reader
+ .pread(&mut [ReadRequest::new(256, &mut output)])
+ .unwrap_err()
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
+ assert!(error.to_string().contains("short read"));
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn worker_read_is_safe_on_multi_thread_runtime() {
+ let data = Bytes::from(vec![3u8; 1024]);
+ let source: Arc<dyn FileRead> = TrackingRead::new(data.clone());
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+ let output = tokio::task::spawn_blocking(move || {
+ let mut output = [0u8; 16];
+ reader
+ .pread(&mut [ReadRequest::new(128, &mut output)])
+ .unwrap();
+ output
+ })
+ .await
+ .unwrap();
+ assert_eq!(&output, &data[128..144]);
+ }
+
+ #[tokio::test]
+ async fn worker_read_is_safe_on_current_thread_runtime() {
+ let data = Bytes::from(vec![5u8; 1024]);
+ let source: Arc<dyn FileRead> = TrackingRead::new(data.clone());
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+ let (sender, receiver) = tokio::sync::oneshot::channel();
+ let worker = std::thread::spawn(move || {
+ let mut output = [0u8; 16];
+ let result = reader.pread(&mut [ReadRequest::new(256, &mut
output)]);
+ let _ = sender.send((result, output));
+ });
+ let (result, output) = receiver.await.unwrap();
+ result.unwrap();
+ worker.join().unwrap();
+ assert_eq!(&output, &data[256..272]);
+ }
+
+ #[test]
+ fn async_reads_run_on_the_calling_runtime() {
+ let runtime = tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ .unwrap();
+ let caller_runtime_id = runtime.handle().id();
+ let data = Bytes::from(vec![6u8; 1024]);
+ let tracking = Arc::new(RuntimeTrackingRead {
+ data: data.clone(),
+ runtime_id: Mutex::new(None),
+ });
+ let source: Arc<dyn FileRead> = tracking.clone();
+ let mut reader = VindexFileReader::new(
+ source,
+ runtime.handle().clone(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+ let (sender, receiver) = tokio::sync::oneshot::channel();
+ let worker = std::thread::spawn(move || {
+ let mut output = [0u8; 16];
+ let result = reader.pread(&mut [ReadRequest::new(128, &mut
output)]);
+ let _ = sender.send((result, output));
+ });
+
+ let (result, output) = runtime.block_on(receiver).unwrap();
+ result.unwrap();
+ worker.join().unwrap();
+ assert_eq!(&output, &data[128..144]);
+ assert_eq!(
+ *tracking.runtime_id.lock().unwrap(),
+ Some(caller_runtime_id)
+ );
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn reader_can_clone_for_parallel_search() {
+ let data = Bytes::from(vec![7u8; 1024]);
+ let source: Arc<dyn FileRead> = TrackingRead::new(data.clone());
+ let reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ "index".to_string(),
+ );
+
+ assert_eq!(reader.read_capabilities(),
SeekReadCapabilities::default());
+ let cloned = reader.try_clone_reader().unwrap().unwrap();
+ assert_eq!(cloned.read_capabilities(),
SeekReadCapabilities::default());
+ }
+
+ #[test]
+ fn local_fs_read_completes_with_one_host_blocking_thread() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let path = temp_dir.path().join("index.bin");
+ let data = vec![9u8; 4096];
+ std::fs::write(&path, &data).unwrap();
+ let runtime = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(1)
+ .max_blocking_threads(1)
+ .enable_all()
+ .build()
+ .unwrap();
+
+ runtime.block_on(async {
+ let file_io = FileIO::from_path(path.to_string_lossy())
+ .unwrap()
+ .build()
+ .unwrap();
+ let input =
file_io.new_input(path.to_string_lossy().as_ref()).unwrap();
+ let source: Arc<dyn FileRead> =
Arc::new(input.reader().await.unwrap());
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ path.to_string_lossy().into_owned(),
+ );
+
+ let (sender, receiver) = tokio::sync::oneshot::channel();
+ let worker = std::thread::spawn(move || {
+ let mut output = [0u8; 32];
+ let result = reader.pread(&mut [ReadRequest::new(512, &mut
output)]);
+ let _ = sender.send((result, output));
+ });
+ let (result, output) =
tokio::time::timeout(Duration::from_secs(5), receiver)
+ .await
+ .expect("vindex range read deadlocked on the host blocking
pool")
+ .unwrap();
+ result.unwrap();
+ worker.join().unwrap();
+ assert_eq!(output, [9u8; 32]);
+ });
+ }
+
+ #[test]
+ fn local_fs_read_is_safe_on_current_thread_runtime() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let path = temp_dir.path().join("index.bin");
+ let data = vec![11u8; 4096];
+ std::fs::write(&path, &data).unwrap();
+ let worker = std::thread::spawn(move || {
+ let runtime = tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ .unwrap();
+ runtime.block_on(async {
+ let file_io = FileIO::from_path(path.to_string_lossy())
+ .unwrap()
+ .build()
+ .unwrap();
+ let input =
file_io.new_input(path.to_string_lossy().as_ref()).unwrap();
+ let source: Arc<dyn FileRead> =
Arc::new(input.reader().await.unwrap());
+ let mut reader = VindexFileReader::new(
+ source,
+ tokio::runtime::Handle::current(),
+ data.len() as u64,
+ path.to_string_lossy().into_owned(),
+ );
+ let (sender, receiver) = tokio::sync::oneshot::channel();
+ let search = std::thread::spawn(move || {
+ let mut output = [0u8; 32];
+ let result = reader.pread(&mut [ReadRequest::new(1024,
&mut output)]);
+ let _ = sender.send((result, output));
+ });
+ let (result, output) =
tokio::time::timeout(Duration::from_secs(5), receiver)
+ .await
+ .expect("vindex range read hung on a current-thread
runtime")
+ .unwrap();
+ result.unwrap();
+ search.join().unwrap();
+ assert_eq!(output, [11u8; 32]);
+ });
+ });
+ worker.join().unwrap();
+ }
+}
diff --git a/crates/paimon/src/vindex/reader.rs
b/crates/paimon/src/vindex/reader.rs
index 74551154..ecc39a57 100644
--- a/crates/paimon/src/vindex/reader.rs
+++ b/crates/paimon/src/vindex/reader.rs
@@ -20,17 +20,63 @@ use paimon_vindex_core::distance::MetricType;
use paimon_vindex_core::index::{
VectorIndexMetadata, VectorIndexReader as VIndexReader, VectorSearchParams,
};
+use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities};
use std::collections::BinaryHeap;
use std::collections::HashMap;
-use std::io::{Cursor, Read, Seek, SeekFrom};
+use std::io;
const DEFAULT_NPROBE: usize = 16;
const NPROBE_PARAMETER: &str = "ivf.nprobe";
+trait ErasedSeekRead: Send {
+ fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) ->
io::Result<()>;
+
+ fn try_clone_erased(&self) -> io::Result<Option<Box<dyn ErasedSeekRead>>>;
+
+ fn capabilities_erased(&self) -> SeekReadCapabilities;
+}
+
+impl<T: SeekRead + 'static> ErasedSeekRead for T {
+ fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) ->
io::Result<()> {
+ SeekRead::pread(self, ranges)
+ }
+
+ fn try_clone_erased(&self) -> io::Result<Option<Box<dyn ErasedSeekRead>>> {
+ Ok(SeekRead::try_clone_reader(self)?
+ .map(|reader| Box::new(reader) as Box<dyn ErasedSeekRead>))
+ }
+
+ fn capabilities_erased(&self) -> SeekReadCapabilities {
+ SeekRead::read_capabilities(self)
+ }
+}
+
+struct VindexInput(Box<dyn ErasedSeekRead>);
+
+impl VindexInput {
+ fn new<S: SeekRead + 'static>(source: S) -> Self {
+ Self(Box::new(source))
+ }
+}
+
+impl SeekRead for VindexInput {
+ fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> {
+ self.0.pread_erased(ranges)
+ }
+
+ fn try_clone_reader(&self) -> io::Result<Option<Self>> {
+ Ok(self.0.try_clone_erased()?.map(Self))
+ }
+
+ fn read_capabilities(&self) -> SeekReadCapabilities {
+ self.0.capabilities_erased()
+ }
+}
+
pub struct VindexVectorGlobalIndexReader {
io_meta: GlobalIndexIOMeta,
options: HashMap<String, String>,
- reader: Option<VIndexReader<Cursor<Vec<u8>>>>,
+ reader: Option<VIndexReader<VindexInput>>,
metadata: Option<VectorIndexMetadata>,
}
@@ -44,7 +90,7 @@ impl VindexVectorGlobalIndexReader {
}
}
- pub fn visit_vector_search<S: Read + Seek + Send + 'static>(
+ pub fn visit_vector_search<S: SeekRead + 'static>(
&mut self,
vector_search: &VectorSearch,
stream_fn: impl FnOnce(&str) -> crate::Result<S>,
@@ -53,7 +99,7 @@ impl VindexVectorGlobalIndexReader {
self.search(vector_search)
}
- pub fn visit_batch_vector_search<S: Read + Seek + Send + 'static>(
+ pub fn visit_batch_vector_search<S: SeekRead + 'static>(
&mut self,
vector_searches: &[VectorSearch],
stream_fn: impl FnOnce(&str) -> crate::Result<S>,
@@ -84,7 +130,7 @@ impl VindexVectorGlobalIndexReader {
search_vindex(reader, metadata, &self.options, vector_search)
}
- fn ensure_loaded<S: Read + Seek + Send + 'static>(
+ fn ensure_loaded<S: SeekRead + 'static>(
&mut self,
stream_fn: impl FnOnce(&str) -> crate::Result<S>,
) -> crate::Result<()> {
@@ -92,26 +138,13 @@ impl VindexVectorGlobalIndexReader {
return Ok(());
}
- let mut stream = stream_fn(&self.io_meta.file_path)?;
- stream
- .seek(SeekFrom::Start(0))
- .map_err(|e| crate::Error::UnexpectedError {
- message: format!("Failed to seek vindex stream to start: {}",
e),
- source: Some(Box::new(e)),
- })?;
- let mut bytes = Vec::with_capacity(self.io_meta.file_size as usize);
- stream
- .read_to_end(&mut bytes)
- .map_err(|e| crate::Error::UnexpectedError {
- message: format!("Failed to read vindex stream: {}", e),
- source: Some(Box::new(e)),
- })?;
-
- let mut reader =
- VIndexReader::open(Cursor::new(bytes)).map_err(|e|
crate::Error::DataInvalid {
+ let source = stream_fn(&self.io_meta.file_path)?;
+ let mut reader =
VIndexReader::open(VindexInput::new(source)).map_err(|e| {
+ crate::Error::DataInvalid {
message: format!("Failed to open paimon-vindex-core reader:
{}", e),
source: Some(Box::new(e)),
- })?;
+ }
+ })?;
let metadata = reader.metadata();
reader
.optimize_for_search()
@@ -127,7 +160,7 @@ impl VindexVectorGlobalIndexReader {
}
fn search_vindex(
- reader: &mut VIndexReader<Cursor<Vec<u8>>>,
+ reader: &mut VIndexReader<impl SeekRead>,
metadata: &VectorIndexMetadata,
options: &HashMap<String, String>,
vector_search: &VectorSearch,
@@ -266,6 +299,102 @@ fn int_parameter(
#[cfg(test)]
mod tests {
use super::*;
+ use crate::io::FileRead;
+ use crate::vindex::range_reader::VindexFileReader;
+ use async_trait::async_trait;
+ use bytes::Bytes;
+ use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer,
VectorIndexWriter};
+ use paimon_vindex_core::io::{PosWriter, SeekReadCapabilities};
+ use std::ops::Range;
+ use std::sync::atomic::{AtomicUsize, Ordering};
+ use std::sync::{Arc, Mutex};
+
+ const TEST_DIMENSION: usize = 16;
+
+ #[derive(Clone)]
+ struct CloneableSeekRead {
+ capabilities: SeekReadCapabilities,
+ }
+
+ impl SeekRead for CloneableSeekRead {
+ fn pread(&mut self, _ranges: &mut [ReadRequest<'_>]) -> io::Result<()>
{
+ Ok(())
+ }
+
+ fn try_clone_reader(&self) -> io::Result<Option<Self>> {
+ Ok(Some(self.clone()))
+ }
+
+ fn read_capabilities(&self) -> SeekReadCapabilities {
+ self.capabilities
+ }
+ }
+
+ struct TrackingIndexRead {
+ data: Bytes,
+ ranges: Mutex<Vec<Range<u64>>>,
+ bytes_read: AtomicUsize,
+ }
+
+ impl TrackingIndexRead {
+ fn new(data: Bytes) -> Arc<Self> {
+ Arc::new(Self {
+ data,
+ ranges: Mutex::new(Vec::new()),
+ bytes_read: AtomicUsize::new(0),
+ })
+ }
+
+ fn ranges(&self) -> Vec<Range<u64>> {
+ self.ranges.lock().unwrap().clone()
+ }
+ }
+
+ #[async_trait]
+ impl FileRead for TrackingIndexRead {
+ async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
+ self.bytes_read
+ .fetch_add((range.end - range.start) as usize,
Ordering::SeqCst);
+ self.ranges.lock().unwrap().push(range.clone());
+ Ok(self.data.slice(range.start as usize..range.end as usize))
+ }
+ }
+
+ fn build_ivf_flat_index() -> Bytes {
+ let vector_count = 8192usize;
+ let mut vectors = Vec::with_capacity(vector_count * TEST_DIMENSION);
+ for row in 0..vector_count {
+ let cluster = (row % 16) as f32 * 100.0;
+ for dimension in 0..TEST_DIMENSION {
+ vectors.push(cluster + dimension as f32 * 0.01 + row as f32 *
0.000001);
+ }
+ }
+ let ids: Vec<i64> = (0..vector_count as i64).collect();
+ let options = HashMap::from([
+ ("index.type".to_string(), "ivf_flat".to_string()),
+ ("dimension".to_string(), TEST_DIMENSION.to_string()),
+ ("nlist".to_string(), "16".to_string()),
+ ("metric".to_string(), "l2".to_string()),
+ ]);
+ let config = VectorIndexConfig::from_options(&options).unwrap();
+ let training = VectorIndexTrainer::train(config, &vectors,
vector_count).unwrap();
+ let mut writer = VectorIndexWriter::new(training);
+ writer.add_vectors(&ids, &vectors, vector_count).unwrap();
+ let mut output = Vec::new();
+ writer.write(&mut PosWriter::new(&mut output)).unwrap();
+ Bytes::from(output)
+ }
+
+ fn query() -> VectorSearch {
+ VectorSearch::new(
+ (0..TEST_DIMENSION)
+ .map(|dimension| dimension as f32 * 0.01)
+ .collect(),
+ 10,
+ "embedding".to_string(),
+ )
+ .unwrap()
+ }
#[test]
fn test_convert_distance_to_score() {
@@ -317,4 +446,58 @@ mod tests {
options.insert(NPROBE_PARAMETER.to_string(), "abc".to_string());
assert!(int_parameter(&options, NPROBE_PARAMETER,
DEFAULT_NPROBE).is_err());
}
+
+ #[test]
+ fn erased_input_forwards_clone_and_capabilities() {
+ let capabilities = SeekReadCapabilities {
+ estimated_random_read_latency_nanos: 123,
+ preferred_window_bytes: 64 * 1024,
+ max_ranges_per_pread: 7,
+ };
+ let input = VindexInput::new(CloneableSeekRead { capabilities });
+
+ assert_eq!(input.read_capabilities(), capabilities);
+ let cloned = input.try_clone_reader().unwrap().unwrap();
+ assert_eq!(cloned.read_capabilities(), capabilities);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn scalar_search_range_reads_instead_of_loading_the_whole_index() {
+ let index = build_ivf_flat_index();
+ let tracking = TrackingIndexRead::new(index.clone());
+ let source: Arc<dyn FileRead> = tracking.clone();
+ let index_size = index.len();
+ let runtime = tokio::runtime::Handle::current();
+
+ let result = tokio::task::spawn_blocking(move || {
+ let source = VindexFileReader::new(
+ source,
+ runtime,
+ index_size as u64,
+ "scalar.index".to_string(),
+ );
+ let io_meta =
+ GlobalIndexIOMeta::new("scalar.index".to_string(), index_size
as u64, Vec::new());
+ let options = HashMap::from([(NPROBE_PARAMETER.to_string(),
"1".to_string())]);
+ let mut reader = VindexVectorGlobalIndexReader::new(io_meta,
options);
+ reader.visit_vector_search(&query(), |_| Ok(source))
+ })
+ .await
+ .unwrap()
+ .unwrap();
+
+ assert!(result.is_some());
+ let bytes_read = tracking.bytes_read.load(Ordering::SeqCst);
+ assert!(
+ bytes_read < index_size / 2,
+ "nprobe=1 should read substantially less than the full index:
read={bytes_read}, file={index_size}"
+ );
+ assert!(
+ tracking
+ .ranges()
+ .iter()
+ .all(|range| range.start != 0 || range.end != index_size as
u64),
+ "range search unexpectedly read the entire index"
+ );
+ }
}
