From d9d802bbc1987f26498d8056e144eb72cd3bd558 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 20:24:38 +0800 Subject: [PATCH 1/6] [core] Add scheduler execution for snapshot-scoped search plans --- .../python/pypaimon_rust/datafusion.pyi | 8 +- bindings/python/src/full_text_search.rs | 36 + bindings/python/src/hybrid_search.rs | 19 +- bindings/python/src/lib.rs | 1 + bindings/python/src/search_task.rs | 181 ++++ bindings/python/src/table.rs | 12 + bindings/python/src/vector_search.rs | 41 +- crates/paimon/src/table/de_vector_read.rs | 848 +++++++++++++----- .../paimon/src/table/de_vector_read/tests.rs | 378 ++++++++ .../src/table/full_text_search_builder.rs | 193 +++- .../full_text_search_builder/execution.rs | 225 +++++ .../full_text_search_builder/parity_tests.rs | 256 ++++++ .../table/full_text_search_builder/read.rs | 26 +- .../paimon/src/table/hybrid_search_builder.rs | 55 +- .../hybrid_search_builder/parity_tests.rs | 77 ++ crates/paimon/src/table/mod.rs | 2 + crates/paimon/src/table/search_task.rs | 447 +++++++++ crates/paimon/src/table/vector_read.rs | 33 +- 18 files changed, 2555 insertions(+), 283 deletions(-) create mode 100644 bindings/python/src/search_task.rs create mode 100644 crates/paimon/src/table/full_text_search_builder/execution.rs create mode 100644 crates/paimon/src/table/search_task.rs diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 4d4dcefb6..407d0cad1 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -321,6 +321,7 @@ class PartitionStat: def total_size_bytes(self) -> int: ... class Table: + def _execute_search_task(self, task: bytes, context: bytes) -> bytes: ... @staticmethod def from_resolved_schema( location: str, schema_json: str, *, database: str = "default", @@ -946,10 +947,10 @@ class VectorScanPlan: def snapshot_id(self) -> Optional[int]: ... class VectorRead: - def read_plan(self, plan: VectorScanPlan) -> "GlobalIndexResult": ... + def read_plan(self, plan: VectorScanPlan, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> "GlobalIndexResult": ... class BatchVectorRead: - def read_batch_plan(self, plan: VectorScanPlan) -> List["GlobalIndexResult"]: ... + def read_batch_plan(self, plan: VectorScanPlan, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> List["GlobalIndexResult"]: ... class GlobalIndexResultReadBuilder: def with_projection(self, columns: List[str]) -> "GlobalIndexResultReadBuilder": ... @@ -972,6 +973,7 @@ class HybridSearchBuilder: def add_full_text_route(self, field_name: str, query: str, limit: int, weight: float = 1.0, options: Optional[Dict[str, str]] = None) -> HybridSearchBuilder: ... def execute_local(self) -> GlobalIndexResult: ... + def execute(self, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> GlobalIndexResult: ... class FullTextSearchBuilder: def with_query(self, field_name: str, query: str) -> FullTextSearchBuilder: ... @@ -990,7 +992,7 @@ class FullTextScanPlan: def snapshot_id(self) -> Optional[int]: ... class FullTextRead: - def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult: ... + def read_plan(self, plan: FullTextScanPlan, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> GlobalIndexResult: ... class ShardTableUpdator: def arrow_reader(self) -> RecordBatchReader: diff --git a/bindings/python/src/full_text_search.rs b/bindings/python/src/full_text_search.rs index acdc76e94..c102ac7b7 100644 --- a/bindings/python/src/full_text_search.rs +++ b/bindings/python/src/full_text_search.rs @@ -129,6 +129,30 @@ impl PyFullTextSearchBuilder { .map(PyGlobalIndexResult::new) .map_err(to_py_err) } + + #[pyo3(signature = (executor=None))] + fn execute( + &self, + py: Python<'_>, + executor: Option>, + ) -> PyResult { + let Some(callback) = executor else { + return self.execute_local(py); + }; + let builder = self.config.builder(&self.table).map_err(to_py_err)?; + let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); + let result = py.detach(|| { + runtime().block_on(async { + let read = builder.new_read()?; + let plan = builder.new_scan()?.scan().await?; + read.read_with_executor(plan, &executor).await + }) + }); + if let Some(error) = executor.take_error() { + return Err(error); + } + result.map(PyGlobalIndexResult::new).map_err(to_py_err) + } } #[pyclass(name = "FullTextScan", module = "pypaimon_rust.datafusion")] @@ -172,11 +196,23 @@ pub struct PyFullTextRead { } #[pymethods] impl PyFullTextRead { + #[pyo3(signature = (plan, executor=None))] fn read_plan( &self, py: Python<'_>, plan: &PyFullTextScanPlan, + executor: Option>, ) -> PyResult { + if let Some(callback) = executor { + let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); + let result = py.detach(|| { + runtime().block_on(self.inner.read_with_executor(plan.inner.clone(), &executor)) + }); + if let Some(error) = executor.take_error() { + return Err(error); + } + return result.map(PyGlobalIndexResult::new).map_err(to_py_err); + } py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) .map(PyGlobalIndexResult::new) .map_err(to_py_err) diff --git a/bindings/python/src/hybrid_search.rs b/bindings/python/src/hybrid_search.rs index 0370b585a..c43fb7cc7 100644 --- a/bindings/python/src/hybrid_search.rs +++ b/bindings/python/src/hybrid_search.rs @@ -177,10 +177,27 @@ impl PyHybridSearchBuilder { Ok(slf) } - fn execute_local(&self, py: Python<'_>) -> PyResult { + #[pyo3(signature = (executor=None))] + fn execute( + &self, + py: Python<'_>, + executor: Option>, + ) -> PyResult { let builder = self.builder(); + if let Some(callback) = executor { + let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); + let result = py.detach(|| runtime().block_on(builder.execute_with_executor(&executor))); + if let Some(error) = executor.take_error() { + return Err(error); + } + return result.map(PyGlobalIndexResult::new).map_err(to_py_err); + } py.detach(|| runtime().block_on(builder.execute_scored())) .map(PyGlobalIndexResult::new) .map_err(to_py_err) } + + fn execute_local(&self, py: Python<'_>) -> PyResult { + self.execute(py, None) + } } diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index d579ee28a..0b5ee863b 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -34,6 +34,7 @@ mod oss_cpp; mod predicate; mod read; mod schema; +mod search_task; mod shard_update; mod system_table; mod table; diff --git a/bindings/python/src/search_task.rs b/bindings/python/src/search_task.rs new file mode 100644 index 000000000..b265b82dd --- /dev/null +++ b/bindings/python/src/search_task.rs @@ -0,0 +1,181 @@ +// 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 paimon::table::{SearchTask, SearchTaskExecutor, SearchTaskResult}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; +use std::sync::Mutex; + +pub(crate) struct PythonSearchTaskExecutor { + callback: Py, + error: Mutex>, +} + +impl PythonSearchTaskExecutor { + pub(crate) fn new(callback: Py) -> Self { + Self { + callback, + error: Mutex::new(None), + } + } + pub(crate) fn take_error(&self) -> Option { + self.error.lock().unwrap().take() + } +} + +#[async_trait::async_trait] +impl SearchTaskExecutor for PythonSearchTaskExecutor { + async fn execute(&self, tasks: Vec) -> paimon::Result> { + let (context, bytes) = SearchTask::serialize_batch(&tasks)?; + let callback = Python::attach(|py| self.callback.clone_ref(py)); + // A callback may wait for remote workers. Yield this route's future so + // the hybrid driver can poll and dispatch its other routes as well. + let result: PyResult> = tokio::task::spawn_blocking(move || { + Python::attach(|py| { + let input = PyList::new(py, bytes.iter().map(|task| PyBytes::new(py, task)))?; + let output = callback + .bind(py) + .call1((PyBytes::new(py, &context), input))?; + let result = output + .try_iter()? + .map(|value| { + let bytes: Vec = value?.extract()?; + SearchTaskResult::deserialize(&bytes).map_err(crate::error::to_py_err) + }) + .collect::>>(); + if result.is_err() { + // A retained scheduler generator still owns pending workers. + // Run its finally block, preserving the primary Python error. + if let Ok(close) = output.getattr("close") { + let _ = close.call0(); + } + } + result + }) + }) + .await + .map_err(|error| paimon::Error::UnexpectedError { + message: "Search task callback join failed".into(), + source: Some(Box::new(error)), + })?; + result.map_err(|error| { + let mut first = self.error.lock().unwrap(); + if first.is_none() { + *first = Some(error); + } + paimon::Error::UnexpectedError { + message: "Search task executor failed".into(), + source: None, + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::types::PyDict; + + #[tokio::test(flavor = "current_thread")] + async fn scheduled_python_callbacks_yield_for_other_routes() { + let callback = Python::attach(|py| { + let globals = PyDict::new(py); + py.run( + c" +import threading +barrier = threading.Barrier(2) +def execute(context, tasks): + barrier.wait(timeout=5) + return [] +", + Some(&globals), + None, + ) + .unwrap(); + globals.get_item("execute").unwrap().unwrap().unbind() + }); + let executor = PythonSearchTaskExecutor::new(callback); + // Both routes must dispatch before either finishes, even with only one + // async runtime thread. A blocking inline callback breaks the barrier. + let (left, right) = + futures::join!(executor.execute(Vec::new()), executor.execute(Vec::new())); + assert!(left.unwrap().is_empty()); + assert!(right.unwrap().is_empty()); + assert!(executor.take_error().is_none()); + } + + #[tokio::test] + async fn scheduled_python_executor_preserves_cancellation_exception_identity() { + let (callback, failure) = Python::attach(|py| { + let globals = PyDict::new(py); + py.run( + c" +failure = KeyboardInterrupt('cancelled search') +def execute(context, tasks): + raise failure +", + Some(&globals), + None, + ) + .unwrap(); + ( + globals.get_item("execute").unwrap().unwrap().unbind(), + globals.get_item("failure").unwrap().unwrap().unbind(), + ) + }); + let executor = PythonSearchTaskExecutor::new(callback); + assert!(executor.execute(Vec::new()).await.is_err()); + let error = executor.take_error().unwrap(); + Python::attach(|py| assert!(error.value(py).is(failure.bind(py)))); + } + + #[tokio::test] + async fn scheduled_python_executor_closes_result_iterator_after_decode_error() { + let (callback, globals) = Python::attach(|py| { + let globals = PyDict::new(py); + py.run( + c" +closed = [] +def execute(context, tasks): + try: + yield b'invalid task result' + finally: + closed.append(True) +", + Some(&globals), + None, + ) + .unwrap(); + ( + globals.get_item("execute").unwrap().unwrap().unbind(), + globals.unbind(), + ) + }); + let executor = PythonSearchTaskExecutor::new(callback); + assert!(executor.execute(Vec::new()).await.is_err()); + Python::attach(|py| { + let closed: Vec = globals + .bind(py) + .get_item("closed") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!(closed, vec![true]); + }); + } +} diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs index f0e290a04..69ef1e83c 100644 --- a/bindings/python/src/table.rs +++ b/bindings/python/src/table.rs @@ -49,6 +49,18 @@ impl PyTable { #[pymethods] impl PyTable { + fn _execute_search_task<'py>( + &self, + py: Python<'py>, + task: &[u8], + context: &[u8], + ) -> PyResult> { + let result = py + .detach(|| runtime().block_on(self.inner.execute_search_task(task, context))) + .map_err(to_py_err)?; + Ok(pyo3::types::PyBytes::new(py, &result)) + } + /// Construct a filesystem table from a Java-format TableSchema JSON document. /// /// Preserves the caller's resolved schema and complete table options without diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index 8b7d42904..1c2d19245 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -367,10 +367,27 @@ pub struct PyVectorRead { #[pymethods] impl PyVectorRead { - fn read_plan(&self, py: Python<'_>, plan: &PyVectorScanPlan) -> PyResult { - py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - .map(PyGlobalIndexResult::new) - .map_err(to_py_err) + #[pyo3(signature = (plan, executor=None))] + fn read_plan( + &self, + py: Python<'_>, + plan: &PyVectorScanPlan, + executor: Option>, + ) -> PyResult { + let Some(callback) = executor else { + return py + .detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + .map(PyGlobalIndexResult::new) + .map_err(to_py_err); + }; + let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); + let result = py.detach(|| { + runtime().block_on(self.inner.read_with_executor(plan.inner.clone(), &executor)) + }); + if let Some(error) = executor.take_error() { + return Err(error); + } + result.map(PyGlobalIndexResult::new).map_err(to_py_err) } } @@ -381,12 +398,26 @@ pub struct PyBatchVectorRead { #[pymethods] impl PyBatchVectorRead { + #[pyo3(signature = (plan, executor=None))] fn read_batch_plan( &self, py: Python<'_>, plan: &PyVectorScanPlan, + executor: Option>, ) -> PyResult> { - py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + let result = if let Some(callback) = executor { + let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); + let result = py.detach(|| { + runtime().block_on(self.inner.read_with_executor(plan.inner.clone(), &executor)) + }); + if let Some(error) = executor.take_error() { + return Err(error); + } + result + } else { + py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + }; + result .map(|results| results.into_iter().map(PyGlobalIndexResult::new).collect()) .map_err(to_py_err) } diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index a96f2f8be..de3513eda 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -34,6 +34,9 @@ use crate::table::global_index_types::normalize_queryable_global_index_type; use crate::table::index_file_path::IndexFileLocation; use crate::table::pk_vector_position_read::SEARCH_SCORE_COLUMN; use crate::table::row_id_predicate::intersect_sorted_ranges; +use crate::table::search_task::{ + execute_tasks, SearchTask, SearchTaskExecutor, SEARCH_TASK_WINDOW, +}; use crate::table::vector_read::Read; use crate::table::vector_search_common::{ configured_refine_factor, configured_vector_index_type, indexed_search_limit, @@ -60,6 +63,7 @@ use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::VectorIndexReader as VIndexReader; use paimon_vindex_core::io::SeekRead; use roaring::RoaringTreemap; +use serde::{Deserialize, Serialize}; use std::borrow::Cow; use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -187,6 +191,23 @@ impl Read for DeVectorRead { type Plan = DeVectorScanPlan; async fn read(&self, plan: DeVectorScanPlan) -> crate::Result> { + self.read_with_executor(plan, None).await + } +} + +impl DeVectorRead { + pub(super) async fn read_with_executor( + &self, + plan: DeVectorScanPlan, + executor: Option<&dyn SearchTaskExecutor>, + ) -> crate::Result> { + if executor.is_some() + && !CoreOptions::new(plan.table.schema().options()).data_evolution_enabled() + { + return Err(crate::Error::Unsupported { + message: "Scheduled vector search requires data evolution".into(), + }); + } let DeVectorScanPlan { table, index_entries, @@ -227,7 +248,7 @@ impl Read for DeVectorRead { search.set_shared_include_row_ids(Arc::clone(&include_row_ids)); } } - evaluate_batch_vector_search( + evaluate_batch_vector_search_with_executor( VectorSearchEvaluation { table: Some(pinned_table), file_io: pinned_table.file_io(), @@ -239,6 +260,7 @@ impl Read for DeVectorRead { }, &index_entries, &vector_searches, + executor, ) .await? }; @@ -271,6 +293,195 @@ impl Read for DeVectorRead { } } +async fn execute_vector_index_work( + evaluation: VectorSearchEvaluation<'_>, + work: VectorIndexWork, + metric_check: VectorMetricCheck, + range_read_limiter: Option, + batch_index_parallelism: usize, +) -> crate::Result<(Vec, IndexSearchTiming)> { + let timing_enabled = vector_search_timing_enabled(); + let concurrency = CoreOptions::new(evaluation.table_options).global_index_thread_num()?; + let global_meta = work + .entry + .index_file + .global_index_meta + .as_ref() + .ok_or_else(|| { + crate::table::search_task::invalid("Vector task is missing index metadata") + })?; + let backend = VectorIndexBackend::from_index_type(&work.entry.index_file.index_type) + .ok_or_else(|| crate::table::search_task::invalid("Unsupported vector task index"))?; + let table_path = evaluation.table_path.trim_end_matches('/'); + let path = IndexFileLocation::Global { table_path }.resolve( + &work.entry.index_file.file_name, + work.entry.index_file.external_path.as_deref(), + ); + let file_name = work.entry.index_file.file_name.clone(); + let file_size = work.entry.index_file.file_size as u64; + let index_meta_bytes = global_meta.index_meta.clone().unwrap_or_default(); + let row_range_start = global_meta.row_range_start; + let vector_searches = work.queries; + let mut options = evaluation.table_options.clone(); + options.extend(vector_searches[0].options.clone()); + let input = evaluation.file_io.new_input(&path); + if vector_searches.iter().all(|search| { + search + .effective_include_row_ids() + .is_some_and(|row_ids| row_ids.is_empty()) + }) { + return Ok(( + vec![ScoredRowIds::empty(); vector_searches.len()], + IndexSearchTiming::default(), + )); + } + let permit_start = timing_enabled.then(Instant::now); + let permit = acquire_process_global_search_permit(concurrency).await?; + let permit_wait = permit_start.map_or(Duration::ZERO, |start| start.elapsed()); + let input = input?; + let query_count = vector_searches.len(); + let mut file_reader_open = Duration::ZERO; + let mut full_file_read = None; + let io_meta = GlobalIndexIOMeta::new(file_name.clone(), file_size, index_meta_bytes); + let results = match backend { + VectorIndexBackend::Lumina => { + let read_start = timing_enabled.then(Instant::now); + let data = input.read().await.map_err(|e| crate::Error::DataInvalid { + message: format!( + "Failed to read {} index file '{}': {}", + backend.error_name(), + file_name, + e + ), + source: None, + })?; + if let Some(start) = read_start { + full_file_read = Some((start.elapsed(), data.len())); + } + let metric = LuminaIndexMeta::deserialize(&io_meta.metadata)?.metric()?; + metric_check.check(RawVectorMetric::from_lumina(metric))?; + execute_global_index_with_guard( + "Lumina global-index batch search task failed", + permit, + move || { + let mut reader = LuminaVectorGlobalIndexReader::new(io_meta, options); + reader.visit_batch_vector_search(&vector_searches, |_| Ok(Cursor::new(data))) + }, + ) + .await? + } + VectorIndexBackend::Vindex => match tokio::runtime::Handle::try_current() { + Ok(runtime) => { + let file_reader_open_start = timing_enabled.then(Instant::now); + let file_reader = input + .reader() + .await + .map_err(|e| crate::Error::DataInvalid { + message: format!( + "Failed to open vindex file '{}' for range reads: {}", + file_name, e + ), + source: None, + })?; + file_reader_open = + file_reader_open_start.map_or(Duration::ZERO, |start| start.elapsed()); + let source = VindexFileReader::new_with_limiter( + Arc::new(file_reader), + runtime, + range_read_limiter.expect("Vindex range-read limiter"), + file_size, + file_name.clone(), + ); + let range_io_stats = source.range_io_stats(); + let results = execute_vindex_searches( + io_meta, + options, + vector_searches, + source, + file_name.clone(), + VectorIndexSearchContext { + parallelism: batch_index_parallelism, + metrics: metric_check, + }, + permit, + ) + .await?; + if let Some(stats) = range_io_stats { + log_vindex_range_io_stats(&file_name, query_count, &stats); + } + results + } + Err(_) if query_count > 1 => { + let read_start = timing_enabled.then(Instant::now); + let data = input.read().await.map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to read vindex index file '{}': {}", file_name, e), + source: None, + })?; + if let Some(start) = read_start { + full_file_read = Some((start.elapsed(), data.len())); + } + execute_vindex_searches( + io_meta, + options, + vector_searches, + Cursor::new(data), + file_name.clone(), + VectorIndexSearchContext { + parallelism: batch_index_parallelism, + metrics: metric_check, + }, + permit, + ) + .await? + } + Err(error) => { + return Err(crate::Error::UnexpectedError { + message: "Vector index range reader requires a Tokio runtime".to_string(), + source: Some(Box::new(error)), + }); + } + }, + }; + if let Some((read, returned_bytes)) = full_file_read { + log::debug!( + target: "paimon::vector_search", + "event=paimon_vector_full_file_io backend={} file={} nq={} requested_bytes={} returned_bytes={} read_ms={:.3}", + backend.error_name(), + file_name, + query_count, + file_size, + returned_bytes, + read.as_secs_f64() * 1000.0, + ); + } + if results.len() != query_count { + return Err(crate::Error::DataInvalid { + message: format!( + "Batch vector search backend returned {} results for {} query vectors", + results.len(), + query_count + ), + source: None, + }); + } + + Ok::<_, crate::Error>(( + results + .into_iter() + .map(|result| match result { + Some(scored_map) => { + ScoredRowIds::from_scored_map(scored_map).offset(row_range_start) + } + None => ScoredRowIds::empty(), + }) + .collect::>(), + IndexSearchTiming { + permit_wait, + file_reader_open, + }, + )) +} + const RAW_SCORE_MATRIX_MIN_QUERY_COUNT: usize = 4; const RAW_SCORE_MATRIX_TARGET_ELEMENTS: usize = 1 << 20; @@ -461,10 +672,21 @@ async fn scalar_matched_row_ids( )) } +#[cfg(test)] async fn evaluate_batch_vector_search( evaluation: VectorSearchEvaluation<'_>, index_entries: &[IndexManifestEntry], vector_searches: &[VectorSearch], +) -> crate::Result> { + evaluate_batch_vector_search_with_executor(evaluation, index_entries, vector_searches, None) + .await +} + +async fn evaluate_batch_vector_search_with_executor( + evaluation: VectorSearchEvaluation<'_>, + index_entries: &[IndexManifestEntry], + vector_searches: &[VectorSearch], + executor: Option<&dyn SearchTaskExecutor>, ) -> crate::Result> { let timing_enabled = vector_search_timing_enabled(); let total_start = timing_enabled.then(Instant::now); @@ -682,230 +904,104 @@ async fn evaluate_batch_vector_search( Some(RangeReadLimiter::new(range_read_concurrency)), ) }; - let futures: Vec<_> = vector_search_plans + let mut works = vector_search_plans .into_iter() .map(|(entry, shared_local_filter)| { - let range_read_limiter = range_read_limiter.clone(); - let global_meta = entry.index_file.global_index_meta.as_ref().unwrap(); - let backend = VectorIndexBackend::from_index_type(&entry.index_file.index_type) - .expect("filtered vector index type"); - let path = IndexFileLocation::Global { table_path } - .resolve(&entry.index_file.file_name, entry.index_file.external_path.as_deref()); - let file_name = entry.index_file.file_name.clone(); - let file_size = entry.index_file.file_size as u64; - let index_meta_bytes = global_meta.index_meta.clone().unwrap_or_default(); - let row_range_start = global_meta.row_range_start; - let row_range_end = global_meta.row_range_end; + let meta = entry.index_file.global_index_meta.as_ref().unwrap(); let index_limit = search_limit_with_deleted_rows( index_search_limit, - row_range_start, - row_range_end, + meta.row_range_start, + meta.row_range_end, deleted_row_index.as_ref(), ) .min(i32::MAX as usize); - let mut vector_searches = indexed_searches.to_vec(); - for vector_search in &mut vector_searches { - vector_search.limit = index_limit; - } - let mut options = evaluation.table_options.clone(); - options.extend(search_options.clone()); - let input = evaluation.file_io.new_input(&path); - let metric_check = metric_check.clone(); - async move { - if let Some(local_filter) = shared_local_filter { - let local_filter = Arc::new(local_filter); - for vector_search in &mut vector_searches { - vector_search - .set_shared_include_row_ids(Arc::clone(&local_filter)); - } - } else { - for vector_search in &mut vector_searches { - if let Some(include_row_ids) = - vector_search.effective_include_row_ids() - { - vector_search.set_shared_include_row_ids(Arc::new( - localize_include_row_ids( - include_row_ids, - row_range_start, - row_range_end, - )?, - )); - } - } - } - if vector_searches.iter().all(|search| { - search + let mut searches = indexed_searches.to_vec(); + let shared_local_filter = shared_local_filter.map(Arc::new); + for search in &mut searches { + search.limit = index_limit; + let local = match shared_local_filter.as_ref() { + Some(local) => Some(Arc::clone(local)), + None => search .effective_include_row_ids() - .is_some_and(|row_ids| row_ids.is_empty()) - }) { - return Ok(( - vec![ScoredRowIds::empty(); vector_searches.len()], - IndexSearchTiming::default(), - )); - } - let permit_start = timing_enabled.then(Instant::now); - let permit = acquire_process_global_search_permit(concurrency).await?; - let permit_wait = - permit_start.map_or(Duration::ZERO, |start| start.elapsed()); - let input = input?; - let query_count = vector_searches.len(); - let mut file_reader_open = Duration::ZERO; - let mut full_file_read = None; - let io_meta = - GlobalIndexIOMeta::new(file_name.clone(), file_size, index_meta_bytes); - let results = match backend { - VectorIndexBackend::Lumina => { - let read_start = timing_enabled.then(Instant::now); - let data = input.read().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to read {} index file '{}': {}", - backend.error_name(), - file_name, - e - ), - source: None, - } - })?; - if let Some(start) = read_start { - full_file_read = Some((start.elapsed(), data.len())); - } - let metric = LuminaIndexMeta::deserialize(&io_meta.metadata)?.metric()?; - metric_check.check(RawVectorMetric::from_lumina(metric))?; - execute_global_index_with_guard( - "Lumina global-index batch search task failed", - permit, - move || { - let mut reader = - LuminaVectorGlobalIndexReader::new(io_meta, options); - reader.visit_batch_vector_search(&vector_searches, |_| { - Ok(Cursor::new(data)) - }) - }, - ) - .await? - } - VectorIndexBackend::Vindex => { - match tokio::runtime::Handle::try_current() { - Ok(runtime) => { - let file_reader_open_start = - timing_enabled.then(Instant::now); - let file_reader = input.reader().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to open vindex file '{}' for range reads: {}", - file_name, e - ), - source: None, - } - })?; - file_reader_open = file_reader_open_start - .map_or(Duration::ZERO, |start| start.elapsed()); - let source = VindexFileReader::new_with_limiter( - Arc::new(file_reader), - runtime, - range_read_limiter.expect("Vindex range-read limiter"), - file_size, - file_name.clone(), - ); - let range_io_stats = source.range_io_stats(); - let results = execute_vindex_searches( - io_meta, - options, - vector_searches, - source, - file_name.clone(), - VectorIndexSearchContext { parallelism: batch_index_parallelism, metrics: metric_check }, - permit, - ) - .await?; - if let Some(stats) = range_io_stats { - log_vindex_range_io_stats( - &file_name, - query_count, - &stats, - ); - } - results - } - Err(_) if query_count > 1 => { - let read_start = timing_enabled.then(Instant::now); - let data = input.read().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to read vindex index file '{}': {}", - file_name, e - ), - source: None, - } - })?; - if let Some(start) = read_start { - full_file_read = Some((start.elapsed(), data.len())); - } - execute_vindex_searches( - io_meta, - options, - vector_searches, - Cursor::new(data), - file_name.clone(), - VectorIndexSearchContext { parallelism: batch_index_parallelism, metrics: metric_check }, - permit, - ) - .await? - } - Err(error) => { - return Err(crate::Error::UnexpectedError { - message: - "Vector index range reader requires a Tokio runtime" - .to_string(), - source: Some(Box::new(error)), - }); - } - } - } + .map(|include| { + localize_include_row_ids( + include, + meta.row_range_start, + meta.row_range_end, + ) + }) + .transpose()? + .map(Arc::new), }; - if let Some((read, returned_bytes)) = full_file_read { - log::debug!( - target: "paimon::vector_search", - "event=paimon_vector_full_file_io backend={} file={} nq={} requested_bytes={} returned_bytes={} read_ms={:.3}", - backend.error_name(), - file_name, - query_count, - file_size, - returned_bytes, - read.as_secs_f64() * 1000.0, - ); + if let Some(local) = local { + search.set_shared_include_row_ids(local); } - if results.len() != query_count { - return Err(crate::Error::DataInvalid { - message: format!( - "Batch vector search backend returned {} results for {} query vectors", - results.len(), - query_count - ), - source: None, - }); + } + Ok::<_, crate::Error>(VectorIndexWork { + entry: entry.clone(), + queries: searches, + }) + }); + let results = if let Some(executor) = executor { + let table = evaluation.table.ok_or_else(|| { + crate::table::search_task::invalid("Remote vector search needs a table") + })?; + let mut results = Vec::new(); + loop { + let tasks = works + .by_ref() + .take(SEARCH_TASK_WINDOW) + .map(|work| { + let work = work?; + let work = VectorIndexTask { + entry: work.entry, + queries: TaskVectorQueries::capture(&work.queries)?, + }; + SearchTask::new( + table, + "vector-index", + serde_json::to_value(work).map_err(|error| { + crate::table::search_task::invalid(error.to_string()) + })?, + ) + }) + .collect::>>()?; + if tasks.is_empty() { + break; + } + for value in execute_tasks(executor, tasks).await? { + let result: VectorTaskScores = serde_json::from_value(value) + .map_err(|error| crate::table::search_task::invalid(error.to_string()))?; + if let Some(metric) = result.metric { + metric_check.check(metric)?; } - - Ok::<_, crate::Error>(( - results - .into_iter() - .map(|result| match result { - Some(scored_map) => ScoredRowIds::from_scored_map(scored_map) - .offset(row_range_start), - None => ScoredRowIds::empty(), - }) - .collect::>(), - IndexSearchTiming { - permit_wait, - file_reader_open, - }, - )) + let rows = result.restore()?; + if rows.len() != vector_searches.len() { + return Err(crate::table::search_task::invalid( + "Vector index task returned wrong query count", + )); + } + results.push((rows, IndexSearchTiming::default())); } - }) - .collect(); - - let results = drain_indexed_jobs(futures.into_iter(), concurrency).await?; + } + results + } else { + let futures = works.map(|work| { + let metric_check = metric_check.clone(); + let range_read_limiter = range_read_limiter.clone(); + async move { + execute_vector_index_work( + evaluation, + work?, + metric_check, + range_read_limiter, + batch_index_parallelism, + ) + .await + } + }); + drain_indexed_jobs(futures, concurrency).await? + }; index_search = index_search_start.map_or(Duration::ZERO, |start| start.elapsed()); let merge_start = timing_enabled.then(Instant::now); for (per_entry, entry_timing) in &results { @@ -928,6 +1024,7 @@ async fn evaluate_batch_vector_search( vector_searches, merged, index_search_limit, + executor, ) .await?; refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); @@ -995,12 +1092,13 @@ async fn evaluate_batch_vector_search( .await?; let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); metric_check.check(metric)?; - let (raw_results, raw_timing) = read_raw_batch_vector_search( + let (raw_results, raw_timing) = read_raw_batch_vector_search_with_executor( table, vector_searches, &raw_ranges, metric, evaluation.filter, + executor, ) .await?; if let Some(raw_timing) = raw_timing { @@ -1224,6 +1322,7 @@ fn attach_scores_by_row_id( Ok(vec![out]) } +#[allow(clippy::too_many_arguments)] async fn maybe_rerank_indexed_batch_results( evaluation: VectorSearchEvaluation<'_>, index_entries: &[IndexManifestEntry], @@ -1232,6 +1331,7 @@ async fn maybe_rerank_indexed_batch_results( vector_searches: &[VectorSearch], results: Vec, index_search_limit: usize, + executor: Option<&dyn SearchTaskExecutor>, ) -> crate::Result> { let timing_enabled = vector_search_timing_enabled(); let total_start = timing_enabled.then(Instant::now); @@ -1280,8 +1380,15 @@ async fn maybe_rerank_indexed_batch_results( metric_check.check(metric)?; let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); - let (results, raw_timing) = - read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric, None).await?; + let (results, raw_timing) = read_raw_batch_vector_search_with_executor( + table, + &candidate_searches, + &raw_ranges, + metric, + None, + executor, + ) + .await?; if let (Some(total_start), Some(raw_timing)) = (total_start, raw_timing) { log::debug!( target: "paimon::vector_search", @@ -1462,7 +1569,7 @@ pub(super) async fn detail_data_ranges_for_table(table: &Table) -> crate::Result Ok(merge_row_ranges(ranges)) } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] enum RawVectorMetric { L2, Cosine, @@ -1719,12 +1826,58 @@ struct RawVectorReadTiming { row_count: usize, } -async fn read_raw_batch_vector_search( +// Merge each window before scheduling more splits. Holding all split results +// would multiply the query top-K state by the entire table's split count. +async fn execute_raw_vector_tasks( + table: &Table, + splits: &[crate::table::DataSplit], + vector_searches: &[VectorSearch], + metric: RawVectorMetric, + executor: &dyn SearchTaskExecutor, +) -> crate::Result> { + let queries = TaskVectorQueries::capture(vector_searches)?; + let mut merged = vec![ScoredRowIds::empty(); vector_searches.len()]; + for window in splits.chunks(SEARCH_TASK_WINDOW) { + let tasks = window + .iter() + .map(|split| { + let work = VectorRowsTask { + splits: vec![split.clone()], + queries: queries.clone(), + metric, + }; + SearchTask::new( + table, + "vector-rows", + serde_json::to_value(work) + .map_err(|error| crate::table::search_task::invalid(error.to_string()))?, + ) + }) + .collect::>>()?; + for value in execute_tasks(executor, tasks).await? { + let result: VectorTaskScores = serde_json::from_value(value) + .map_err(|error| crate::table::search_task::invalid(error.to_string()))?; + let result = result.restore()?; + if result.len() != merged.len() { + return Err(crate::table::search_task::invalid( + "Vector task returned wrong query count", + )); + } + for ((merged, result), search) in merged.iter_mut().zip(result).zip(vector_searches) { + *merged = merged.or(&result).top_k(search.limit); + } + } + } + Ok(merged) +} + +async fn read_raw_batch_vector_search_with_executor( table: &Table, vector_searches: &[VectorSearch], raw_ranges: &[RowRange], metric: RawVectorMetric, filter: Option<&Predicate>, + executor: Option<&dyn SearchTaskExecutor>, ) -> crate::Result<(Vec, Option)> { let timing_enabled = vector_search_timing_enabled(); let total_start = timing_enabled.then(Instant::now); @@ -1755,7 +1908,7 @@ async fn read_raw_batch_vector_search( let mut read_builder = table.new_read_builder(); read_builder .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])? - .with_row_ranges(raw_ranges); + .with_row_ranges(raw_ranges.clone()); if let Some(filter) = filter { read_builder.with_filter(filter.clone()); } @@ -1777,6 +1930,31 @@ async fn read_raw_batch_vector_search( }), )); } + if let Some(executor) = executor { + let mut searches = vector_searches.to_vec(); + if let Some(filter) = filter { + // Resolve once for the whole pinned raw range, not once per worker. + // Transport row IDs rather than narrowing supported predicate types. + let matched = + matching_row_ids_for_filter(table, filter, Some(raw_ranges.clone())).await?; + let matched = Arc::new(matched); + let shared = shared_batch_include_row_ids(&searches) + .map(|include| Arc::new(include.as_ref() & matched.as_ref())); + for search in &mut searches { + let eligible = match &shared { + Some(shared) => Arc::clone(shared), + None => match search.effective_include_row_ids() { + Some(include) => Arc::new(include & matched.as_ref()), + None => Arc::clone(&matched), + }, + }; + search.set_shared_include_row_ids(eligible); + } + } + let merged = + execute_raw_vector_tasks(table, plan.splits(), &searches, metric, executor).await?; + return Ok((merged, None)); + } let read = read_builder.new_read()?; let mut stream = read.to_arrow(plan.splits())?; @@ -2439,3 +2617,235 @@ fn compute_raw_vector_score(query: &[f32], stored: &[f32], metric: RawVectorMetr #[cfg(test)] mod tests; + +#[derive(Clone, Serialize, Deserialize)] +struct TaskVectorSearch { + vector: Vec, + limit: usize, + field: String, + options: HashMap, + include: Option, +} +#[derive(Clone, Serialize, Deserialize)] +struct TaskVectorQueries { + queries: Vec, + filters: Vec>, +} +impl TaskVectorQueries { + fn capture(searches: &[VectorSearch]) -> crate::Result { + let mut filters = Vec::new(); + let mut positions = HashMap::<*const RoaringTreemap, usize>::new(); + let queries = searches + .iter() + .map(|search| { + let include = search + .effective_include_row_ids() + .map(|rows| { + let key = rows as *const RoaringTreemap; + if let Some(index) = positions.get(&key) { + return Ok::<_, crate::Error>(*index); + } + let index = filters.len(); + let mut bytes = Vec::new(); + rows.serialize_into(&mut bytes).map_err(|error| { + crate::table::search_task::invalid(error.to_string()) + })?; + filters.push(bytes); + positions.insert(key, index); + Ok(index) + }) + .transpose()?; + Ok(TaskVectorSearch { + vector: search.vector.iter().map(|value| value.to_bits()).collect(), + limit: search.limit, + field: search.field_name.clone(), + options: search.options.clone(), + include, + }) + }) + .collect::>>()?; + Ok(Self { queries, filters }) + } + fn restore(&self) -> crate::Result> { + let filters = self + .filters + .iter() + .map(|bytes| { + RoaringTreemap::deserialize_from(bytes.as_slice()) + .map(Arc::new) + .map_err(|error| crate::table::search_task::invalid(error.to_string())) + }) + .collect::>>()?; + self.queries + .iter() + .map(|query| { + let mut search = VectorSearch::new( + query + .vector + .iter() + .map(|bits| f32::from_bits(*bits)) + .collect(), + query.limit, + query.field.clone(), + )? + .with_options(query.options.clone()); + if let Some(index) = query.include { + let shared = filters.get(index).ok_or_else(|| { + crate::table::search_task::invalid( + "Vector task bitmap reference is invalid", + ) + })?; + search.set_shared_include_row_ids(Arc::clone(shared)); + } + Ok(search) + }) + .collect() + } +} +#[derive(Serialize, Deserialize)] +struct VectorIndexTask { + entry: IndexManifestEntry, + queries: TaskVectorQueries, +} +struct VectorIndexWork { + entry: IndexManifestEntry, + queries: Vec, +} +#[derive(Serialize, Deserialize)] +struct VectorRowsTask { + splits: Vec, + queries: TaskVectorQueries, + metric: RawVectorMetric, +} +#[derive(Serialize, Deserialize)] +struct TaskScoredRows { + row_ids: Vec, + score_bits: Vec, +} +#[derive(Serialize, Deserialize)] +struct VectorTaskScores { + metric: Option, + results: Vec, +} +impl VectorTaskScores { + fn capture(results: Vec, metric: Option) -> Self { + Self { + metric, + results: results + .into_iter() + .map(|result| TaskScoredRows { + row_ids: result.row_ids, + score_bits: result.scores.into_iter().map(f32::to_bits).collect(), + }) + .collect(), + } + } + fn restore(self) -> crate::Result> { + self.results + .into_iter() + .map(|result| { + if result.row_ids.len() != result.score_bits.len() { + return Err(crate::table::search_task::invalid( + "Vector task scores are not aligned", + )); + } + Ok(ScoredRowIds::new( + result.row_ids, + result.score_bits.into_iter().map(f32::from_bits).collect(), + )) + }) + .collect() + } +} + +pub(super) async fn execute_search_task( + table: &Table, + kind: &str, + payload: &serde_json::Value, +) -> crate::Result { + let scores = match kind { + "vector-index" => { + let work: VectorIndexTask = serde_json::from_value(payload.clone()) + .map_err(|error| crate::table::search_task::invalid(error.to_string()))?; + if work.queries.queries.is_empty() { + return Err(crate::table::search_task::invalid( + "Vector task has no queries", + )); + } + let metric_check = VectorMetricCheck { + persisted: Arc::new(std::sync::Mutex::new(None)), + requested: None, + }; + let evaluation = VectorSearchEvaluation { + table: Some(table), + file_io: table.file_io(), + table_path: table.location(), + table_options: table.schema().options(), + schema_fields: table.schema().fields(), + next_row_id: None, + filter: None, + }; + let core = CoreOptions::new(table.schema().options()); + let (parallelism, range_concurrency) = + vindex_concurrency_limits(&core, 1, core.global_index_thread_num()?)?; + let (results, _) = execute_vector_index_work( + evaluation, + VectorIndexWork { + entry: work.entry, + queries: work.queries.restore()?, + }, + metric_check.clone(), + Some(RangeReadLimiter::new(range_concurrency)), + parallelism, + ) + .await?; + let metric = *metric_check + .persisted + .lock() + .map_err(|_| crate::table::search_task::invalid("Vector metric lock poisoned"))?; + VectorTaskScores::capture(results, metric) + } + "vector-rows" => { + let work: VectorRowsTask = serde_json::from_value(payload.clone()) + .map_err(|error| crate::table::search_task::invalid(error.to_string()))?; + let searches = work.queries.restore()?; + if searches.is_empty() { + return Err(crate::table::search_task::invalid( + "Vector task has no queries", + )); + } + let mut builder = table.new_read_builder(); + builder.with_projection(&[searches[0].field_name.as_str(), ROW_ID_FIELD_NAME])?; + let read = builder.new_read()?; + let mut stream = read.to_arrow(&work.splits)?; + let scoring = RawScoringPlan::new(&searches, work.metric); + let mut top_k = searches + .iter() + .map(|search| RawScoreTopK::new(search.limit)) + .collect::>(); + while let Some(batch) = stream.try_next().await? { + collect_raw_batch_vector_batch( + &batch, + &searches, + work.metric, + &scoring, + &mut top_k, + )?; + } + VectorTaskScores::capture( + top_k + .into_iter() + .map(RawScoreTopK::into_search_result) + .collect(), + Some(work.metric), + ) + } + _ => { + return Err(crate::table::search_task::invalid( + "Unknown vector search task", + )) + } + }; + serde_json::to_value(scores) + .map_err(|error| crate::table::search_task::invalid(error.to_string())) +} diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs index a2718e81e..a7f56c8ad 100644 --- a/crates/paimon/src/table/de_vector_read/tests.rs +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -239,6 +239,36 @@ fn shared_raw_filter_does_not_expand_row_query_associations() { assert_eq!(plan.shared_filter_groups[0].query_indices.len(), 128); } +#[test] +fn scheduled_vector_transport_restores_one_shared_raw_filter_per_bitmap() { + let shared = Arc::new(RoaringTreemap::from_iter(0..100_000)); + let mut searches = (0..128) + .map(|_| VectorSearch::new(vec![1.0, 0.0], 2, "embedding".into()).unwrap()) + .collect::>(); + for search in &mut searches { + search.set_shared_include_row_ids(Arc::clone(&shared)); + } + let wire = TaskVectorQueries::capture(&searches).unwrap(); + assert_eq!( + wire.filters.len(), + 1, + "encode a shared bitmap once per query batch" + ); + let restored = wire.restore().unwrap(); + let restored_shared = shared_batch_include_row_ids(&restored).unwrap(); + assert_eq!(restored_shared.as_ref(), shared.as_ref()); + let scoring = RawScoringPlan::new(&restored, RawVectorMetric::L2); + assert!( + scoring.candidate_query_indices.is_empty(), + "shared filters must not expand row × query associations" + ); + assert_eq!(scoring.shared_filter_groups.len(), 1); + assert_eq!( + scoring.shared_filter_groups[0].query_indices.len(), + searches.len() + ); +} + #[test] fn shared_raw_filter_prunes_unindexed_ranges_before_reading() { let shared = Arc::new(RoaringTreemap::from_iter([7, 1_000, 1_001, 900_000])); @@ -1785,3 +1815,351 @@ fn scalar_unindexed_coverage_intersects_all_predicate_fields() { id.index_file.index_type = "btree".to_string(); } } + +struct WireSearchExecutor { + table: Table, + kinds: std::sync::Mutex>, + fail: bool, +} + +#[async_trait::async_trait] +impl crate::table::SearchTaskExecutor for WireSearchExecutor { + async fn execute( + &self, + tasks: Vec, + ) -> crate::Result> { + if self.fail { + return Err(crate::Error::UnexpectedError { + message: "scheduler cancelled".into(), + source: None, + }); + } + let (context, frames) = crate::table::SearchTask::serialize_batch(&tasks)?; + let mut results = Vec::new(); + for (task, frame) in tasks.iter().zip(frames) { + let value: serde_json::Value = serde_json::from_slice(&task.serialize()?).unwrap(); + self.kinds + .lock() + .unwrap() + .push(value["kind"].as_str().unwrap().into()); + results.push(crate::table::SearchTaskResult::deserialize( + &self.table.execute_search_task(&frame, &context).await?, + )?); + } + // Completion order must never decide index precedence or query order. + results.reverse(); + Ok(results) + } +} + +#[tokio::test] +async fn scheduled_batch_vector_matches_local_filter_refine_and_pinned_plan() { + let table = de_vector_table().await; + let mut builder = table.new_batch_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(2) + .with_filter(id_gt_filter(&table, 1)) + .with_options(HashMap::from([("refine_factor".into(), "2".into())])); + let plan = builder.new_scan().unwrap().plan().await.unwrap(); + let reader = builder.new_read().unwrap(); + let expected = reader.read(plan.clone()).await.unwrap(); + let executor = WireSearchExecutor { + table: table.clone(), + kinds: std::sync::Mutex::new(Vec::new()), + fail: false, + }; + let actual = reader.read_with_executor(plan, &executor).await.unwrap(); + for (expected, actual) in expected.iter().zip(&actual) { + assert_eq!( + expected.row_ids().unwrap().row_ids, + actual.row_ids().unwrap().row_ids + ); + assert_eq!( + expected.row_ids().unwrap().scores, + actual.row_ids().unwrap().scores + ); + } + let kinds = executor.kinds.lock().unwrap(); + assert!(kinds.iter().any(|kind| kind == "vector-index")); + assert!(kinds.iter().any(|kind| kind == "vector-rows")); +} + +#[tokio::test] +async fn scheduled_vector_failure_does_not_replay_locally() { + let table = de_vector_table().await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1); + let plan = builder.new_scan().unwrap().plan().await.unwrap(); + let reader = builder.new_read().unwrap(); + let executor = WireSearchExecutor { + table, + kinds: std::sync::Mutex::new(Vec::new()), + fail: true, + }; + let error = reader + .read_with_executor(plan, &executor) + .await + .err() + .unwrap(); + assert!(error.to_string().contains("scheduler cancelled"), "{error}"); +} + +struct CandidateSearchExecutor { + table: Table, + metric: RawVectorMetric, + index_ordinal: std::sync::atomic::AtomicUsize, + index_batch_sizes: std::sync::Mutex>, +} +#[async_trait::async_trait] +impl crate::table::SearchTaskExecutor for CandidateSearchExecutor { + async fn execute( + &self, + tasks: Vec, + ) -> crate::Result> { + let mut results = Vec::new(); + let batch_size = tasks.len(); + let mut recorded_batch = false; + for task in tasks { + let wire: serde_json::Value = serde_json::from_slice(&task.serialize()?).unwrap(); + if wire["kind"] == "vector-index" { + if !recorded_batch { + assert!(batch_size <= SEARCH_TASK_WINDOW); + self.index_batch_sizes.lock().unwrap().push(batch_size); + recorded_batch = true; + } + let ordinal = self + .index_ordinal + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let rows = if ordinal == 0 { vec![0, 2] } else { vec![1] }; + let scores = if ordinal == 0 { + vec![10.0_f32, 9.0] + } else { + vec![8.0] + }; + let payload = serde_json::to_value(VectorTaskScores::capture( + vec![ScoredRowIds::new(rows, scores)], + Some(self.metric), + )) + .unwrap(); + let result = + serde_json::json!({"version": 1, "id": wire["id"], "payload": payload}); + results.push(crate::table::SearchTaskResult::deserialize( + &serde_json::to_vec(&result).unwrap(), + )?); + } else { + results.push(task.execute(&self.table).await?); + } + } + results.reverse(); + Ok(results) + } +} + +#[tokio::test] +async fn scheduled_vector_selects_global_candidates_before_rerank() { + let table = de_vector_table().await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![0.0, 1.0]) + .with_limit(1) + .with_options(HashMap::from([("refine_factor".into(), "2".into())])); + let mut plan = builder.new_scan().unwrap().plan().await.unwrap(); + let crate::table::vector_scan::VectorScanWork::DataEvolution(work) = &mut plan.work else { + unreachable!() + }; + let duplicate = work + .index_entries + .iter() + .find(|entry| entry.index_file.index_type == IVF_FLAT_IDENTIFIER) + .unwrap() + .clone(); + for _ in 0..SEARCH_TASK_WINDOW + 1 { + work.index_entries.push(duplicate.clone()); + } + let metric = resolve_raw_vector_metric( + table.file_io(), + table.location(), + table.schema().options(), + &HashMap::new(), + &work.index_entries, + find_field_id_by_name(table.schema().fields(), "embedding").unwrap(), + "embedding", + ) + .await + .unwrap(); + let executor = CandidateSearchExecutor { + table: table.clone(), + metric, + index_ordinal: std::sync::atomic::AtomicUsize::new(0), + index_batch_sizes: std::sync::Mutex::new(Vec::new()), + }; + let reader = builder.new_read().unwrap(); + let result = reader.read_with_executor(plan, &executor).await.unwrap(); + // Global ANN candidates are 0 and 2. Exact nearest row 1 is outside that + // budget; reranking each shard before merging would incorrectly admit it. + assert_eq!(result.row_ids().unwrap().row_ids, vec![2]); + let expected = match metric { + RawVectorMetric::L2 => 0.5, + RawVectorMetric::Cosine => 1.0_f32 / 2.0_f32.sqrt(), + RawVectorMetric::InnerProduct => 1.0, + }; + assert!((result.row_ids().unwrap().scores[0] - expected).abs() < 1e-6); + assert_eq!( + *executor.index_batch_sizes.lock().unwrap(), + vec![SEARCH_TASK_WINDOW, 2] + ); +} + +struct WindowSearchExecutor { + batch_sizes: std::sync::Mutex>, + ordinal: std::sync::atomic::AtomicUsize, +} + +#[async_trait::async_trait] +impl crate::table::SearchTaskExecutor for WindowSearchExecutor { + async fn execute( + &self, + tasks: Vec, + ) -> crate::Result> { + assert!(tasks.len() <= SEARCH_TASK_WINDOW); + self.batch_sizes.lock().unwrap().push(tasks.len()); + let mut results = Vec::new(); + for task in tasks { + let ordinal = self + .ordinal + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let wire: serde_json::Value = serde_json::from_slice(&task.serialize()?).unwrap(); + let payload = VectorTaskScores::capture( + vec![ScoredRowIds::new( + vec![ordinal as u64], + vec![ordinal as f32], + )], + None, + ); + results.push(crate::table::SearchTaskResult::deserialize( + &serde_json::to_vec(&serde_json::json!({ + "version": 1, "id": wire["id"], "payload": payload + })) + .unwrap(), + )?); + } + results.reverse(); + Ok(results) + } +} + +#[tokio::test] +async fn scheduled_raw_vector_merges_bounded_windows_without_losing_late_winners() { + let table = de_vector_table().await; + let mut builder = table.new_read_builder(); + builder + .with_projection(&["embedding", ROW_ID_FIELD_NAME]) + .unwrap(); + let plan = builder.new_scan().plan().await.unwrap(); + let count = SEARCH_TASK_WINDOW * 3 + 1; + let splits = vec![plan.splits()[0].clone(); count]; + let queries = vec![VectorSearch::new(vec![1.0, 0.0], 3, "embedding".into()).unwrap()]; + let executor = WindowSearchExecutor { + batch_sizes: std::sync::Mutex::new(Vec::new()), + ordinal: std::sync::atomic::AtomicUsize::new(0), + }; + let results = + execute_raw_vector_tasks(&table, &splits, &queries, RawVectorMetric::L2, &executor) + .await + .unwrap(); + assert_eq!( + *executor.batch_sizes.lock().unwrap(), + vec![ + SEARCH_TASK_WINDOW, + SEARCH_TASK_WINDOW, + SEARCH_TASK_WINDOW, + 1 + ] + ); + assert_eq!( + results[0].row_ids, + vec![count as u64 - 1, count as u64 - 2, count as u64 - 3] + ); +} + +#[tokio::test] +async fn scheduled_raw_vector_matches_filtered_pinned_stream() { + let table = de_vector_table().await; + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let pinned = table.copy_with_pinned_snapshot(Some(&snapshot)); + let filter = id_gt_filter(&pinned, 1); + let queries = vec![ + VectorSearch::new(vec![1.0, 0.0], 2, "embedding".into()).unwrap(), + VectorSearch::new(vec![0.0, 1.0], 2, "embedding".into()).unwrap(), + ]; + let ranges = vec![RowRange::new(0, 2)]; + let expected = read_raw_batch_vector_search_with_executor( + &pinned, + &queries, + &ranges, + RawVectorMetric::L2, + Some(&filter), + None, + ) + .await + .unwrap() + .0; + + let element = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut values = + arrow_array::builder::ListBuilder::new(Float32Builder::new()).with_field(element.clone()); + values.values().append_slice(&[1.0, 0.0]); + values.append(true); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element), true), + ])), + vec![ + Arc::new(Int32Array::from(vec![4])), + Arc::new(values.finish()), + ], + ) + .unwrap(); + let mut write = crate::table::TableWrite::new(&table, "scheduled-test".into()).unwrap(); + write.write_arrow_batch(&batch).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + crate::table::TableCommit::new(table.clone(), "scheduled-test".into()) + .commit(messages) + .await + .unwrap(); + + let executor = WireSearchExecutor { + table, + kinds: std::sync::Mutex::new(Vec::new()), + fail: false, + }; + let actual = read_raw_batch_vector_search_with_executor( + &pinned, + &queries, + &ranges, + RawVectorMetric::L2, + Some(&filter), + Some(&executor), + ) + .await + .unwrap() + .0; + assert_eq!(expected.len(), actual.len()); + for (expected, actual) in expected.iter().zip(actual) { + assert_eq!(expected.row_ids, actual.row_ids); + assert_eq!(expected.scores, actual.scores); + assert!(actual.row_ids.iter().all(|row| [1, 2].contains(row))); + } +} diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 4b7e4c799..858b93d42 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -46,9 +46,11 @@ use serde_json::json; use std::collections::{HashMap, HashSet}; const FULL_TEXT_INDEX_TYPE: &str = "full-text"; +mod execution; mod filter; mod read; mod scan; +pub(super) use execution::execute_search_task; pub use read::FullTextRead; pub use scan::{FullTextScan, FullTextScanPlan}; @@ -221,10 +223,20 @@ struct FullTextSearchEvaluation<'a> { row_filter: Option<&'a Predicate>, } +#[cfg(test)] async fn evaluate_full_text_search( evaluation: FullTextSearchEvaluation<'_>, index_entries: &[IndexManifestEntry], search: &FullTextSearch, +) -> crate::Result { + evaluate_full_text_search_with_executor(evaluation, index_entries, search, None).await +} + +async fn evaluate_full_text_search_with_executor( + evaluation: FullTextSearchEvaluation<'_>, + index_entries: &[IndexManifestEntry], + search: &FullTextSearch, + executor: Option<&dyn crate::table::SearchTaskExecutor>, ) -> crate::Result { let table_path = evaluation.table_path.trim_end_matches('/'); let core_options = CoreOptions::new(evaluation.table_options); @@ -282,38 +294,53 @@ async fn evaluate_full_text_search( fulltext_entries, indexed_search.include_row_ids.as_ref(), )?; - let results = plans - .into_iter() - .map(|plan| { - let entry = plan.entry; - let global_meta = entry.index_file.global_index_meta.as_ref().unwrap(); - let path = IndexFileLocation::Global { table_path }.resolve( - &entry.index_file.file_name, - entry.index_file.external_path.as_deref(), - ); - let file_name = entry.index_file.file_name.clone(); - let query_text = search.query_text.clone(); - let local_filter = plan.local_filter; - let row_range_start = global_meta.row_range_start; - let row_range_end = global_meta.row_range_end; - // Compound query leaves must see every candidate in the shard - // before Boolean/boost composition and the final global Top-K. - let limit = candidate_limit(row_range_start, row_range_end)?; - Ok::<_, crate::Error>(async move { - let input = evaluation.file_io.new_input(&path)?; - let reader = Box::new(input.reader().await?); - let result = - search_full_text_index(reader, file_name, query_text, limit, local_filter) - .await?; - Ok::<_, crate::Error>(search_result_from_core(result)?.offset(row_range_start)) + let results = if let Some(executor) = executor { + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Scheduled full-text search requires table context".into(), + source: None, + })?; + execution::search_indexes(table, plans, &search.query_text, executor).await? + } else { + let results = plans + .into_iter() + .map(|plan| { + let entry = plan.entry; + let global_meta = entry.index_file.global_index_meta.as_ref().unwrap(); + let path = IndexFileLocation::Global { table_path }.resolve( + &entry.index_file.file_name, + entry.index_file.external_path.as_deref(), + ); + let file_name = entry.index_file.file_name.clone(); + let query_text = search.query_text.clone(); + let local_filter = plan.local_filter; + let row_range_start = global_meta.row_range_start; + let row_range_end = global_meta.row_range_end; + // Compound query leaves must see every candidate in the shard + // before Boolean/boost composition and the final global Top-K. + let limit = candidate_limit(row_range_start, row_range_end)?; + Ok::<_, crate::Error>(async move { + let input = evaluation.file_io.new_input(&path)?; + let reader = Box::new(input.reader().await?); + let result = search_full_text_index( + reader, + file_name, + query_text, + limit, + local_filter, + ) + .await?; + Ok::<_, crate::Error>( + search_result_from_core(result)?.offset(row_range_start), + ) + }) }) - }) - .collect::>>()? - .into_iter(); - let results = futures::stream::iter(results) - .buffered(core_options.global_index_thread_num()?) - .try_collect::>() - .await?; + .collect::>>()? + .into_iter(); + futures::stream::iter(results) + .buffered(core_options.global_index_thread_num()?) + .try_collect::>() + .await? + }; for r in &results { merged = merged.or(r); } @@ -352,14 +379,26 @@ async fn evaluate_full_text_search( matched, )); } - let raw_result = read_raw_full_text_search_filtered( - table, - &raw_search, - &raw_ranges, - evaluation.table_options, - evaluation.partition_filter, - ) - .await?; + let raw_result = if let Some(executor) = executor { + execution::search_raw_rows( + table, + &raw_search, + &raw_ranges, + evaluation.table_options, + evaluation.partition_filter, + executor, + ) + .await? + } else { + read_raw_full_text_search_filtered( + table, + &raw_search, + &raw_ranges, + evaluation.table_options, + evaluation.partition_filter, + ) + .await? + }; merged = merged.without_row_ranges(&raw_ranges)?; merged = merged.or(&raw_result); } @@ -494,6 +533,45 @@ async fn read_raw_full_text_search_filtered( if raw_ranges.is_empty() { return Ok(SearchResult::empty()); } + let splits = plan_raw_full_text_splits(table, search, raw_ranges, partition_filter).await?; + read_raw_full_text_splits( + table, + search, + raw_ranges, + table_options, + &splits, + partition_filter, + ) + .await +} + +async fn plan_raw_full_text_splits( + table: &Table, + search: &FullTextSearch, + raw_ranges: &[RowRange], + partition_filter: Option<&Predicate>, +) -> crate::Result> { + let mut read_builder = table.new_read_builder(); + read_builder + .with_projection(&[search.field_name.as_str(), ROW_ID_FIELD_NAME])? + .with_row_ranges(merge_row_ranges(raw_ranges.to_vec())); + if let Some(filter) = partition_filter { + read_builder.with_filter(filter.clone()); + } + Ok(read_builder.new_scan().plan().await?.splits().to_vec()) +} + +async fn read_raw_full_text_splits( + table: &Table, + search: &FullTextSearch, + raw_ranges: &[RowRange], + table_options: &HashMap, + splits: &[crate::table::DataSplit], + partition_filter: Option<&Predicate>, +) -> crate::Result { + if raw_ranges.is_empty() || splits.is_empty() { + return Ok(SearchResult::empty()); + } let raw_ranges = merge_row_ranges(raw_ranges.to_vec()); let row_range_start = raw_ranges.first().map(RowRange::from).unwrap_or(0); let row_range_end = raw_ranges @@ -508,12 +586,8 @@ async fn read_raw_full_text_search_filtered( if let Some(filter) = partition_filter { read_builder.with_filter(filter.clone()); } - let plan = read_builder.new_scan().plan().await?; - if plan.splits().is_empty() { - return Ok(SearchResult::empty()); - } let read = read_builder.new_read()?; - let mut stream = read.to_arrow(plan.splits())?; + let mut stream = read.to_arrow(splits)?; let config = raw_full_text_config(table_options)?; let mut writer = tokio::task::spawn_blocking(move || { @@ -1362,6 +1436,37 @@ mod tests { // (d) The PK full-text path produces physical positions, not global row-ids: // Only the ranges-only `execute` API rejects the physical-position result. + #[tokio::test] + async fn scheduled_full_text_rejects_primary_key_plans() { + struct NoScheduling; + #[async_trait::async_trait] + impl crate::table::SearchTaskExecutor for NoScheduling { + async fn execute( + &self, + _: Vec, + ) -> crate::Result> { + panic!("a primary-key plan must fail before scheduling"); + } + } + let table = pk_full_text_table("pk_ft_scheduled_unsupported", &[]); + let mut builder = table.new_full_text_search_builder(); + builder + .with_query("body", r#"{"match":{"query":"alpha"}}"#) + .with_limit(10); + let plan = builder.new_scan().unwrap().scan().await.unwrap(); + let error = builder + .new_read() + .unwrap() + .read_with_executor(plan, &NoScheduling) + .await + .unwrap_err(); + assert!( + matches!(error, crate::Error::Unsupported { ref message } + if message.contains("data-evolution")), + "{error}" + ); + } + #[tokio::test] async fn pk_full_text_ranges_reject_physical_positions() { let table = pk_full_text_table("pk_ft_execute_loud", &[]); diff --git a/crates/paimon/src/table/full_text_search_builder/execution.rs b/crates/paimon/src/table/full_text_search_builder/execution.rs new file mode 100644 index 000000000..e7bd3e5a4 --- /dev/null +++ b/crates/paimon/src/table/full_text_search_builder/execution.rs @@ -0,0 +1,225 @@ +// 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. + +//! Serializable full-text shard work; global filtering and ranking stay in core. + +use super::*; +use crate::table::search_task::{execute_tasks, SEARCH_TASK_WINDOW}; +use crate::table::{DataSplit, SearchTask, SearchTaskExecutor}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +const INDEX: &str = "fulltext-index"; +const ROWS: &str = "fulltext-rows"; + +#[derive(Serialize, Deserialize)] +struct IndexWork { + entry: IndexManifestEntry, + query: String, + local_filter: Option>, +} + +#[derive(Serialize, Deserialize)] +struct RowsWork { + query: String, + field: String, + limit: usize, + include_row_ids: Option>, + ranges: Vec, + splits: Vec, + options: HashMap, +} + +#[derive(Serialize, Deserialize)] +struct Hits { + row_ids: Vec, + scores: Vec, +} + +impl Hits { + fn from_search(result: SearchResult) -> Self { + Self { + row_ids: result.row_ids, + scores: result.scores, + } + } + + fn into_search(self) -> crate::Result { + if self.row_ids.len() != self.scores.len() || self.scores.iter().any(|s| !s.is_finite()) { + return Err(invalid("Full-text task returned invalid row IDs or scores")); + } + Ok(SearchResult::new(self.row_ids, self.scores)) + } +} + +fn invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +fn encode(value: T) -> crate::Result { + serde_json::to_value(value) + .map_err(|error| invalid(format!("Cannot encode full-text task: {error}"))) +} + +fn decode(value: &Value) -> crate::Result { + serde_json::from_value(value.clone()) + .map_err(|error| invalid(format!("Invalid full-text task payload: {error}"))) +} + +fn encode_bitmap(ids: &RoaringTreemap) -> crate::Result> { + let mut bytes = Vec::new(); + ids.serialize_into(&mut bytes) + .map_err(|error| invalid(format!("Cannot encode full-text task row IDs: {error}")))?; + Ok(bytes) +} + +fn decode_bitmap(bytes: Vec) -> crate::Result { + RoaringTreemap::deserialize_from(bytes.as_slice()) + .map_err(|error| invalid(format!("Invalid full-text task row IDs: {error}"))) +} + +pub(super) async fn search_indexes( + table: &Table, + plans: Vec, + query: &str, + executor: &dyn SearchTaskExecutor, +) -> crate::Result> { + let mut plans = plans.into_iter(); + let mut results = Vec::new(); + loop { + // Materialize only the next window's query and bitmap payloads. A large + // index manifest must not retain a serialized task for every shard. + let tasks = plans + .by_ref() + .take(SEARCH_TASK_WINDOW) + .map(|plan| { + SearchTask::new( + table, + INDEX, + encode(IndexWork { + entry: plan.entry, + query: query.to_string(), + local_filter: plan.local_filter.as_ref().map(encode_bitmap).transpose()?, + })?, + ) + }) + .collect::>>()?; + if tasks.is_empty() { + break; + } + for value in execute_tasks(executor, tasks).await? { + results.push(decode::(&value)?.into_search()?); + } + } + Ok(results) +} + +pub(super) async fn search_raw_rows( + table: &Table, + search: &FullTextSearch, + ranges: &[RowRange], + options: &HashMap, + partition_filter: Option<&Predicate>, + executor: &dyn SearchTaskExecutor, +) -> crate::Result { + let splits = plan_raw_full_text_splits(table, search, ranges, partition_filter).await?; + if splits.is_empty() { + return Ok(SearchResult::empty()); + } + // A full-text raw fallback is one BM25 corpus, even when the scan has many + // data splits. Rebuilding one index per split changes relevance scores. + let task = SearchTask::new( + table, + ROWS, + encode(RowsWork { + query: search.query_text.clone(), + field: search.field_name.clone(), + limit: search.limit, + include_row_ids: search + .include_row_ids + .as_ref() + .map(encode_bitmap) + .transpose()?, + ranges: ranges.to_vec(), + splits, + options: options.clone(), + })?, + )?; + let mut values = execute_tasks(executor, vec![task]).await?; + decode::(&values.remove(0))?.into_search() +} + +pub(crate) async fn execute_search_task( + table: &Table, + kind: &str, + payload: &Value, +) -> crate::Result { + match kind { + INDEX => { + let work: IndexWork = decode(payload)?; + let meta = work + .entry + .index_file + .global_index_meta + .as_ref() + .ok_or_else(|| invalid("Full-text index task is missing global index metadata"))?; + let limit = candidate_limit(meta.row_range_start, meta.row_range_end)?; + let path = IndexFileLocation::Global { + table_path: table.location().trim_end_matches('/'), + } + .resolve( + &work.entry.index_file.file_name, + work.entry.index_file.external_path.as_deref(), + ); + let input = table.file_io().new_input(&path)?; + let reader = Box::new(input.reader().await?); + let result = search_full_text_index( + reader, + work.entry.index_file.file_name.clone(), + work.query, + limit, + work.local_filter.map(decode_bitmap).transpose()?, + ) + .await?; + encode(Hits::from_search( + search_result_from_core(result)?.offset(meta.row_range_start), + )) + } + ROWS => { + let work: RowsWork = decode(payload)?; + let mut search = FullTextSearch::new(work.query, work.limit, work.field)?; + search.include_row_ids = work.include_row_ids.map(decode_bitmap).transpose()?; + let result = read_raw_full_text_splits( + table, + &search, + &work.ranges, + &work.options, + &work.splits, + // Partition pruning already selected these exact source splits. + None, + ) + .await?; + encode(Hits::from_search(result)) + } + _ => Err(invalid(format!( + "Unknown full-text search task kind: {kind}" + ))), + } +} diff --git a/crates/paimon/src/table/full_text_search_builder/parity_tests.rs b/crates/paimon/src/table/full_text_search_builder/parity_tests.rs index d37882857..9903a1483 100644 --- a/crates/paimon/src/table/full_text_search_builder/parity_tests.rs +++ b/crates/paimon/src/table/full_text_search_builder/parity_tests.rs @@ -655,3 +655,259 @@ async fn scalar_search_pre_filters_exclude_composite_indexes() { ); } } + +struct WireExecutor { + table: Table, + work: std::sync::Mutex>, + batch_sizes: std::sync::Mutex>, + corrupt_scores: bool, +} + +#[async_trait::async_trait] +impl crate::table::SearchTaskExecutor for WireExecutor { + async fn execute( + &self, + tasks: Vec, + ) -> crate::Result> { + self.batch_sizes.lock().unwrap().push(tasks.len()); + let mut results = Vec::new(); + let (context, frames) = crate::table::SearchTask::serialize_batch(&tasks)?; + for (task, frame) in tasks.into_iter().zip(frames) { + let bytes = task.serialize()?; + self.work + .lock() + .unwrap() + .push(serde_json::from_slice(&bytes).unwrap()); + let mut output: serde_json::Value = + serde_json::from_slice(&self.table.execute_search_task(&frame, &context).await?) + .unwrap(); + if self.corrupt_scores { + output["payload"]["scores"] = json!([]); + } + results.push(crate::table::SearchTaskResult::deserialize( + &serde_json::to_vec(&output).unwrap(), + )?); + } + // Completion order belongs to the scheduler. Core re-associates the + // results before deterministic merge and floating-point accumulation. + results.reverse(); + Ok(results) + } +} + +#[tokio::test] +async fn scheduled_full_text_preserves_filters_global_top_k_raw_corpus_and_snapshot() { + let table = new_table(false).copy_with_options(HashMap::from([( + "source.split.target-size".into(), + "1b".into(), + )])); + append(&table, 0, 8).await; + index(&table, "full-text", "text").await; + append(&table, 8, 4).await; + append(&table, 12, 4).await; + let pb = PredicateBuilder::new(table.schema().fields()); + let mut builder = query(&table); + builder + .with_filter(pb.equal("label", Datum::String("keep".into())).unwrap()) + .with_include_row_ids([1, 5, 9, 13].into_iter().collect()) + .with_limit(3); + let plan = builder.new_scan().unwrap().scan().await.unwrap(); + let snapshot = plan.snapshot_id().unwrap(); + let read = builder.new_read().unwrap(); + // A worker starts from a current table. Its task must restore the captured + // snapshot and use the exact data splits selected for the raw corpus. + append(&table, 16, 4).await; + let expected = read.read(plan.clone()).await.unwrap(); + let executor = WireExecutor { + table: table.clone(), + work: Default::default(), + batch_sizes: Default::default(), + corrupt_scores: false, + }; + let actual = read.read_with_executor(plan, &executor).await.unwrap(); + assert_eq!(actual.snapshot_id(), Some(snapshot)); + let actual = actual.into_row_ids().unwrap(); + let expected = expected.into_row_ids().unwrap(); + assert_eq!(actual.row_ids, expected.row_ids); + assert_eq!(actual.scores, expected.scores); + assert_eq!(actual.row_ids, vec![1, 5, 9]); + let work = executor.work.lock().unwrap(); + assert_eq!( + work.iter() + .filter(|task| task["kind"] == "fulltext-index") + .count(), + 2 + ); + let raw = work + .iter() + .filter(|task| task["kind"] == "fulltext-rows") + .collect::>(); + assert_eq!( + raw.len(), + 1, + "raw BM25 relevance requires one shared corpus" + ); + assert!(raw[0]["payload"]["splits"].as_array().unwrap().len() >= 2); + assert!(work.iter().all(|task| task["snapshot"]["id"] == snapshot)); +} + +#[tokio::test] +async fn scheduled_full_text_rejects_malformed_worker_scores() { + let table = new_table(false); + append(&table, 0, 4).await; + index(&table, "full-text", "text").await; + let builder = query(&table); + let plan = builder.new_scan().unwrap().scan().await.unwrap(); + let executor = WireExecutor { + table: table.clone(), + work: Default::default(), + batch_sizes: Default::default(), + corrupt_scores: true, + }; + let error = builder + .new_read() + .unwrap() + .read_with_executor(plan, &executor) + .await + .unwrap_err(); + assert!( + matches!(error, crate::Error::DataInvalid { ref message, .. } + if message.contains("invalid row IDs or scores")), + "{error}" + ); +} + +#[tokio::test] +async fn scheduled_full_text_rejects_foreign_plan_before_scheduling() { + let table = new_table(false); + let other = new_table(false); + let builder = query(&table); + let plan = query(&other).new_scan().unwrap().scan().await.unwrap(); + let executor = WireExecutor { + table: table.clone(), + work: Default::default(), + batch_sizes: Default::default(), + corrupt_scores: false, + }; + let error = builder + .new_read() + .unwrap() + .read_with_executor(plan, &executor) + .await + .unwrap_err(); + assert!( + matches!(error, crate::Error::DataInvalid { ref message, .. } + if message.contains("different table")), + "{error}" + ); + assert!(executor.work.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn scheduled_full_text_keeps_raw_partition_pruning() { + let table = new_table(true); + append(&table, 0, 4).await; + index(&table, "full-text", "text").await; + append(&table, 4, 4).await; // An uncovered partition outside the query. + append_labels(&table, 8, &[Some("keep"); 4]).await; // More rows in pt=0. + let pb = PredicateBuilder::new(table.schema().fields()); + let mut builder = query(&table); + builder + .with_partition_filter(pb.equal("pt", Datum::Int(0)).unwrap()) + .unwrap(); + let plan = builder.new_scan().unwrap().scan().await.unwrap(); + let read = builder.new_read().unwrap(); + let expected = read + .read(plan.clone()) + .await + .unwrap() + .into_row_ids() + .unwrap(); + let executor = WireExecutor { + table: table.clone(), + work: Default::default(), + batch_sizes: Default::default(), + corrupt_scores: false, + }; + let actual = read + .read_with_executor(plan, &executor) + .await + .unwrap() + .into_row_ids() + .unwrap(); + assert_eq!(actual.row_ids, expected.row_ids); + assert_eq!(actual.scores, expected.scores); + let mut ids = actual.row_ids; + ids.sort_unstable(); + assert_eq!(ids, vec![0, 1, 2, 8, 9, 10, 11]); +} + +#[tokio::test] +async fn scheduled_full_text_bounds_index_task_windows_and_keeps_all_candidates() { + use crate::table::search_task::SEARCH_TASK_WINDOW; + + let table = new_table(false).copy_with_options(HashMap::from([ + ("global-index.row-count-per-shard".into(), "1".into()), + ("full-text-index.search-mode".into(), "fast".into()), + ])); + // Include the null-only shard: planning still schedules its empty index. + // Two full windows and one last shard must retain every known matching row. + let row_count = SEARCH_TASK_WINDOW * 2 + 1; + append(&table, 0, row_count as i32).await; + index(&table, "full-text", "text").await; + let mut builder = query(&table); + builder.with_limit(row_count); + let plan = builder.new_scan().unwrap().scan().await.unwrap(); + let files = plan + .entries + .iter() + .filter(|entry| { + entry.kind == FileKind::Add && entry.index_file.index_type == FULL_TEXT_INDEX_TYPE + }) + .map(|entry| entry.index_file.file_name.clone()) + .collect::>(); + let read = builder.new_read().unwrap(); + let expected = read + .read(plan.clone()) + .await + .unwrap() + .into_row_ids() + .unwrap(); + let executor = WireExecutor { + table: table.clone(), + work: Default::default(), + batch_sizes: Default::default(), + corrupt_scores: false, + }; + let actual = read + .read_with_executor(plan, &executor) + .await + .unwrap() + .into_row_ids() + .unwrap(); + assert_eq!(actual.row_ids, expected.row_ids); + assert_eq!(actual.scores, expected.scores); + let mut ids = actual.row_ids; + ids.sort_unstable(); + assert_eq!( + ids, + (0..row_count as u64) + .filter(|id| *id != 3) + .collect::>() + ); + assert_eq!( + *executor.batch_sizes.lock().unwrap(), + vec![SEARCH_TASK_WINDOW, SEARCH_TASK_WINDOW, 1] + ); + let work = executor.work.lock().unwrap(); + let actual_files = work + .iter() + .map(|task| { + serde_json::from_value::(task["payload"]["entry"].clone()) + .unwrap() + .index_file + .file_name + }) + .collect::>(); + assert_eq!(actual_files, files, "windows must retain manifest order"); +} diff --git a/crates/paimon/src/table/full_text_search_builder/read.rs b/crates/paimon/src/table/full_text_search_builder/read.rs index d24227cd4..a73146f96 100644 --- a/crates/paimon/src/table/full_text_search_builder/read.rs +++ b/crates/paimon/src/table/full_text_search_builder/read.rs @@ -84,6 +84,24 @@ impl FullTextRead { pub async fn read( &self, plan: FullTextScanPlan, + ) -> crate::Result { + self.read_internal(plan, None).await + } + + /// Execute shard work through a scheduler while retaining the plan's + /// snapshot, filtering, and final global ranking in core. + pub async fn read_with_executor( + &self, + plan: FullTextScanPlan, + executor: &dyn crate::table::SearchTaskExecutor, + ) -> crate::Result { + self.read_internal(plan, Some(executor)).await + } + + async fn read_internal( + &self, + plan: FullTextScanPlan, + executor: Option<&dyn crate::table::SearchTaskExecutor>, ) -> crate::Result { CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; CoreOptions::new(plan.table.schema().options()).ensure_read_authorized()?; @@ -93,6 +111,11 @@ impl FullTextRead { source: None, }); } + if executor.is_some() && !self.table.schema().core_options().data_evolution_enabled() { + return Err(crate::Error::Unsupported { + message: "Scheduled full-text search requires a data-evolution table".into(), + }); + } let table = std::sync::Arc::new(plan.table.clone()); if self.context.primary_key_field_id.is_some() { let empty = super::super::pk_full_text_scan::PrimaryKeyFullTextScanPlan { @@ -131,7 +154,7 @@ impl FullTextRead { } let mut search = self.search.clone(); search.include_row_ids = self.context.include_row_ids.clone(); - let hits = evaluate_full_text_search( + let hits = evaluate_full_text_search_with_executor( FullTextSearchEvaluation { table: Some(&plan.table), file_io: plan.table.file_io(), @@ -144,6 +167,7 @@ impl FullTextRead { }, &plan.entries, &search, + executor, ) .await?; Ok(crate::table::GlobalIndexResult::from_global_search( diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index 7970e8dee..ce57a9fb8 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -316,8 +316,29 @@ impl<'a> HybridSearchBuilder<'a> { } pub async fn execute_scored(&self) -> crate::Result { + self.execute_scored_internal(None).await + } + + /// Schedule each route's shard work while preserving one shared snapshot + /// and the configured core fusion ranker. + pub async fn execute_with_executor( + &self, + executor: &dyn super::SearchTaskExecutor, + ) -> crate::Result { + self.execute_scored_internal(Some(executor)).await + } + + async fn execute_scored_internal( + &self, + executor: Option<&dyn super::SearchTaskExecutor>, + ) -> crate::Result { let core = CoreOptions::new(self.table.schema().options()); core.ensure_read_authorized()?; + if executor.is_some() && !core.data_evolution_enabled() { + return Err(crate::Error::Unsupported { + message: "Scheduled hybrid search requires a data-evolution table".into(), + }); + } let limit = self.validate_search()?; // A primary-key hybrid fuses PHYSICAL positions, not global row ids, so a @@ -342,7 +363,7 @@ impl<'a> HybridSearchBuilder<'a> { // selectors independently in the routes can observe different versions. let (table, hits) = match self.resolve_pinned_route_table().await? { Some(table) => { - let hits = self.execute_global_routes(&table, limit).await?; + let hits = self.execute_global_routes(&table, limit, executor).await?; (table, hits) } None => ( @@ -427,19 +448,29 @@ impl<'a> HybridSearchBuilder<'a> { &self, table: &Table, limit: usize, + executor: Option<&dyn super::SearchTaskExecutor>, ) -> crate::Result { // Bound concurrent route work while retaining input order for float // accumulation and keeping each route's weight and limit attached to it. let route_results = stream::iter(self.routes.clone()) .map(|route| async move { let result = match route.kind { - HybridSearchRouteKind::Vector => self - .vector_route_builder(table, &route) - .execute() - .await? - .into_row_ids()?, + HybridSearchRouteKind::Vector => { + let builder = self.vector_route_builder(table, &route); + let result = if let Some(executor) = executor { + let plan = builder.new_scan()?.plan().await?; + builder + .new_read()? + .read_with_executor(plan, executor) + .await? + } else { + builder.execute().await? + }; + result.into_row_ids()? + } HybridSearchRouteKind::FullText => { - execute_full_text_route(table, &route, self.filter.as_ref()).await? + execute_full_text_route(table, &route, self.filter.as_ref(), executor) + .await? } }; Ok::<_, crate::Error>(WeightedRouteResult { @@ -1049,9 +1080,16 @@ async fn execute_full_text_route( table: &Table, route: &HybridSearchRoute, filter: Option<&Predicate>, + executor: Option<&dyn super::SearchTaskExecutor>, ) -> crate::Result { let builder = full_text_route_builder(table, route, filter); - let result = builder.execute_scored().await?; + let result = if let Some(executor) = executor { + let read = builder.new_read()?; + let plan = builder.new_scan()?.scan().await?; + read.read_with_executor(plan, executor).await? + } else { + builder.execute_scored().await? + }; result.into_row_ids() } @@ -1060,6 +1098,7 @@ async fn execute_full_text_route( _table: &Table, _route: &HybridSearchRoute, _filter: Option<&Predicate>, + _executor: Option<&dyn super::SearchTaskExecutor>, ) -> crate::Result { Err(crate::Error::ConfigInvalid { message: "Full-text hybrid routes require the fulltext feature".to_string(), diff --git a/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs index 8c1031018..d3ad4acf4 100644 --- a/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs +++ b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs @@ -417,3 +417,80 @@ async fn empty_tables_still_validate_route_and_partition_configuration() { .with_limit(4); assert!(missing.execute_scored().await.is_err()); } + +struct WireExecutor { + table: Table, + kinds: std::sync::Mutex>, +} + +#[async_trait::async_trait] +impl crate::table::SearchTaskExecutor for WireExecutor { + async fn execute( + &self, + tasks: Vec, + ) -> crate::Result> { + let mut results = Vec::new(); + let (context, frames) = crate::table::SearchTask::serialize_batch(&tasks)?; + for (task, frame) in tasks.into_iter().zip(frames) { + let bytes = task.serialize()?; + let metadata: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + self.kinds + .lock() + .unwrap() + .push(metadata["kind"].as_str().unwrap().to_string()); + results.push(crate::table::SearchTaskResult::deserialize( + &self.table.execute_search_task(&frame, &context).await?, + )?); + } + results.reverse(); + Ok(results) + } +} + +#[tokio::test] +async fn scheduled_hybrid_shares_snapshot_and_preserves_all_core_rankers() { + let table = table(); + append(&table, 0).await; + append(&table, 4).await; + index(&table, "ivf-flat", "embedding").await; + index(&table, "full-text", "text").await; + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let pb = PredicateBuilder::new(table.schema().fields()); + let keep = pb.equal("label", Datum::String("keep".into())).unwrap(); + // The table used by workers has a later commit than the builder's shared + // view. Every route must keep the captured view through task transport. + append(&table, 8).await; + for ranker in [ + HybridSearchRanker::Rrf, + HybridSearchRanker::WeightedScore, + HybridSearchRanker::Mrr, + ] { + let mut builder = mixed_query(&table, 3); + builder + .with_snapshot(Some(&snapshot)) + .with_filter(keep.clone()) + .with_ranker(ranker.as_str()) + .unwrap(); + let expected = builder.execute_scored().await.unwrap(); + let executor = WireExecutor { + table: table.clone(), + kinds: Default::default(), + }; + let actual = builder.execute_with_executor(&executor).await.unwrap(); + assert_eq!(actual.snapshot_id(), Some(snapshot.id())); + let actual = actual.into_row_ids().unwrap(); + let expected = expected.into_row_ids().unwrap(); + assert_eq!(actual.row_ids, expected.row_ids, "ranker {ranker:?}"); + assert_eq!(actual.scores, expected.scores, "ranker {ranker:?}"); + assert!(!actual.row_ids.is_empty()); + assert!(actual.row_ids.iter().all(|id| [1, 2, 5, 6].contains(id))); + let kinds = executor.kinds.lock().unwrap(); + assert!(kinds.iter().any(|kind| kind.starts_with("vector-"))); + assert!(kinds.iter().any(|kind| kind == "fulltext-index")); + } +} diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index a8b35eb12..0862146c7 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -150,6 +150,8 @@ mod row_sidecar; mod scan_stats; pub(crate) mod schema_manager; mod search_scan_split; +mod search_task; +pub use search_task::{SearchTask, SearchTaskExecutor, SearchTaskResult}; mod shard_table_updater; pub(crate) mod snapshot_commit; mod snapshot_manager; diff --git a/crates/paimon/src/table/search_task.rs b/crates/paimon/src/table/search_task.rs new file mode 100644 index 000000000..fdd9866fd --- /dev/null +++ b/crates/paimon/src/table/search_task.rs @@ -0,0 +1,447 @@ +// 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. + +//! Opaque, snapshot-scoped search work for local or remote schedulers. + +use super::Table; +use crate::spec::{CoreOptions, Snapshot, TableSchema}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +const VERSION: u8 = 1; +pub(super) const SEARCH_TASK_WINDOW: usize = 16; + +pub(super) fn invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +pub(super) fn encode(value: &T) -> crate::Result> { + serde_json::to_vec(value).map_err(|error| invalid(format!("Invalid search task: {error}"))) +} + +pub(super) fn decode(value: &[u8]) -> crate::Result { + serde_json::from_slice(value).map_err(|error| invalid(format!("Invalid search task: {error}"))) +} + +/// Work created by Core. Schedulers transport it without interpreting its contents. +#[derive(Clone, Serialize, Deserialize)] +pub struct SearchTask { + version: u8, + id: String, + location: String, + branch: String, + schema: TableSchema, + snapshot: Option, + kind: String, + payload: Value, +} + +impl SearchTask { + pub(super) fn new(table: &Table, kind: &str, payload: Value) -> crate::Result { + Ok(Self { + version: VERSION, + id: uuid::Uuid::new_v4().to_string(), + location: table.location().trim_end_matches('/').into(), + branch: table.branch().into(), + schema: table.schema().clone(), + // None is an explicitly captured empty table, never "latest". + snapshot: table.travel_snapshot().cloned(), + kind: kind.into(), + payload, + }) + } + + pub fn serialize(&self) -> crate::Result> { + encode(self) + } + + /// Share repeated immutable task inputs once per scheduler invocation. + /// Query matrices, schema/snapshot data and repeated bitmaps stay opaque. + pub fn serialize_batch(tasks: &[Self]) -> crate::Result<(Vec, Vec>)> { + let values = tasks + .iter() + .map(|task| serde_json::to_value(task).map_err(|error| invalid(error.to_string()))) + .collect::>>()?; + let mut counts = HashMap::, usize>::new(); + fn count(value: &Value, counts: &mut HashMap, usize>) -> crate::Result<()> { + if value.is_array() || value.is_object() { + let bytes = encode(value)?; + if bytes.len() >= 256 { + *counts.entry(bytes).or_default() += 1; + } + match value { + Value::Array(values) => { + for value in values { + count(value, counts)?; + } + } + Value::Object(values) => { + for value in values.values() { + count(value, counts)?; + } + } + _ => (), + } + } + Ok(()) + } + for value in &values { + count(value, &mut counts)?; + } + let context_id = uuid::Uuid::new_v4().to_string(); + let mut shared = Vec::new(); + let mut positions = HashMap::, usize>::new(); + fn share( + value: &mut Value, + path: &str, + counts: &HashMap, usize>, + positions: &mut HashMap, usize>, + shared: &mut Vec, + references: &mut Vec<(String, usize)>, + ) -> crate::Result<()> { + if value.is_array() || value.is_object() { + let bytes = encode(value)?; + if counts.get(&bytes).is_some_and(|count| *count > 1) { + let index = match positions.get(&bytes) { + Some(index) => *index, + None => { + let index = shared.len(); + shared.push(value.clone()); + positions.insert(bytes, index); + index + } + }; + references.push((path.into(), index)); + *value = Value::Null; + return Ok(()); + } + match value { + Value::Array(values) => { + for (index, value) in values.iter_mut().enumerate() { + share( + value, + &format!("{path}/{index}"), + counts, + positions, + shared, + references, + )?; + } + } + Value::Object(values) => { + for (key, value) in values.iter_mut() { + let escaped = key.replace('~', "~0").replace('/', "~1"); + share( + value, + &format!("{path}/{escaped}"), + counts, + positions, + shared, + references, + )?; + } + } + _ => (), + } + } + Ok(()) + } + let frames = values + .into_iter() + .map(|mut task| { + let mut references = Vec::new(); + share( + &mut task, + "", + &counts, + &mut positions, + &mut shared, + &mut references, + )?; + encode(&SharedTaskFrame { + version: VERSION, + context_id: context_id.clone(), + task, + references, + }) + }) + .collect::>>()?; + Ok(( + encode(&SharedTaskContext { + version: VERSION, + id: context_id, + values: shared, + })?, + frames, + )) + } + + pub fn deserialize_with_context(bytes: &[u8], context: &[u8]) -> crate::Result { + let context: SharedTaskContext = decode(context)?; + let mut frame: SharedTaskFrame = decode(bytes)?; + if context.version != VERSION || frame.version != VERSION || context.id != frame.context_id + { + return Err(invalid("Search task shared context does not match")); + } + for (path, index) in frame.references { + let value = context + .values + .get(index) + .ok_or_else(|| invalid("Invalid search task shared reference"))?; + *frame + .task + .pointer_mut(&path) + .ok_or_else(|| invalid("Invalid search task shared path"))? = value.clone(); + } + let task: Self = + serde_json::from_value(frame.task).map_err(|error| invalid(error.to_string()))?; + if task.version != VERSION { + return Err(invalid("Unsupported search task version")); + } + Ok(task) + } + + pub fn deserialize(bytes: &[u8]) -> crate::Result { + let task: Self = decode(bytes)?; + if task.version != VERSION { + return Err(invalid("Unsupported search task version")); + } + Ok(task) + } + + pub async fn execute(&self, table: &Table) -> crate::Result { + // The worker's runtime restrictions take precedence over transported options. + CoreOptions::new(table.schema().options()).ensure_read_authorized()?; + CoreOptions::new(self.schema.options()).ensure_read_authorized()?; + if self.location != table.location().trim_end_matches('/') || self.branch != table.branch() + { + return Err(invalid( + "Search task belongs to a different table or branch", + )); + } + let pinned = table + .copy_with_resolved_schema(self.schema.clone(), &self.branch)? + .copy_with_pinned_snapshot(self.snapshot.as_ref()); + let payload = match self.kind.as_str() { + "vector-index" | "vector-rows" => { + super::de_vector_read::execute_search_task(&pinned, &self.kind, &self.payload) + .await? + } + #[cfg(feature = "fulltext")] + kind if kind.starts_with("fulltext-") => { + super::full_text_search_builder::execute_search_task(&pinned, kind, &self.payload) + .await? + } + _ => return Err(invalid("Unsupported search task kind")), + }; + Ok(SearchTaskResult { + version: VERSION, + id: self.id.clone(), + payload, + }) + } +} + +/// An opaque worker outcome, tied to the exact task that produced it. +#[derive(Clone, Serialize, Deserialize)] +pub struct SearchTaskResult { + version: u8, + id: String, + payload: Value, +} + +impl SearchTaskResult { + pub fn serialize(&self) -> crate::Result> { + encode(self) + } + pub fn deserialize(bytes: &[u8]) -> crate::Result { + let result: Self = decode(bytes)?; + if result.version != VERSION { + return Err(invalid("Unsupported search result version")); + } + Ok(result) + } +} + +/// A scheduler owns task placement and cancellation; Core owns search semantics. +/// On failure it must cancel or drain outstanding work before returning. +#[async_trait::async_trait] +pub trait SearchTaskExecutor: Send + Sync { + async fn execute(&self, tasks: Vec) -> crate::Result>; +} + +pub(super) async fn execute_tasks( + executor: &dyn SearchTaskExecutor, + tasks: Vec, +) -> crate::Result> { + if tasks.is_empty() { + return Ok(Vec::new()); + } + let ids = tasks.iter().map(|task| task.id.clone()).collect::>(); + let results = executor.execute(tasks).await?; + if results.len() != ids.len() { + return Err(invalid("Search executor returned the wrong result count")); + } + let mut by_id = HashMap::with_capacity(results.len()); + for result in results { + if result.version != VERSION || by_id.insert(result.id, result.payload).is_some() { + return Err(invalid( + "Search executor returned a duplicate or invalid result", + )); + } + } + ids.into_iter() + .map(|id| { + by_id + .remove(&id) + .ok_or_else(|| invalid("Search executor returned a result for another task")) + }) + .collect() +} + +#[derive(Serialize, Deserialize)] +struct SharedTaskContext { + version: u8, + id: String, + values: Vec, +} +#[derive(Serialize, Deserialize)] +struct SharedTaskFrame { + version: u8, + context_id: String, + task: Value, + references: Vec<(String, usize)>, +} + +impl Table { + /// Execute one opaque task and its shared inputs using this worker's credentials. + pub async fn execute_search_task(&self, task: &[u8], context: &[u8]) -> crate::Result> { + SearchTask::deserialize_with_context(task, context)? + .execute(self) + .await? + .serialize() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::table::vector_search_test_utils::vector_test_table; + + #[test] + fn scheduled_transport_shares_large_inputs_and_rejects_foreign_context() { + let table = vector_test_table(); + let matrix = (0..32) + .map(|row| { + (0..128) + .map(|column| (row * 128 + column) as u32) + .collect::>() + }) + .collect::>(); + let tasks = (0..20) + .map(|entry| { + SearchTask::new( + &table, + "vector-index", + serde_json::json!({"queries": matrix, "entry": entry}), + ) + .unwrap() + }) + .collect::>(); + let original_size = tasks + .iter() + .map(|task| task.serialize().unwrap().len()) + .sum::(); + let (context, frames) = SearchTask::serialize_batch(&tasks).unwrap(); + let wire_size = context.len() + frames.iter().map(Vec::len).sum::(); + assert!( + wire_size < original_size / 4, + "shared {wire_size} versus copied {original_size}" + ); + for (task, frame) in tasks.iter().zip(&frames) { + assert_eq!( + task.serialize().unwrap(), + SearchTask::deserialize_with_context(frame, &context) + .unwrap() + .serialize() + .unwrap() + ); + } + let (foreign_context, _) = SearchTask::serialize_batch(&tasks).unwrap(); + assert!(SearchTask::deserialize_with_context(&frames[0], &foreign_context).is_err()); + } + + struct ReplyExecutor { + invalid: bool, + } + #[async_trait::async_trait] + impl SearchTaskExecutor for ReplyExecutor { + async fn execute(&self, tasks: Vec) -> crate::Result> { + let mut results = tasks + .into_iter() + .map(|task| SearchTaskResult { + version: VERSION, + id: task.id, + payload: task.payload, + }) + .collect::>(); + if self.invalid { + results[1].id = results[0].id.clone(); + } + results.reverse(); + Ok(results) + } + } + + #[tokio::test] + async fn scheduled_results_reorder_by_identity_and_reject_duplicates() { + let table = vector_test_table(); + let tasks = (0..2) + .map(|value| SearchTask::new(&table, "test", Value::from(value)).unwrap()) + .collect::>(); + assert_eq!( + execute_tasks(&ReplyExecutor { invalid: false }, tasks.clone()) + .await + .unwrap(), + vec![Value::from(0), Value::from(1)] + ); + assert!(execute_tasks(&ReplyExecutor { invalid: true }, tasks) + .await + .unwrap_err() + .to_string() + .contains("duplicate")); + } + + #[tokio::test] + async fn scheduled_worker_preserves_runtime_authorization_restrictions() { + let table = vector_test_table(); + let task = SearchTask::new(&table, "vector-index", Value::Null).unwrap(); + let schema = table.schema().copy_with_options(HashMap::from([( + "query-auth.enabled".into(), + "true".into(), + )])); + let restricted = table + .copy_with_resolved_schema(schema, table.branch()) + .unwrap(); + let error = task.execute(&restricted).await.err().unwrap(); + assert!(error.to_string().contains("query-auth.enabled"), "{error}"); + } +} diff --git a/crates/paimon/src/table/vector_read.rs b/crates/paimon/src/table/vector_read.rs index 5feef4816..294e4899b 100644 --- a/crates/paimon/src/table/vector_read.rs +++ b/crates/paimon/src/table/vector_read.rs @@ -23,7 +23,7 @@ use crate::table::pk_vector_read::PkVectorRead; use crate::table::pk_vector_search_params::PkVectorSearchParams; use crate::table::vector_scan::{PlanContext, VectorScanPlan, VectorScanWork}; use crate::table::vector_search_common::{take_only_result, targets_primary_key_column}; -use crate::table::Table; +use crate::table::{SearchTaskExecutor, Table}; use crate::vector_search::GlobalIndexResult; use std::collections::HashMap; use std::future::Future; @@ -50,6 +50,16 @@ impl VectorRead { pub async fn read(&self, plan: VectorScanPlan) -> crate::Result { take_only_result(self.batch.read(plan).await?, "vector search") } + pub async fn read_with_executor( + &self, + plan: VectorScanPlan, + executor: &dyn SearchTaskExecutor, + ) -> crate::Result { + take_only_result( + self.batch.read_with_executor(plan, executor).await?, + "vector search", + ) + } } /// Searches a common plan with multiple queries, preserving input order and arity. @@ -89,6 +99,20 @@ impl BatchVectorRead { } pub async fn read(&self, plan: VectorScanPlan) -> crate::Result> { + self.read_impl(plan, None).await + } + pub async fn read_with_executor( + &self, + plan: VectorScanPlan, + executor: &dyn SearchTaskExecutor, + ) -> crate::Result> { + self.read_impl(plan, Some(executor)).await + } + async fn read_impl( + &self, + plan: VectorScanPlan, + executor: Option<&dyn SearchTaskExecutor>, + ) -> crate::Result> { if self.context != plan.context { return Err(crate::Error::DataInvalid { message: "vector plan and reader must use the same table, column and pre-filter" @@ -98,9 +122,14 @@ impl BatchVectorRead { } match (&self.reader, plan.work) { (VectorReadKind::DataEvolution(reader), VectorScanWork::DataEvolution(plan)) => { - reader.read(*plan).await + reader.read_with_executor(*plan, executor).await } (VectorReadKind::PrimaryKey(reader), VectorScanWork::PrimaryKey(plan)) => { + if executor.is_some() { + return Err(crate::Error::Unsupported { + message: "Distributed primary-key vector search is not supported".into(), + }); + } reader.read(*plan).await } _ => Err(crate::Error::DataInvalid { From 49033b50320cb82ca7873133e4c3413bafcee559 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 21:14:43 +0800 Subject: [PATCH 2/6] Construct search plans from snapshot sources and internalize scheduling --- bindings/c/src/vector_scan.rs | 6 +- bindings/python/Cargo.toml | 2 +- .../python/pypaimon_rust/datafusion.pyi | 19 +- bindings/python/src/full_text_search.rs | 73 ++-- bindings/python/src/hybrid_search.rs | 33 +- bindings/python/src/search_task.rs | 2 +- bindings/python/src/table.rs | 8 +- bindings/python/src/vector_search.rs | 375 +++++++++++++++--- crates/paimon/Cargo.toml | 3 + .../src/table/batch_vector_search_builder.rs | 1 + crates/paimon/src/table/de_vector_read.rs | 160 ++++++-- .../paimon/src/table/de_vector_read/tests.rs | 21 +- crates/paimon/src/table/de_vector_scan.rs | 11 +- .../paimon/src/table/de_vector_scan/tests.rs | 364 +++++++++++++++++ .../src/table/full_text_search_builder.rs | 122 +++++- .../full_text_search_builder/execution.rs | 3 + .../full_text_search_builder/parity_tests.rs | 190 ++++++++- .../table/full_text_search_builder/read.rs | 165 ++++++-- .../table/full_text_search_builder/scan.rs | 269 +++++++------ .../paimon/src/table/hybrid_search_builder.rs | 3 +- .../hybrid_search_builder/parity_tests.rs | 73 ++++ crates/paimon/src/table/internal.rs | 67 ++++ crates/paimon/src/table/mod.rs | 9 +- crates/paimon/src/table/pk_vector_scan.rs | 194 +++++++++ crates/paimon/src/table/search_scan_split.rs | 206 +++++++++- crates/paimon/src/table/search_task.rs | 29 +- crates/paimon/src/table/vector_read.rs | 93 ++++- crates/paimon/src/table/vector_scan.rs | 364 ++++++++++------- .../paimon/src/table/vector_search_builder.rs | 1 + .../tests/pk_vector_bucket_split_read_test.rs | 9 +- 30 files changed, 2312 insertions(+), 563 deletions(-) create mode 100644 crates/paimon/src/table/internal.rs diff --git a/bindings/c/src/vector_scan.rs b/bindings/c/src/vector_scan.rs index fbc45eed0..5d4c50c94 100644 --- a/bindings/c/src/vector_scan.rs +++ b/bindings/c/src/vector_scan.rs @@ -64,7 +64,9 @@ pub unsafe extern "C" fn paimon_vector_scan_plan( } /// Construct a common read plan from already-decoded Java PK bucket splits. -/// No snapshot or index manifest is read. Handles are borrowed and copied; +/// Capture the supplied snapshot ID, or the scan's configured snapshot for +/// unversioned buckets, without replanning the index manifest. +/// Handles are borrowed and copied; /// callers may free them as soon as this returns. Failure leaves inputs intact. /// Empty input, DE scans and mixed snapshots are rejected. /// # Safety @@ -111,7 +113,7 @@ pub unsafe extern "C" fn paimon_vector_scan_plan_from_bucket_splits( } decoded.push((&*((*split).inner as *const BucketVectorSearchSplit)).clone()); } - wrap_plan(scan.plan_from_bucket_splits(decoded)) + wrap_plan(runtime().block_on(scan.plan_from_bucket_splits(decoded))) } /// Free a vector scan. Null is accepted. diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml index 733b7b6d3..08063569f 100644 --- a/bindings/python/Cargo.toml +++ b/bindings/python/Cargo.toml @@ -36,7 +36,7 @@ chrono = "0.4.38" datafusion = { workspace = true } datafusion-ffi = { workspace = true } futures = "0.3" -paimon = { path = "../../crates/paimon", features = ["storage-all"] } +paimon = { path = "../../crates/paimon", features = ["storage-all", "internal-api"] } paimon-datafusion = { path = "../../crates/integrations/datafusion", features = ["fulltext"] } paimon-vindex-core = "0.5.0" pyo3 = { version = "0.28", features = ["abi3-py310", "chrono"] } diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 407d0cad1..391c5c9de 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -321,7 +321,6 @@ class PartitionStat: def total_size_bytes(self) -> int: ... class Table: - def _execute_search_task(self, task: bytes, context: bytes) -> bytes: ... @staticmethod def from_resolved_schema( location: str, schema_json: str, *, database: str = "default", @@ -943,14 +942,17 @@ class VectorScan: def scan(self) -> "VectorScanPlan": ... class VectorScanPlan: - def _metadata(self) -> str: ... + def __init__(self, snapshot_json: Optional[str], splits_json: str) -> None: + """Capture resolved Java Snapshot JSON and exported search source splits; None pins empty.""" + ... + def metadata(self) -> str: ... def snapshot_id(self) -> Optional[int]: ... class VectorRead: - def read_plan(self, plan: VectorScanPlan, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> "GlobalIndexResult": ... + def read_plan(self, plan: VectorScanPlan) -> "GlobalIndexResult": ... class BatchVectorRead: - def read_batch_plan(self, plan: VectorScanPlan, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> List["GlobalIndexResult"]: ... + def read_batch_plan(self, plan: VectorScanPlan) -> List["GlobalIndexResult"]: ... class GlobalIndexResultReadBuilder: def with_projection(self, columns: List[str]) -> "GlobalIndexResultReadBuilder": ... @@ -973,7 +975,7 @@ class HybridSearchBuilder: def add_full_text_route(self, field_name: str, query: str, limit: int, weight: float = 1.0, options: Optional[Dict[str, str]] = None) -> HybridSearchBuilder: ... def execute_local(self) -> GlobalIndexResult: ... - def execute(self, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> GlobalIndexResult: ... + def execute(self) -> GlobalIndexResult: ... class FullTextSearchBuilder: def with_query(self, field_name: str, query: str) -> FullTextSearchBuilder: ... @@ -988,11 +990,14 @@ class FullTextScan: def scan(self) -> FullTextScanPlan: ... class FullTextScanPlan: - def _metadata(self) -> str: ... + def __init__(self, snapshot_json: Optional[str], splits_json: str) -> None: + """Capture resolved Java Snapshot JSON and exported search source splits; None pins empty.""" + ... + def metadata(self) -> str: ... def snapshot_id(self) -> Optional[int]: ... class FullTextRead: - def read_plan(self, plan: FullTextScanPlan, executor: Optional[Callable[[bytes, List[bytes]], Iterable[bytes]]] = None) -> GlobalIndexResult: ... + def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult: ... class ShardTableUpdator: def arrow_reader(self) -> RecordBatchReader: diff --git a/bindings/python/src/full_text_search.rs b/bindings/python/src/full_text_search.rs index c102ac7b7..383fb49e7 100644 --- a/bindings/python/src/full_text_search.rs +++ b/bindings/python/src/full_text_search.rs @@ -129,30 +129,6 @@ impl PyFullTextSearchBuilder { .map(PyGlobalIndexResult::new) .map_err(to_py_err) } - - #[pyo3(signature = (executor=None))] - fn execute( - &self, - py: Python<'_>, - executor: Option>, - ) -> PyResult { - let Some(callback) = executor else { - return self.execute_local(py); - }; - let builder = self.config.builder(&self.table).map_err(to_py_err)?; - let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); - let result = py.detach(|| { - runtime().block_on(async { - let read = builder.new_read()?; - let plan = builder.new_scan()?.scan().await?; - read.read_with_executor(plan, &executor).await - }) - }); - if let Some(error) = executor.take_error() { - return Err(error); - } - result.map(PyGlobalIndexResult::new).map_err(to_py_err) - } } #[pyclass(name = "FullTextScan", module = "pypaimon_rust.datafusion")] @@ -174,14 +150,21 @@ pub struct PyFullTextScanPlan { } #[pymethods] impl PyFullTextScanPlan { - fn _metadata(&self, py: Python<'_>) -> PyResult { - let splits = py - .detach(|| runtime().block_on(self.inner.splits())) - .map_err(to_py_err)?; + #[new] + #[pyo3(signature = (snapshot_json, splits_json))] + fn new(snapshot_json: Option<&str>, splits_json: &str) -> PyResult { + let (snapshot, splits) = + crate::vector_search::search_plan_from_json(snapshot_json, splits_json)?; + Ok(Self { + inner: FullTextScanPlan::new(snapshot, splits).map_err(to_py_err)?, + }) + } + + fn metadata(&self) -> PyResult { crate::vector_search::search_plan_metadata_json( self.inner.snapshot(), self.inner.snapshot_id(), - splits, + self.inner.splits(), ) } @@ -196,25 +179,33 @@ pub struct PyFullTextRead { } #[pymethods] impl PyFullTextRead { - #[pyo3(signature = (plan, executor=None))] fn read_plan( &self, py: Python<'_>, plan: &PyFullTextScanPlan, - executor: Option>, ) -> PyResult { - if let Some(callback) = executor { - let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); - let result = py.detach(|| { - runtime().block_on(self.inner.read_with_executor(plan.inner.clone(), &executor)) - }); - if let Some(error) = executor.take_error() { - return Err(error); - } - return result.map(PyGlobalIndexResult::new).map_err(to_py_err); - } py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) .map(PyGlobalIndexResult::new) .map_err(to_py_err) } + + fn _read_plan_with_executor( + &self, + py: Python<'_>, + plan: &PyFullTextScanPlan, + executor: Py, + ) -> PyResult { + let executor = crate::search_task::PythonSearchTaskExecutor::new(executor); + let result = py.detach(|| { + runtime().block_on(paimon::table::internal::read_full_text_with_executor( + &self.inner, + plan.inner.clone(), + &executor, + )) + }); + if let Some(error) = executor.take_error() { + return Err(error); + } + result.map(PyGlobalIndexResult::new).map_err(to_py_err) + } } diff --git a/bindings/python/src/hybrid_search.rs b/bindings/python/src/hybrid_search.rs index c43fb7cc7..374e5ac54 100644 --- a/bindings/python/src/hybrid_search.rs +++ b/bindings/python/src/hybrid_search.rs @@ -177,27 +177,32 @@ impl PyHybridSearchBuilder { Ok(slf) } - #[pyo3(signature = (executor=None))] - fn execute( + fn execute(&self, py: Python<'_>) -> PyResult { + let builder = self.builder(); + py.detach(|| runtime().block_on(builder.execute_scored())) + .map(PyGlobalIndexResult::new) + .map_err(to_py_err) + } + + fn _execute_with_executor( &self, py: Python<'_>, - executor: Option>, + executor: Py, ) -> PyResult { let builder = self.builder(); - if let Some(callback) = executor { - let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); - let result = py.detach(|| runtime().block_on(builder.execute_with_executor(&executor))); - if let Some(error) = executor.take_error() { - return Err(error); - } - return result.map(PyGlobalIndexResult::new).map_err(to_py_err); + let executor = crate::search_task::PythonSearchTaskExecutor::new(executor); + let result = py.detach(|| { + runtime().block_on(paimon::table::internal::execute_hybrid_with_executor( + &builder, &executor, + )) + }); + if let Some(error) = executor.take_error() { + return Err(error); } - py.detach(|| runtime().block_on(builder.execute_scored())) - .map(PyGlobalIndexResult::new) - .map_err(to_py_err) + result.map(PyGlobalIndexResult::new).map_err(to_py_err) } fn execute_local(&self, py: Python<'_>) -> PyResult { - self.execute(py, None) + self.execute(py) } } diff --git a/bindings/python/src/search_task.rs b/bindings/python/src/search_task.rs index b265b82dd..0ab5eb6f0 100644 --- a/bindings/python/src/search_task.rs +++ b/bindings/python/src/search_task.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use paimon::table::{SearchTask, SearchTaskExecutor, SearchTaskResult}; +use paimon::table::internal::{SearchTask, SearchTaskExecutor, SearchTaskResult}; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyList}; use std::sync::Mutex; diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs index 69ef1e83c..cf8bafa8b 100644 --- a/bindings/python/src/table.rs +++ b/bindings/python/src/table.rs @@ -56,7 +56,13 @@ impl PyTable { context: &[u8], ) -> PyResult> { let result = py - .detach(|| runtime().block_on(self.inner.execute_search_task(task, context))) + .detach(|| { + runtime().block_on(paimon::table::internal::execute_search_task( + &self.inner, + task, + context, + )) + }) .map_err(to_py_err)?; Ok(pyo3::types::PyBytes::new(py, &result)) } diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index 1c2d19245..5e91f7cc6 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -19,7 +19,10 @@ use std::collections::HashMap; use std::sync::Arc; use paimon::spec::Predicate; -use paimon::table::{BatchVectorRead, Table, VectorRead, VectorScan, VectorScanPlan}; +use paimon::table::{ + BatchVectorRead, DataSplit, RowRange, SearchScanSplit, Table, VectorRead, VectorScan, + VectorScanPlan, +}; use paimon_datafusion::runtime::runtime; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -31,9 +34,8 @@ use crate::predicate::dict_to_table_predicate; pub(crate) fn search_plan_metadata_json( snapshot: Option<&paimon::spec::Snapshot>, snapshot_id: Option, - splits: Vec, + splits: &[SearchScanSplit], ) -> PyResult { - use paimon::table::SearchScanSplit; use serde_json::json; let files = |files: &[paimon::spec::IndexFileMeta]| { files @@ -41,17 +43,17 @@ pub(crate) fn search_plan_metadata_json( .map(|file| file.to_serialized_row_data().map_err(to_py_err)) .collect::>>() }; - let splits = splits.into_iter().map(|split| Ok(match split { + let splits = splits.iter().map(|split| Ok(match split { SearchScanSplit::Index { column, row_range, index_files, scalar_index_files } => json!({ "kind": "index", "column": column, "row_range": row_range, - "index_files": files(&index_files)?, "scalar_index_files": files(&scalar_index_files)?, + "index_files": files(index_files)?, "scalar_index_files": files(scalar_index_files)?, }), SearchScanSplit::Raw { row_ranges, scalar_index_files, index_type } => json!({ - "kind": "raw", "row_ranges": row_ranges, "scalar_index_files": files(&scalar_index_files)?, "index_type": index_type, + "kind": "raw", "row_ranges": row_ranges, "scalar_index_files": files(scalar_index_files)?, "index_type": index_type, }), SearchScanSplit::PrimaryKey { data_split, payloads, uncovered_data_files, row_ranges_by_file } => json!({ "kind": "primary-key", "data_split": data_split.serialize_split_v1().map_err(to_py_err)?, - "payloads": files(&payloads)?, "uncovered_data_files": uncovered_data_files, + "payloads": files(payloads)?, "uncovered_data_files": uncovered_data_files, "row_ranges_by_file": row_ranges_by_file, }), })).collect::>>()?; @@ -61,6 +63,105 @@ pub(crate) fn search_plan_metadata_json( .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) } +/// Decode the same Java metadata codecs used by source-plan export. The binding +/// does not infer work from the snapshot or bind a plan to a live reader. +pub(crate) fn search_plan_from_json( + snapshot_json: Option<&str>, + splits_json: &str, +) -> PyResult<(Option, Vec)> { + use pyo3::exceptions::PyValueError; + use serde_json::Value; + + let invalid = |error: serde_json::Error| PyValueError::new_err(error.to_string()); + let snapshot = snapshot_json + .map(serde_json::from_str) + .transpose() + .map_err(invalid)?; + let values: Vec = serde_json::from_str(splits_json).map_err(invalid)?; + let field = |value: &Value, name: &str| -> PyResult { + value + .get(name) + .cloned() + .ok_or_else(|| PyValueError::new_err(format!("Search split is missing '{name}'"))) + }; + let files = |value: Value| -> PyResult> { + let bytes: Vec> = serde_json::from_value(value).map_err(invalid)?; + bytes + .iter() + .map(|bytes| { + paimon::spec::IndexFileMeta::from_serialized_row_data(bytes).map_err(to_py_err) + }) + .collect() + }; + let validate_ranges = |ranges: &[RowRange]| -> PyResult<()> { + if ranges + .iter() + .any(|range| range.from() < 0 || range.from() > range.to()) + { + return Err(PyValueError::new_err( + "Search row ranges must be nonnegative and ordered", + )); + } + Ok(()) + }; + let splits = values + .iter() + .map(|value| -> PyResult { + let kind: String = serde_json::from_value(field(value, "kind")?).map_err(invalid)?; + Ok(match kind.as_str() { + "index" => { + let row_range: RowRange = + serde_json::from_value(field(value, "row_range")?).map_err(invalid)?; + validate_ranges(std::slice::from_ref(&row_range))?; + SearchScanSplit::Index { + column: serde_json::from_value(field(value, "column")?).map_err(invalid)?, + row_range, + index_files: files(field(value, "index_files")?)?, + scalar_index_files: files(field(value, "scalar_index_files")?)?, + } + } + "raw" => { + let row_ranges: Vec = + serde_json::from_value(field(value, "row_ranges")?).map_err(invalid)?; + validate_ranges(&row_ranges)?; + SearchScanSplit::Raw { + row_ranges, + scalar_index_files: files(field(value, "scalar_index_files")?)?, + index_type: serde_json::from_value(field(value, "index_type")?) + .map_err(invalid)?, + } + } + "primary-key" => { + let bytes: Vec = + serde_json::from_value(field(value, "data_split")?).map_err(invalid)?; + let row_ranges_by_file: HashMap> = + serde_json::from_value(field(value, "row_ranges_by_file")?) + .map_err(invalid)?; + for ranges in row_ranges_by_file.values() { + validate_ranges(ranges)?; + } + SearchScanSplit::PrimaryKey { + data_split: DataSplit::deserialize_split_v1(&bytes).map_err(to_py_err)?, + payloads: files(field(value, "payloads")?)?, + uncovered_data_files: serde_json::from_value(field( + value, + "uncovered_data_files", + )?) + .map_err(invalid)?, + row_ranges_by_file, + } + } + _ => { + return Err(PyValueError::new_err(format!( + "Unknown search split kind '{kind}'" + ))) + } + }) + }) + .collect::>>()?; + Ok((snapshot, splits)) +} + #[derive(Clone, Default)] struct SearchConfig { column: Option, @@ -188,15 +289,7 @@ impl PyVectorSearchBuilder { .map_err(to_py_err)? .new_scan() .map_err(to_py_err)?; - Ok(PyVectorScan { - inner, - options: self - .config - .options - .iter() - .flat_map(|options| options.clone()) - .collect(), - }) + Ok(PyVectorScan { inner }) } fn new_vector_search_read(&self) -> PyResult { let mut builder = self.config.single(&self.table).map_err(to_py_err)?; @@ -291,15 +384,7 @@ impl PyBatchVectorSearchBuilder { .map_err(to_py_err)? .new_scan() .map_err(to_py_err)?; - Ok(PyVectorScan { - inner, - options: self - .config - .options - .iter() - .flat_map(|options| options.clone()) - .collect(), - }) + Ok(PyVectorScan { inner }) } fn new_batch_vector_search_read(&self) -> PyResult { let mut builder = self.config.batch(&self.table).map_err(to_py_err)?; @@ -324,7 +409,6 @@ impl PyBatchVectorSearchBuilder { #[pyclass(name = "VectorScan", module = "pypaimon_rust.datafusion")] pub struct PyVectorScan { inner: VectorScan, - options: HashMap, } #[pymethods] @@ -333,26 +417,32 @@ impl PyVectorScan { let inner = py .detach(|| runtime().block_on(self.inner.plan())) .map_err(to_py_err)?; - Ok(PyVectorScanPlan { - inner, - options: self.options.clone(), - }) + Ok(PyVectorScanPlan { inner }) } } #[pyclass(name = "VectorScanPlan", module = "pypaimon_rust.datafusion")] pub struct PyVectorScanPlan { inner: VectorScanPlan, - options: HashMap, } #[pymethods] impl PyVectorScanPlan { - fn _metadata(&self, py: Python<'_>) -> PyResult { - let splits = py - .detach(|| runtime().block_on(self.inner.splits(&self.options))) - .map_err(to_py_err)?; - search_plan_metadata_json(self.inner.snapshot(), self.inner.snapshot_id(), splits) + #[new] + #[pyo3(signature = (snapshot_json, splits_json))] + fn new(snapshot_json: Option<&str>, splits_json: &str) -> PyResult { + let (snapshot, splits) = search_plan_from_json(snapshot_json, splits_json)?; + Ok(Self { + inner: VectorScanPlan::new(snapshot, splits).map_err(to_py_err)?, + }) + } + + fn metadata(&self) -> PyResult { + search_plan_metadata_json( + self.inner.snapshot(), + self.inner.snapshot_id(), + self.inner.splits(), + ) } fn snapshot_id(&self) -> Option { @@ -367,22 +457,25 @@ pub struct PyVectorRead { #[pymethods] impl PyVectorRead { - #[pyo3(signature = (plan, executor=None))] - fn read_plan( + fn read_plan(&self, py: Python<'_>, plan: &PyVectorScanPlan) -> PyResult { + py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + .map(PyGlobalIndexResult::new) + .map_err(to_py_err) + } + + fn _read_plan_with_executor( &self, py: Python<'_>, plan: &PyVectorScanPlan, - executor: Option>, + executor: Py, ) -> PyResult { - let Some(callback) = executor else { - return py - .detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - .map(PyGlobalIndexResult::new) - .map_err(to_py_err); - }; - let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); + let executor = crate::search_task::PythonSearchTaskExecutor::new(executor); let result = py.detach(|| { - runtime().block_on(self.inner.read_with_executor(plan.inner.clone(), &executor)) + runtime().block_on(paimon::table::internal::read_vector_with_executor( + &self.inner, + plan.inner.clone(), + &executor, + )) }); if let Some(error) = executor.take_error() { return Err(error); @@ -398,27 +491,187 @@ pub struct PyBatchVectorRead { #[pymethods] impl PyBatchVectorRead { - #[pyo3(signature = (plan, executor=None))] fn read_batch_plan( &self, py: Python<'_>, plan: &PyVectorScanPlan, - executor: Option>, ) -> PyResult> { - let result = if let Some(callback) = executor { - let executor = crate::search_task::PythonSearchTaskExecutor::new(callback); - let result = py.detach(|| { - runtime().block_on(self.inner.read_with_executor(plan.inner.clone(), &executor)) - }); - if let Some(error) = executor.take_error() { - return Err(error); - } - result - } else { - py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - }; + py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + .map(|results| results.into_iter().map(PyGlobalIndexResult::new).collect()) + .map_err(to_py_err) + } + + fn _read_batch_plan_with_executor( + &self, + py: Python<'_>, + plan: &PyVectorScanPlan, + executor: Py, + ) -> PyResult> { + let executor = crate::search_task::PythonSearchTaskExecutor::new(executor); + let result = py.detach(|| { + runtime().block_on(paimon::table::internal::read_batch_vector_with_executor( + &self.inner, + plan.inner.clone(), + &executor, + )) + }); + if let Some(error) = executor.take_error() { + return Err(error); + } result .map(|results| results.into_iter().map(PyGlobalIndexResult::new).collect()) .map_err(to_py_err) } } + +#[cfg(test)] +mod source_plan_tests { + use super::*; + use crate::full_text_search::PyFullTextScanPlan; + use paimon::spec::{BinaryRow, CommitKind, IndexFileMeta, Snapshot}; + use pyo3::types::PyModule; + use serde_json::{json, Value}; + + fn snapshot_json() -> String { + serde_json::to_string( + &Snapshot::builder() + .version(3) + .id(7) + .schema_id(2) + .base_manifest_list("base".into()) + .delta_manifest_list("delta".into()) + .commit_user("test".into()) + .commit_identifier(1) + .commit_kind(CommitKind::APPEND) + .time_millis(0) + .next_row_id(Some(10)) + .build(), + ) + .unwrap() + } + + fn index_file() -> IndexFileMeta { + IndexFileMeta { + index_type: "lucene-vector".into(), + file_name: "index-1".into(), + file_size: 1024, + row_count: 10, + deletion_vectors_ranges: None, + external_path: Some("s3://bucket/index-1".into()), + global_index_meta: None, + } + } + + #[test] + fn source_metadata_reuses_java_index_and_split_codecs() { + Python::attach(|_| { + let split = DataSplit::builder() + .with_snapshot(7) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("file:/tmp/bucket-0".into()) + .with_total_buckets(1) + .with_data_files(vec![]) + .build() + .unwrap(); + let sources = vec![ + SearchScanSplit::Index { + column: "embedding".into(), + row_range: RowRange::new(0, 9), + index_files: vec![index_file()], + scalar_index_files: vec![], + }, + SearchScanSplit::Raw { + row_ranges: vec![RowRange::new(2, 4)], + scalar_index_files: vec![index_file()], + index_type: Some("lucene-vector".into()), + }, + SearchScanSplit::PrimaryKey { + data_split: split, + payloads: vec![index_file()], + uncovered_data_files: vec!["data-1.parquet".into()], + row_ranges_by_file: HashMap::from([( + "data-1.parquet".into(), + vec![RowRange::new(1, 2)], + )]), + }, + ]; + let snapshot: Snapshot = serde_json::from_str(&snapshot_json()).unwrap(); + let metadata = search_plan_metadata_json(Some(&snapshot), Some(7), &sources).unwrap(); + let value: Value = serde_json::from_str(&metadata).unwrap(); + let (restored_snapshot, restored) = + search_plan_from_json(Some(&snapshot_json()), &value["splits"].to_string()) + .unwrap(); + assert_eq!(restored_snapshot.as_ref(), Some(&snapshot)); + assert_eq!( + serde_json::from_str::( + &search_plan_metadata_json(restored_snapshot.as_ref(), Some(7), &restored) + .unwrap() + ) + .unwrap(), + value, + ); + assert_eq!( + value["splits"][0]["index_files"][0], + json!(index_file().to_serialized_row_data().unwrap()) + ); + }); + } + + #[test] + fn python_source_plan_constructors_capture_snapshot_and_splits() { + Python::attach(|py| { + let module = PyModule::new(py, "plans").unwrap(); + module.add_class::().unwrap(); + module.add_class::().unwrap(); + let sources = json!([{ + "kind": "raw", "row_ranges": [{"from": 2, "to": 4}], + "scalar_index_files": [], "index_type": null, + }]) + .to_string(); + for name in ["VectorScanPlan", "FullTextScanPlan"] { + let class = module.getattr(name).unwrap(); + let plan = class.call1((snapshot_json(), &sources)).unwrap(); + assert_eq!( + plan.call_method0("snapshot_id") + .unwrap() + .extract::() + .unwrap(), + 7 + ); + let metadata: String = plan.call_method0("metadata").unwrap().extract().unwrap(); + let value: Value = serde_json::from_str(&metadata).unwrap(); + assert_eq!(value["splits"].to_string(), sources); + assert_eq!(value["snapshot"]["id"], json!(7)); + let copied = class + .call1((value["snapshot"].to_string(), value["splits"].to_string())) + .unwrap(); + let copied_metadata: String = + copied.call_method0("metadata").unwrap().extract().unwrap(); + assert_eq!(metadata, copied_metadata); + let empty = class.call1((py.None(), "[]")).unwrap(); + assert!(empty.call_method0("snapshot_id").unwrap().is_none()); + assert!(class.call1((py.None(), &sources)).is_err()); + } + }); + } + + #[test] + fn source_plan_decode_rejects_corrupt_codecs_and_ranges() { + Python::attach(|_| { + for value in [ + json!([{"kind": "unknown"}]), + json!([{"kind": "index", "column": "v", "row_range": {"from": 2, "to": 1}, "index_files": [], "scalar_index_files": []}]), + json!([{"kind": "raw", "row_ranges": [{"from": -1, "to": 1}], "scalar_index_files": [], "index_type": null}]), + json!([{"kind": "index", "column": "v", "row_range": {"from": 0, "to": 1}, "index_files": [[1, 2, 3]], "scalar_index_files": []}]), + json!([{"kind": "primary-key", "data_split": [1, 2, 3], "payloads": [], "uncovered_data_files": [], "row_ranges_by_file": {}}]), + ] { + assert!( + search_plan_from_json(Some(&snapshot_json()), &value.to_string()).is_err(), + "{value}" + ); + } + assert!(search_plan_from_json(Some("invalid snapshot"), "[]").is_err()); + }); + } +} diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index 34644e6c1..662a6e296 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -45,6 +45,9 @@ storage-all = [ "storage-oss-cpp", ] fulltext = ["dep:paimon-ftindex-core"] +# Unstable bridge used by the repository's language bindings. It is not part +# of the public table API or the default crate surface. +internal-api = [] vortex = ["dep:vortex"] storage-memory = ["opendal/services-memory"] diff --git a/crates/paimon/src/table/batch_vector_search_builder.rs b/crates/paimon/src/table/batch_vector_search_builder.rs index 2d42bc40c..4853a27f3 100644 --- a/crates/paimon/src/table/batch_vector_search_builder.rs +++ b/crates/paimon/src/table/batch_vector_search_builder.rs @@ -122,6 +122,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { self.include_row_ids.as_ref(), self.prepared_filter.as_ref(), ) + .map(|scan| scan.with_options(&self.options)) } /// Create an owned batch reader; result i belongs to input query i. diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index de3513eda..912f8045f 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -216,6 +216,7 @@ impl DeVectorRead { next_row_id, timing, skip_search, + source_splits, } = plan; let table = Arc::new(table); let make_results = |results: Vec| { @@ -261,6 +262,7 @@ impl DeVectorRead { &index_entries, &vector_searches, executor, + source_splits.as_deref(), ) .await? }; @@ -678,8 +680,14 @@ async fn evaluate_batch_vector_search( index_entries: &[IndexManifestEntry], vector_searches: &[VectorSearch], ) -> crate::Result> { - evaluate_batch_vector_search_with_executor(evaluation, index_entries, vector_searches, None) - .await + evaluate_batch_vector_search_with_executor( + evaluation, + index_entries, + vector_searches, + None, + None, + ) + .await } async fn evaluate_batch_vector_search_with_executor( @@ -687,6 +695,7 @@ async fn evaluate_batch_vector_search_with_executor( index_entries: &[IndexManifestEntry], vector_searches: &[VectorSearch], executor: Option<&dyn SearchTaskExecutor>, + sources: Option<&[crate::table::SearchScanSplit]>, ) -> crate::Result> { let timing_enabled = vector_search_timing_enabled(); let total_start = timing_enabled.then(Instant::now); @@ -743,7 +752,8 @@ async fn evaluate_batch_vector_search_with_executor( }) .collect(); - if vector_entries.is_empty() && search_mode == GlobalIndexSearchMode::Fast { + if sources.is_none() && vector_entries.is_empty() && search_mode == GlobalIndexSearchMode::Fast + { return Ok(vec![ScoredRowIds::empty(); vector_searches.len()]); } @@ -760,6 +770,36 @@ async fn evaluate_batch_vector_search_with_executor( .collect(), ); let mut indexed_searches = Cow::Borrowed(vector_searches); + // Source plans carry global payload addresses, without retaining the + // planner's partition predicate or its serialized partition rows. Restrict + // ANN candidates from fixed-snapshot data metadata; the raw path applies + // the full predicate itself and must keep its original source ranges. + if sources.is_some() && !vector_ranges.is_empty() { + if let (Some(table), Some(filter)) = (evaluation.table, evaluation.filter) { + let (partition, _) = crate::table::bucket_filter::split_partition_and_data_predicates( + filter.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ); + if let Some(partition) = partition { + let partition_ids = Arc::new( + super::search_scan_split::partition_row_ids_for_search_sources( + table, + &partition, + vector_ranges.clone(), + ) + .await?, + ); + for search in indexed_searches.to_mut() { + let selected = match search.effective_include_row_ids() { + Some(include) => Arc::new(partition_ids.as_ref() & include), + None => Arc::clone(&partition_ids), + }; + search.set_shared_include_row_ids(selected); + } + } + } + } let mut route_indexed_to_raw = false; // Partition predicates have already pruned manifest entries. They do not // make an otherwise exact scalar lookup candidate-only. @@ -787,7 +827,11 @@ async fn evaluate_batch_vector_search_with_executor( Some(row_ids) => { let row_ids = Arc::new(row_ids); for search in indexed_searches.to_mut() { - search.set_shared_include_row_ids(Arc::clone(&row_ids)); + if let Some(include) = search.effective_include_row_ids() { + search.set_shared_include_row_ids(Arc::new(row_ids.as_ref() & include)); + } else { + search.set_shared_include_row_ids(Arc::clone(&row_ids)); + } } } None => route_indexed_to_raw = true, @@ -798,8 +842,11 @@ async fn evaluate_batch_vector_search_with_executor( let deleted_row_index = if core_options.data_evolution_enabled() { match evaluation.table { Some(table) => { - let ranges = - deleted_row_ranges_for_data_evolution_dvs(table, index_entries).await?; + let ranges = if sources.is_some() { + super::search_scan_split::deleted_row_ranges_for_search_sources(table).await? + } else { + deleted_row_ranges_for_data_evolution_dvs(table, index_entries).await? + }; (!ranges.is_empty()).then(|| RowRangeIndex::create(ranges)) } None => None, @@ -1030,46 +1077,62 @@ async fn evaluate_batch_vector_search_with_executor( refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); } - if route_indexed_to_raw + if sources.is_some() + || route_indexed_to_raw || search_mode != GlobalIndexSearchMode::Fast || scalar_search_mode.is_some_and(|mode| mode != GlobalIndexSearchMode::Fast) { let raw_fallback_start = timing_enabled.then(Instant::now); - let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail - || scalar_search_mode == Some(GlobalIndexSearchMode::Detail) - { - let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { - message: "Vector raw search in detail mode requires table context".to_string(), - source: None, - })?; - detail_data_ranges_for_table(table).await? + let mut raw_ranges = if let Some(sources) = sources { + sources + .iter() + .filter_map(|source| match source { + crate::table::SearchScanSplit::Raw { row_ranges, .. } => { + Some(row_ranges.as_slice()) + } + _ => None, + }) + .flatten() + .cloned() + .collect::>() } else { - Vec::new() - }; - let field_ids = HashSet::from([field_id]); - let mut raw_ranges = unindexed_ranges_for_global_index_entries( - index_entries, - &field_ids, - search_mode, - evaluation.next_row_id, - &detail_ranges, - is_vector_global_index_file, - ); - if let (Some(filter), Some(scalar_mode)) = (&scalar_filter, scalar_search_mode) { - let mut scalar_ranges = scalar_unindexed_ranges( - evaluation, + let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail + || scalar_search_mode == Some(GlobalIndexSearchMode::Detail) + { + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Vector raw search in detail mode requires table context".to_string(), + source: None, + })?; + detail_data_ranges_for_table(table).await? + } else { + Vec::new() + }; + let field_ids = HashSet::from([field_id]); + let mut ranges = unindexed_ranges_for_global_index_entries( index_entries, - filter, - scalar_mode, + &field_ids, + search_mode, + evaluation.next_row_id, &detail_ranges, + is_vector_global_index_file, ); - // Scalar FULL/DETAIL also recover gaps inside vector coverage. In - // vector FAST they must not expand beyond the indexed corpus. - if search_mode == GlobalIndexSearchMode::Fast { - scalar_ranges = intersect_sorted_ranges(&scalar_ranges, &vector_ranges); + if let (Some(filter), Some(scalar_mode)) = (&scalar_filter, scalar_search_mode) { + let mut scalar_ranges = scalar_unindexed_ranges( + evaluation, + index_entries, + filter, + scalar_mode, + &detail_ranges, + ); + // Scalar FULL/DETAIL also recover gaps inside vector coverage. In + // vector FAST they must not expand beyond the indexed corpus. + if search_mode == GlobalIndexSearchMode::Fast { + scalar_ranges = intersect_sorted_ranges(&scalar_ranges, &vector_ranges); + } + ranges.extend(scalar_ranges); } - raw_ranges.extend(scalar_ranges); - } + ranges + }; if route_indexed_to_raw { raw_ranges.extend(vector_ranges); } @@ -1080,11 +1143,29 @@ async fn evaluate_batch_vector_search_with_executor( source: None, })?; let metric_start = timing_enabled.then(Instant::now); + let mut metric_options = vector_searches[0].options.clone(); + if configured_vector_index_type(evaluation.table_options, &metric_options, field_name) + .is_none() + { + if let Some(index_type) = + sources + .into_iter() + .flatten() + .find_map(|source| match source { + crate::table::SearchScanSplit::Raw { index_type, .. } => { + index_type.as_ref() + } + _ => None, + }) + { + metric_options.insert("index_type".into(), index_type.clone()); + } + } let metric = resolve_raw_vector_metric( evaluation.file_io, table_path, evaluation.table_options, - &vector_searches[0].options, + &metric_options, index_entries, field_id, field_name, @@ -2666,6 +2747,7 @@ impl TaskVectorQueries { .collect::>>()?; Ok(Self { queries, filters }) } + #[cfg(any(feature = "internal-api", test))] fn restore(&self) -> crate::Result> { let filters = self .filters @@ -2728,6 +2810,7 @@ struct VectorTaskScores { results: Vec, } impl VectorTaskScores { + #[cfg(any(feature = "internal-api", test))] fn capture(results: Vec, metric: Option) -> Self { Self { metric, @@ -2758,6 +2841,7 @@ impl VectorTaskScores { } } +#[cfg(any(feature = "internal-api", test))] pub(super) async fn execute_search_task( table: &Table, kind: &str, diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs index a7f56c8ad..218b660b2 100644 --- a/crates/paimon/src/table/de_vector_read/tests.rs +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -1969,25 +1969,28 @@ async fn scheduled_vector_selects_global_candidates_before_rerank() { .with_query_vector(vec![0.0, 1.0]) .with_limit(1) .with_options(HashMap::from([("refine_factor".into(), "2".into())])); - let mut plan = builder.new_scan().unwrap().plan().await.unwrap(); - let crate::table::vector_scan::VectorScanWork::DataEvolution(work) = &mut plan.work else { - unreachable!() - }; - let duplicate = work - .index_entries + let plan = builder.new_scan().unwrap().plan().await.unwrap(); + let mut sources = plan.splits().to_vec(); + let duplicate = sources .iter() - .find(|entry| entry.index_file.index_type == IVF_FLAT_IDENTIFIER) + .find(|source| { + matches!(source, + crate::table::SearchScanSplit::Index { index_files, .. } + if index_files.iter().any(|file| file.index_type == IVF_FLAT_IDENTIFIER)) + }) .unwrap() .clone(); for _ in 0..SEARCH_TASK_WINDOW + 1 { - work.index_entries.push(duplicate.clone()); + sources.push(duplicate.clone()); } + let entries = crate::table::search_scan_split::global_entries_from_splits(&sources).unwrap(); + let plan = crate::table::VectorScanPlan::new(plan.snapshot().cloned(), sources).unwrap(); let metric = resolve_raw_vector_metric( table.file_io(), table.location(), table.schema().options(), &HashMap::new(), - &work.index_entries, + &entries, find_field_id_by_name(table.schema().fields(), "embedding").unwrap(), "embedding", ) diff --git a/crates/paimon/src/table/de_vector_scan.rs b/crates/paimon/src/table/de_vector_scan.rs index 1b2af3329..b9d48b043 100644 --- a/crates/paimon/src/table/de_vector_scan.rs +++ b/crates/paimon/src/table/de_vector_scan.rs @@ -67,7 +67,7 @@ impl Table { include_row_ids: Arc::new(RoaringTreemap::new()), }); }; - let table = self.copy_with_resolved_snapshot(&snapshot).await?; + let table = self.copy_with_pinned_snapshot(Some(&snapshot)); let include_row_ids = matching_row_ids_for_filter(&table, &filter, None).await?; Ok(PreparedVectorSearchFilter { table, @@ -110,6 +110,8 @@ pub(super) struct DeVectorScanPlan { pub(super) timing: Option, pub(super) next_row_id: Option, pub(super) skip_search: bool, + /// Source metadata supplied by a transportable public plan. + pub(super) source_splits: Option>, } impl DeVectorScanPlan { @@ -122,6 +124,7 @@ impl DeVectorScanPlan { filter: None, next_row_id: None, timing, + source_splits: None, } } } @@ -194,11 +197,7 @@ impl Scan for DeVectorScan { }; plan.table = match self.prepared_filter.as_ref() { Some(prepared) => prepared.table().clone(), - None => { - execution_table - .copy_with_resolved_snapshot(&snapshot) - .await? - } + None => execution_table.copy_with_pinned_snapshot(Some(&snapshot)), }; plan.next_row_id = snapshot.next_row_id(); diff --git a/crates/paimon/src/table/de_vector_scan/tests.rs b/crates/paimon/src/table/de_vector_scan/tests.rs index 0e16c0452..3737ede2c 100644 --- a/crates/paimon/src/table/de_vector_scan/tests.rs +++ b/crates/paimon/src/table/de_vector_scan/tests.rs @@ -643,3 +643,367 @@ async fn scalar_coverage_gaps_follow_scalar_and_vector_modes() { } } } + +#[tokio::test] +async fn source_plan_round_trip_allows_reader_queries_limits_and_filters() { + let table = de_vector_table().await; + let mut builder = table.new_vector_search_builder(); + builder.with_vector_column("embedding"); + let exported = builder.new_scan().unwrap().plan().await.unwrap(); + let constructed = + crate::table::VectorScanPlan::new(exported.snapshot().cloned(), exported.splits().to_vec()) + .unwrap(); + let wire = serde_json::to_vec(&constructed).unwrap(); + let restored: crate::table::VectorScanPlan = serde_json::from_slice(&wire).unwrap(); + assert_eq!(restored.snapshot(), exported.snapshot()); + assert_eq!( + serde_json::to_value(restored.splits()).unwrap(), + serde_json::to_value(exported.splits()).unwrap() + ); + + for query in [vec![1.0, 0.0], vec![0.0, 1.0]] { + let mut filtered = table.new_vector_search_builder(); + filtered + .with_vector_column("embedding") + .with_query_vector(query) + .with_limit(1) + .with_filter(id_gt_filter(&table, 1)) + .with_options(HashMap::from([("ivf.nprobe".into(), "2".into())])); + let expected = filtered.execute().await.unwrap(); + let actual = filtered + .new_read() + .unwrap() + .read(restored.clone()) + .await + .unwrap(); + assert_eq!( + actual.row_ids().unwrap().row_ids, + expected.row_ids().unwrap().row_ids + ); + assert_eq!( + actual.row_ids().unwrap().scores, + expected.row_ids().unwrap().scores + ); + } +} + +#[tokio::test] +async fn explicit_raw_sources_and_empty_sources_never_discover_other_work() { + let table = de_vector_table().await; + let snapshot = table.resolve_read_snapshot().await.unwrap().unwrap(); + let source = crate::table::SearchScanSplit::Raw { + row_ranges: vec![crate::table::RowRange::new(1, 1)], + scalar_index_files: vec![], + index_type: Some("ivf-flat".into()), + }; + let plan = crate::table::VectorScanPlan::new(Some(snapshot.clone()), vec![source]).unwrap(); + // The reader's default FAST mode must not reinterpret an explicit raw source + // as absent, nor add the index shards present in the snapshot manifest. + let mut builder = table.new_batch_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(3); + let reader = builder.new_read().unwrap(); + for result in reader.read(plan).await.unwrap() { + assert_eq!(result.row_ids().unwrap().row_ids, vec![1]); + assert_eq!(result.snapshot_id(), Some(snapshot.id())); + } + for snapshot in [Some(snapshot), None] { + let expected_id = snapshot.as_ref().map(crate::spec::Snapshot::id); + let plan = crate::table::VectorScanPlan::new(snapshot, vec![]).unwrap(); + let results = reader.read(plan).await.unwrap(); + assert_eq!(results.len(), 2); + for result in results { + assert!(result.is_empty()); + assert_eq!(result.snapshot_id(), expected_id); + } + } +} + +#[tokio::test] +async fn transported_vector_sources_pin_snapshot_and_keep_bound_read_schema() { + let table = de_vector_table().await; + // The fixture bypasses a catalog and holds schema zero only in memory. + // A renamed reader must be able to load the historical data-file schema. + table + .file_io() + .new_output(&table.schema_manager().schema_path(table.schema().id())) + .unwrap() + .write(serde_json::to_vec(table.schema()).unwrap().into()) + .await + .unwrap(); + let mut old_builder = table.new_vector_search_builder(); + old_builder.with_vector_column("embedding"); + let plan = old_builder.new_scan().unwrap().plan().await.unwrap(); + let schema = table + .schema() + .apply_changes(vec![crate::spec::SchemaChange::rename_column( + "embedding".into(), + "renamed_vector".into(), + )]) + .unwrap(); + let renamed = table + .copy_with_resolved_schema(schema, table.branch()) + .unwrap(); + let mut builder = renamed.new_vector_search_builder(); + builder + .with_vector_column("renamed_vector") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(3) + .with_options(HashMap::from([ + ("refine_factor".into(), "2".into()), + ("ivf.nprobe".into(), "2".into()), + ])); + let renamed_plan = builder.new_scan().unwrap().plan().await.unwrap(); + assert_eq!(renamed_plan.snapshot_id(), plan.snapshot_id()); + let captured_id = plan.snapshot_id().unwrap(); + let wire = serde_json::to_vec(&plan).unwrap(); + // Snapshot JSON is unnecessary after capture, including on another process. + let manager = table.snapshot_manager(); + for id in manager.list_all_ids().await.unwrap() { + table + .file_io() + .delete_file(&manager.snapshot_path(id)) + .await + .unwrap(); + } + let plan: crate::table::VectorScanPlan = serde_json::from_slice(&wire).unwrap(); + let result = builder.new_read().unwrap().read(plan).await.unwrap(); + assert_eq!(result.snapshot_id(), Some(captured_id)); + assert_eq!(result.row_ids().unwrap().row_ids.len(), 3); + let batches: Vec = result + .new_read_builder() + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert!(batches + .iter() + .all(|batch| batch.column_by_name("renamed_vector").is_some() + && batch.column_by_name("embedding").is_none())); +} + +#[test] +fn transported_vector_plan_rejects_missing_snapshot_and_invalid_ranges() { + let raw = crate::table::SearchScanSplit::Raw { + row_ranges: vec![crate::table::RowRange::new(0, 0)], + scalar_index_files: vec![], + index_type: None, + }; + assert!(crate::table::VectorScanPlan::new(None, vec![raw.clone()]).is_err()); + let wire = + serde_json::json!({"snapshot": null, "splits": [serde_json::to_value(raw).unwrap()]}); + assert!(serde_json::from_value::(wire).is_err()); +} + +#[tokio::test] +async fn transported_vector_sources_apply_snapshot_deletion_vectors() { + let table = de_vector_table().await.copy_with_options(HashMap::from([( + "deletion-vectors.enabled".into(), + "true".into(), + )])); + let write = table.new_write_builder(); + let delete = write.new_update().unwrap(); + write + .new_commit() + .commit(delete.delete_by_row_id(vec![0]).await.unwrap()) + .await + .unwrap(); + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(3) + .with_options(HashMap::from([("ivf.nprobe".into(), "2".into())])); + let exported = builder.new_scan().unwrap().plan().await.unwrap(); + let plan = serde_json::from_slice(&serde_json::to_vec(&exported).unwrap()).unwrap(); + let reader = builder.new_read().unwrap(); + let indexed = reader.read(plan).await.unwrap(); + let mut expected = indexed.row_ids().unwrap().row_ids.clone(); + expected.sort_unstable(); + assert_eq!(expected, vec![1, 2]); + let raw = crate::table::VectorScanPlan::new( + exported.snapshot().cloned(), + vec![crate::table::SearchScanSplit::Raw { + row_ranges: vec![crate::table::RowRange::new(0, 2)], + scalar_index_files: vec![], + index_type: Some("ivf-flat".into()), + }], + ) + .unwrap(); + let actual = reader.read(raw).await.unwrap(); + let mut actual = actual.row_ids().unwrap().row_ids.clone(); + actual.sort_unstable(); + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn reader_partition_filter_is_intersected_with_scalar_candidates() { + use crate::spec::{ + ArrayType, DataType, Datum, FloatType, IntType, PredicateBuilder, Schema, TableSchema, + }; + let schema = Schema::builder() + .column("p", DataType::Int(IntType::new())) + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .partition_keys(["p"]) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .option("global-index.enabled", "true") + .option("ivf-flat.dimension", "2") + .option("ivf-flat.nlist", "1") + .build() + .unwrap(); + let table = Table::new( + crate::io::FileIOBuilder::new("memory").build().unwrap(), + crate::catalog::Identifier::new("default", "partition_source"), + "memory:/partition_source".into(), + TableSchema::new(0, &schema), + None, + ); + let element = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vectors = ListBuilder::new(Float32Builder::new()).with_field(element.clone()); + for vector in [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]] { + vectors.values().append_slice(&vector); + vectors.append(true); + } + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("p", ArrowDataType::Int32, false), + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element), true), + ])), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 1])), + Arc::new(Int32Array::from(vec![3, 2, 3])), + Arc::new(vectors.finish()), + ], + ) + .unwrap(); + let mut writer = TableWrite::new(&table, "partition_source".into()).unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + TableCommit::new(table.clone(), "partition_source".into()) + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + for (column, kind) in [("embedding", "ivf-flat"), ("id", "btree")] { + table + .new_global_index_build_builder() + .with_index_column(column) + .with_index_type(kind) + .execute() + .await + .unwrap(); + } + let mut builder = table.new_vector_search_builder(); + builder.with_vector_column("embedding"); + let exported = builder.new_scan().unwrap().plan().await.unwrap(); + let predicates = PredicateBuilder::new(table.schema().fields()); + builder + .with_query_vector(vec![1.0, 0.0]) + .with_limit(3) + .with_filter(crate::spec::Predicate::and(vec![ + predicates.equal("p", Datum::Int(1)).unwrap(), + predicates.greater_than("id", Datum::Int(1)).unwrap(), + ])); + let result = builder.new_read().unwrap().read(exported).await.unwrap(); + let batches: Vec = result + .new_read_builder() + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![2, 3]); +} + +#[tokio::test] +async fn prepared_filter_rejects_a_source_plan_from_another_snapshot() { + let table = de_vector_table().await; + let prepared = table + .prepare_vector_search_filter(id_gt_filter(&table, 1)) + .await + .unwrap(); + let mut builder = table.new_batch_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0]]) + .with_limit(2) + .with_prepared_filter(prepared); + let empty = crate::table::VectorScanPlan::new(None, vec![]).unwrap(); + let error = builder.new_read().unwrap().read(empty).await.err().unwrap(); + assert!( + error.to_string().contains("same captured snapshot"), + "{error}" + ); +} + +#[tokio::test] +async fn transported_vector_sources_keep_pk_positions_separate_from_global_ids() { + let table = de_vector_table().await; + let snapshot = table.resolve_read_snapshot().await.unwrap().unwrap(); + let split = |snapshot_id| crate::table::SearchScanSplit::PrimaryKey { + data_split: crate::table::DataSplitBuilder::new() + .with_snapshot(snapshot_id) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/pk_sources/bucket-0".into()) + .with_data_files(vec![]) + .build() + .unwrap(), + payloads: vec![], + uncovered_data_files: vec![], + row_ranges_by_file: HashMap::new(), + }; + assert!(crate::table::VectorScanPlan::new( + Some(snapshot.clone()), + vec![split(snapshot.id() + 1)] + ) + .is_err()); + let plan = crate::table::VectorScanPlan::new(Some(snapshot.clone()), vec![split(0)]).unwrap(); + let plan = serde_json::from_slice(&serde_json::to_vec(&plan).unwrap()).unwrap(); + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1); + let error = builder.new_read().unwrap().read(plan).await.err().unwrap(); + assert!( + error.to_string().contains("primary-key physical"), + "{error}" + ); + let mixed = crate::table::VectorScanPlan::new( + Some(snapshot), + vec![ + split(0), + crate::table::SearchScanSplit::Raw { + row_ranges: vec![crate::table::RowRange::new(0, 0)], + scalar_index_files: vec![], + index_type: None, + }, + ], + ); + assert!(mixed.is_err()); +} diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 858b93d42..5fc928a69 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -50,6 +50,7 @@ mod execution; mod filter; mod read; mod scan; +#[cfg(any(feature = "internal-api", test))] pub(super) use execution::execute_search_task; pub use read::FullTextRead; pub use scan::{FullTextScan, FullTextScanPlan}; @@ -221,6 +222,7 @@ struct FullTextSearchEvaluation<'a> { next_row_id: Option, partition_filter: Option<&'a Predicate>, row_filter: Option<&'a Predicate>, + source_splits: Option<&'a [crate::table::SearchScanSplit]>, } #[cfg(test)] @@ -261,15 +263,24 @@ async fn evaluate_full_text_search_with_executor( .collect(); let has_fulltext_entries = !fulltext_entries.is_empty(); - if !has_fulltext_entries { + if !has_fulltext_entries && evaluation.source_splits.is_none() { + return Ok(SearchResult::empty()); + } + if evaluation + .source_splits + .is_some_and(|splits| splits.is_empty()) + { return Ok(SearchResult::empty()); } let deleted_row_index = if core_options.data_evolution_enabled() { match evaluation.table { Some(table) => { - let ranges = - deleted_row_ranges_for_data_evolution_dvs(table, index_entries).await?; + let ranges = if evaluation.source_splits.is_some() { + super::search_scan_split::deleted_row_ranges_for_search_sources(table).await? + } else { + deleted_row_ranges_for_data_evolution_dvs(table, index_entries).await? + }; (!ranges.is_empty()).then(|| RowRangeIndex::create(ranges)) } None => None, @@ -288,6 +299,27 @@ async fn evaluate_full_text_search_with_executor( matched, )); } + if let (Some(_), Some(partition)) = (evaluation.source_splits, evaluation.partition_filter) { + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Full-text partition filtering requires table context".into(), + source: None, + })?; + let ranges = fulltext_entries + .iter() + .map(|entry| { + let meta = entry.index_file.global_index_meta.as_ref().unwrap(); + RowRange::new(meta.row_range_start, meta.row_range_end) + }) + .collect(); + let matched = super::search_scan_split::partition_row_ids_for_search_sources( + table, partition, ranges, + ) + .await?; + indexed_search.include_row_ids = Some(filter::intersect_include( + indexed_search.include_row_ids.as_ref(), + matched, + )); + } let mut merged = SearchResult::empty(); if !fulltext_entries.is_empty() { let plans = plan_full_text_index_searches( @@ -346,8 +378,12 @@ async fn evaluate_full_text_search_with_executor( } } - if search_mode != GlobalIndexSearchMode::Fast && has_fulltext_entries { - let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail { + if evaluation.source_splits.is_some() + || search_mode != GlobalIndexSearchMode::Fast && has_fulltext_entries + { + let detail_ranges = if evaluation.source_splits.is_none() + && search_mode == GlobalIndexSearchMode::Detail + { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { message: "Full-text raw search in detail mode requires table context".to_string(), source: None, @@ -357,14 +393,26 @@ async fn evaluate_full_text_search_with_executor( Vec::new() }; let field_ids = HashSet::from([field_id]); - let raw_ranges = unindexed_ranges_for_global_index_entries( - index_entries, - &field_ids, - search_mode, - evaluation.next_row_id, - &detail_ranges, - is_full_text_index_file, - ); + let raw_ranges = if let Some(splits) = evaluation.source_splits { + merge_row_ranges( + splits + .iter() + .flat_map(|split| match split { + crate::table::SearchScanSplit::Raw { row_ranges, .. } => row_ranges.clone(), + _ => vec![], + }) + .collect(), + ) + } else { + unindexed_ranges_for_global_index_entries( + index_entries, + &field_ids, + search_mode, + evaluation.next_row_id, + &detail_ranges, + is_full_text_index_file, + ) + }; if !raw_ranges.is_empty() { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { message: "Full-text raw search requires table context".to_string(), @@ -931,6 +979,7 @@ mod tests { next_row_id: Some(10), partition_filter: None, row_filter: None, + source_splits: None, }, &[], &search, @@ -1086,6 +1135,7 @@ mod tests { next_row_id: Some(102), partition_filter: None, row_filter: None, + source_splits: None, }, &[entry], &search, @@ -1172,6 +1222,7 @@ mod tests { next_row_id: Some(102), partition_filter: None, row_filter: None, + source_splits: None, }, &[entry, missing_empty_filter_entry], &search, @@ -1434,6 +1485,51 @@ mod tests { .is_empty()); } + #[tokio::test] + async fn pk_full_text_source_plan_resolves_unknown_split_snapshot_from_capture() { + let table = pk_full_text_table("pk_ft_unknown_split_snapshot", &[]); + let snapshot = crate::spec::Snapshot::builder() + .version(3) + .id(7) + .schema_id(0) + .base_manifest_list(String::new()) + .delta_manifest_list(String::new()) + .commit_user("capture".into()) + .commit_identifier(1) + .commit_kind(crate::spec::CommitKind::APPEND) + .time_millis(0) + .build(); + let split = |snapshot_id| crate::table::SearchScanSplit::PrimaryKey { + data_split: crate::table::DataSplitBuilder::new() + .with_snapshot(snapshot_id) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{}/bucket-0", table.location())) + .with_data_files(vec![]) + .build() + .unwrap(), + payloads: vec![], + uncovered_data_files: vec![], + row_ranges_by_file: Default::default(), + }; + assert!(FullTextScanPlan::new(Some(snapshot.clone()), vec![split(8)]).is_err()); + assert!(FullTextScanPlan::new(None, vec![split(0)]).is_err()); + let plan = FullTextScanPlan::new(Some(snapshot), vec![split(0)]).unwrap(); + let plan: FullTextScanPlan = + serde_json::from_slice(&serde_json::to_vec(&plan).unwrap()).unwrap(); + let result = table + .new_full_text_search_builder() + .with_query("body", r#"{"match":{"query":"alpha"}}"#) + .with_limit(10) + .new_read() + .unwrap() + .read(plan) + .await + .unwrap(); + assert_eq!(result.snapshot_id(), Some(7)); + assert!(result.positions().unwrap().is_empty()); + } + // (d) The PK full-text path produces physical positions, not global row-ids: // Only the ranges-only `execute` API rejects the physical-position result. #[tokio::test] diff --git a/crates/paimon/src/table/full_text_search_builder/execution.rs b/crates/paimon/src/table/full_text_search_builder/execution.rs index e7bd3e5a4..262608985 100644 --- a/crates/paimon/src/table/full_text_search_builder/execution.rs +++ b/crates/paimon/src/table/full_text_search_builder/execution.rs @@ -51,6 +51,7 @@ struct Hits { } impl Hits { + #[cfg(any(feature = "internal-api", test))] fn from_search(result: SearchResult) -> Self { Self { row_ids: result.row_ids, @@ -90,6 +91,7 @@ fn encode_bitmap(ids: &RoaringTreemap) -> crate::Result> { Ok(bytes) } +#[cfg(any(feature = "internal-api", test))] fn decode_bitmap(bytes: Vec) -> crate::Result { RoaringTreemap::deserialize_from(bytes.as_slice()) .map_err(|error| invalid(format!("Invalid full-text task row IDs: {error}"))) @@ -166,6 +168,7 @@ pub(super) async fn search_raw_rows( decode::(&values.remove(0))?.into_search() } +#[cfg(any(feature = "internal-api", test))] pub(crate) async fn execute_search_task( table: &Table, kind: &str, diff --git a/crates/paimon/src/table/full_text_search_builder/parity_tests.rs b/crates/paimon/src/table/full_text_search_builder/parity_tests.rs index 9903a1483..03fc09a14 100644 --- a/crates/paimon/src/table/full_text_search_builder/parity_tests.rs +++ b/crates/paimon/src/table/full_text_search_builder/parity_tests.rs @@ -467,7 +467,7 @@ async fn partition_filters_apply_to_indexed_and_raw_corpora() { } #[tokio::test] -async fn scan_read_pin_snapshot_and_validate_plan_origin() { +async fn scan_read_pins_snapshot_and_keeps_query_configuration_on_reader() { let table = new_table(false); append(&table, 0, 8).await; index(&table, "full-text", "text").await; @@ -488,14 +488,10 @@ async fn scan_read_pin_snapshot_and_validate_plan_origin() { let pb = PredicateBuilder::new(table.schema().fields()); let mut other = query(&table); other.with_filter(pb.equal("id", Datum::Int(1)).unwrap()); - assert!(other - .new_read() - .unwrap() - .read(plan.clone()) - .await - .unwrap_err() - .to_string() - .contains("different table, column or filter")); + assert_eq!( + ids(other.new_read().unwrap().read(plan.clone()).await.unwrap()), + vec![1] + ); let foreign = new_table(false); assert!(query(&foreign) .new_read() @@ -611,7 +607,8 @@ async fn snapshot_pinning_preserves_schema_only_renames() { vec![0, 1, 2, 4, 5, 6, 7] ); let plan = builder.new_scan().unwrap().scan().await.unwrap(); - assert_eq!(plan.table.schema().id(), table.schema().id()); + let plan: FullTextScanPlan = + serde_json::from_slice(&serde_json::to_vec(&plan).unwrap()).unwrap(); assert_eq!( ids(builder.new_read().unwrap().read(plan).await.unwrap()), vec![0, 1, 2, 4, 5, 6, 7] @@ -778,11 +775,13 @@ async fn scheduled_full_text_rejects_malformed_worker_scores() { } #[tokio::test] -async fn scheduled_full_text_rejects_foreign_plan_before_scheduling() { +async fn scheduled_full_text_rejects_wrong_reader_field_before_scheduling() { let table = new_table(false); - let other = new_table(false); - let builder = query(&table); - let plan = query(&other).new_scan().unwrap().scan().await.unwrap(); + append(&table, 0, 4).await; + index(&table, "full-text", "text").await; + let plan = query(&table).new_scan().unwrap().scan().await.unwrap(); + let mut builder = query(&table); + builder.with_text_column("label"); let executor = WireExecutor { table: table.clone(), work: Default::default(), @@ -797,7 +796,7 @@ async fn scheduled_full_text_rejects_foreign_plan_before_scheduling() { .unwrap_err(); assert!( matches!(error, crate::Error::DataInvalid { ref message, .. } - if message.contains("different table")), + if message.contains("reader's field and index type")), "{error}" ); assert!(executor.work.lock().unwrap().is_empty()); @@ -859,12 +858,13 @@ async fn scheduled_full_text_bounds_index_task_windows_and_keeps_all_candidates( builder.with_limit(row_count); let plan = builder.new_scan().unwrap().scan().await.unwrap(); let files = plan - .entries + .splits() .iter() - .filter(|entry| { - entry.kind == FileKind::Add && entry.index_file.index_type == FULL_TEXT_INDEX_TYPE + .flat_map(|split| match split { + crate::table::SearchScanSplit::Index { index_files, .. } => index_files.clone(), + _ => vec![], }) - .map(|entry| entry.index_file.file_name.clone()) + .map(|file| file.file_name) .collect::>(); let read = builder.new_read().unwrap(); let expected = read @@ -911,3 +911,155 @@ async fn scheduled_full_text_bounds_index_task_windows_and_keeps_all_candidates( .collect::>(); assert_eq!(actual_files, files, "windows must retain manifest order"); } + +#[tokio::test] +async fn data_plans_execute_only_the_supplied_sources_after_roundtrip() { + let table = new_table(false); + append(&table, 0, 8).await; + index(&table, "full-text", "text").await; + append(&table, 8, 4).await; + let builder = query(&table); + let scanned = builder.new_scan().unwrap().scan().await.unwrap(); + let snapshot = scanned.snapshot().cloned(); + let index_source = scanned.splits().iter().find(|split| { + matches!(split, crate::table::SearchScanSplit::Index { row_range, .. } if row_range.from() == 0) + }).unwrap().clone(); + let raw_source = scanned + .splits() + .iter() + .find(|split| matches!(split, crate::table::SearchScanSplit::Raw { .. })) + .unwrap() + .clone(); + let index_only = FullTextScanPlan::new(snapshot.clone(), vec![index_source]).unwrap(); + let raw_only = FullTextScanPlan::new(snapshot.clone(), vec![raw_source]).unwrap(); + let empty = FullTextScanPlan::new(snapshot, vec![]).unwrap(); + append(&table, 12, 4).await; + for (plan, expected) in [ + (index_only, vec![0, 1, 2]), + (raw_only, vec![8, 9, 10, 11]), + (empty, vec![]), + ] { + let encoded = serde_json::to_vec(&plan).unwrap(); + let plan: FullTextScanPlan = serde_json::from_slice(&encoded).unwrap(); + let result = builder.new_read().unwrap().read(plan).await.unwrap(); + assert_eq!(result.snapshot_id(), scanned.snapshot_id()); + assert_eq!(ids(result), expected); + } +} + +#[tokio::test] +async fn data_plan_without_snapshot_never_resolves_latest() { + let table = new_table(false); + append(&table, 0, 4).await; + index(&table, "full-text", "text").await; + let builder = query(&table); + let source = builder.new_scan().unwrap().scan().await.unwrap(); + assert!(FullTextScanPlan::new(None, source.splits().to_vec()).is_err()); + let malformed = serde_json::json!({ "snapshot": null, "splits": source.splits() }); + assert!(serde_json::from_value::(malformed).is_err()); + let plan = FullTextScanPlan::new(None, vec![]).unwrap(); + let result = builder.new_read().unwrap().read(plan).await.unwrap(); + assert_eq!(result.snapshot_id(), None); + assert!(result.is_empty()); +} + +#[tokio::test] +async fn data_plan_applies_reader_partition_filter_to_supplied_sources() { + let table = new_table(true); + append(&table, 0, 8).await; + index(&table, "full-text", "text").await; + append(&table, 8, 4).await; + let source = query(&table).new_scan().unwrap().scan().await.unwrap(); + let plan = FullTextScanPlan::new(source.snapshot().cloned(), source.splits().to_vec()).unwrap(); + let pb = PredicateBuilder::new(table.schema().fields()); + let mut builder = query(&table); + builder + .with_partition_filter(pb.equal("pt", Datum::Int(1)).unwrap()) + .unwrap(); + assert_eq!( + user_ids( + &table, + builder.new_read().unwrap().read(plan).await.unwrap() + ) + .await, + vec![4, 5, 6, 7] + ); +} + +#[tokio::test] +async fn data_plan_does_not_reselect_search_sources_from_snapshot_manifest() { + let table = new_table(false).copy_with_options(HashMap::from([( + "deletion-vectors.enabled".into(), + "false".into(), + )])); + append(&table, 0, 8).await; + index(&table, "full-text", "text").await; + let builder = query(&table); + let scanned = builder.new_scan().unwrap().scan().await.unwrap(); + let mut snapshot = serde_json::to_value(scanned.snapshot().unwrap()).unwrap(); + snapshot["indexManifest"] = serde_json::json!("missing-search-source-manifest"); + let snapshot = serde_json::from_value(snapshot).unwrap(); + let source = scanned + .splits() + .iter() + .find(|split| { + matches!( + split, + crate::table::SearchScanSplit::Index { row_range, .. } if row_range.from() == 0 + ) + }) + .unwrap() + .clone(); + let plan = FullTextScanPlan::new(Some(snapshot), vec![source]).unwrap(); + assert_eq!( + ids(builder.new_read().unwrap().read(plan).await.unwrap()), + vec![0, 1, 2], + ); +} + +#[tokio::test] +async fn data_plan_retains_scalar_metadata_for_a_later_reader_filter() { + let table = new_table(false).copy_with_options(HashMap::from([( + "scalar-index.search-mode".into(), + "fast".into(), + )])); + append(&table, 0, 8).await; + index(&table, "full-text", "text").await; + index(&table, "bitmap", "label").await; + let plan = query(&table).new_scan().unwrap().scan().await.unwrap(); + let pb = PredicateBuilder::new(table.schema().fields()); + let mut builder = query(&table); + builder.with_filter(pb.equal("label", Datum::String("keep".into())).unwrap()); + assert_eq!( + ids(builder.new_read().unwrap().read(plan).await.unwrap()), + vec![1, 5, 7], + ); +} + +#[tokio::test] +async fn index_only_data_plan_applies_reader_partitions_across_commits() { + let table = new_table(true); + append(&table, 0, 4).await; + append(&table, 4, 4).await; + index(&table, "full-text", "text").await; + let captured = query(&table).new_scan().unwrap().scan().await.unwrap(); + let plan = FullTextScanPlan::new( + captured.snapshot().cloned(), + captured + .splits() + .iter() + .filter(|split| matches!(split, crate::table::SearchScanSplit::Index { .. })) + .cloned() + .collect(), + ) + .unwrap(); + let pb = PredicateBuilder::new(table.schema().fields()); + for (partition, expected) in [(0, vec![0, 1, 2]), (1, vec![4, 5, 6, 7])] { + let mut reader = query(&table); + reader + .with_partition_filter(pb.equal("pt", Datum::Int(partition)).unwrap()) + .unwrap(); + let result = reader.new_read().unwrap().read(plan.clone()).await.unwrap(); + assert_eq!(user_ids(&table, result).await, expected); + } +} diff --git a/crates/paimon/src/table/full_text_search_builder/read.rs b/crates/paimon/src/table/full_text_search_builder/read.rs index a73146f96..54f46ddb6 100644 --- a/crates/paimon/src/table/full_text_search_builder/read.rs +++ b/crates/paimon/src/table/full_text_search_builder/read.rs @@ -20,8 +20,8 @@ use super::scan::PlanContext; use super::*; -/// An owned reader can execute queries over plans from the same table, column -/// and filter. Changing a query or limit does not require another scan. +/// An owned reader supplies table access and query configuration for captured +/// source plans. Queries and limits can change without another source scan. pub struct FullTextRead { table: Table, context: PlanContext, @@ -90,7 +90,7 @@ impl FullTextRead { /// Execute shard work through a scheduler while retaining the plan's /// snapshot, filtering, and final global ranking in core. - pub async fn read_with_executor( + pub(crate) async fn read_with_executor( &self, plan: FullTextScanPlan, executor: &dyn crate::table::SearchTaskExecutor, @@ -104,46 +104,101 @@ impl FullTextRead { executor: Option<&dyn crate::table::SearchTaskExecutor>, ) -> crate::Result { CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; - CoreOptions::new(plan.table.schema().options()).ensure_read_authorized()?; - if self.context != plan.context { - return Err(crate::Error::DataInvalid { - message: "Full-text plan belongs to a different table, column or filter".into(), - source: None, - }); - } if executor.is_some() && !self.table.schema().core_options().data_evolution_enabled() { return Err(crate::Error::Unsupported { message: "Scheduled full-text search requires a data-evolution table".into(), }); } - let table = std::sync::Arc::new(plan.table.clone()); + let table = self.table.copy_with_pinned_snapshot(plan.snapshot()); + let table = std::sync::Arc::new(table); if self.context.primary_key_field_id.is_some() { - let empty = super::super::pk_full_text_scan::PrimaryKeyFullTextScanPlan { - snapshot_id: 0, - splits: Vec::new(), + let partition_pruner = self.context.filter.as_ref().and_then(|filter| { + let (partition, _) = + super::super::bucket_filter::split_partition_and_data_predicates( + filter.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ); + partition.map(|partition| { + super::super::partition_filter::PartitionFilter::from_predicate( + partition, + &table.schema().partition_fields(), + ) + }) + }); + let mut splits = Vec::new(); + for split in plan.splits() { + let crate::table::SearchScanSplit::PrimaryKey { + data_split, + payloads, + uncovered_data_files, + row_ranges_by_file, + } = split + else { + return Err(crate::Error::DataInvalid { + message: "Primary-key full-text read requires primary-key sources".into(), + source: None, + }); + }; + if let Some(pruner) = &partition_pruner { + if !pruner.matches_entry(&data_split.partition().to_serialized_bytes())? { + continue; + } + } + if !row_ranges_by_file.is_empty() { + return Err(crate::Error::DataInvalid { + message: "Full-text sources do not support vector row-range restrictions" + .into(), + source: None, + }); + } + let data_split = if data_split.snapshot_id() == 0 { + let mut builder = crate::table::DataSplitBuilder::new() + .with_snapshot(plan.snapshot_id().unwrap_or(0)) + .with_partition(data_split.partition().clone()) + .with_bucket(data_split.bucket()) + .with_bucket_path(data_split.bucket_path().to_string()) + .with_total_buckets(data_split.total_buckets()) + .with_data_files(data_split.data_files().to_vec()) + .with_raw_convertible(data_split.raw_convertible()) + .with_streaming(data_split.is_streaming()); + if let Some(dvs) = data_split.data_deletion_files() { + builder = builder.with_data_deletion_files(dvs.to_vec()); + } + builder.build()? + } else { + data_split.clone() + }; + splits.push( + super::super::pk_full_text_scan::PrimaryKeyFullTextSearchSplit::new( + data_split, + payloads.clone(), + uncovered_data_files.clone(), + )?, + ); + } + let source = super::super::pk_full_text_scan::PrimaryKeyFullTextScanPlan { + snapshot_id: plan.snapshot_id().unwrap_or(0), + splits, }; - let source = plan.primary_key_plan.as_ref().unwrap_or(&empty); let materialize = DataFileReader::new( - plan.table.file_io().clone(), - plan.table.schema_manager().clone(), - plan.table.schema().id(), - plan.table.schema().fields().to_vec(), - plan.table.data_fields_extractor(), - plan.table.schema().fields().to_vec(), + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + table.data_fields_extractor(), + table.schema().fields().to_vec(), Vec::new(), ) - .with_table_options(plan.table.schema().options().clone()); + .with_table_options(table.schema().options().clone()); let read = PrimaryKeyFullTextRead::new( - plan.table.file_io().clone(), + table.file_io().clone(), materialize, - plan.table.location().trim_end_matches('/').to_string(), - plan.table - .schema() - .core_options() - .index_file_in_data_file_dir(), + table.location().trim_end_matches('/').to_string(), + table.schema().core_options().index_file_in_data_file_dir(), ); let result = read - .search_route(source, &self.search.query_text, self.search.limit) + .search_route(&source, &self.search.query_text, self.search.limit) .await?; return crate::table::GlobalIndexResult::from_full_text_primary_key( table, @@ -152,20 +207,54 @@ impl FullTextRead { result.splits, ); } + for split in plan.splits() { + match split { + crate::table::SearchScanSplit::Index { index_files, .. } => { + let field_id = + find_field_id_by_name(table.schema().fields(), &self.context.column); + if index_files.iter().any(|file| { + file.index_type != FULL_TEXT_INDEX_TYPE + || file + .global_index_meta + .as_ref() + .map(|meta| meta.index_field_id) + != field_id + }) { + return Err(crate::Error::DataInvalid { + message: + "Full-text source does not match the reader's field and index type" + .into(), + source: None, + }); + } + } + crate::table::SearchScanSplit::PrimaryKey { .. } => { + return Err(crate::Error::DataInvalid { + message: "Global full-text read requires global row-ID sources".into(), + source: None, + }); + } + _ => {} + } + } + let entries = super::super::search_scan_split::global_entries_from_splits(plan.splits())?; + let (partition_filter, row_filter) = + super::scan::bind_filters(&table, self.context.filter.as_ref())?; let mut search = self.search.clone(); search.include_row_ids = self.context.include_row_ids.clone(); let hits = evaluate_full_text_search_with_executor( FullTextSearchEvaluation { - table: Some(&plan.table), - file_io: plan.table.file_io(), - table_path: plan.table.location(), - table_options: self.table.schema().options(), - schema_fields: self.table.schema().fields(), - next_row_id: plan.next_row_id, - partition_filter: plan.partition_filter.as_ref(), - row_filter: plan.row_filter.as_ref(), + table: Some(&table), + file_io: table.file_io(), + table_path: table.location(), + table_options: table.schema().options(), + schema_fields: table.schema().fields(), + next_row_id: plan.snapshot().and_then(crate::spec::Snapshot::next_row_id), + partition_filter: partition_filter.as_ref(), + row_filter: row_filter.as_ref(), + source_splits: Some(plan.splits()), }, - &plan.entries, + &entries, &search, executor, ) diff --git a/crates/paimon/src/table/full_text_search_builder/scan.rs b/crates/paimon/src/table/full_text_search_builder/scan.rs index 3f7926b43..d26a1a5ad 100644 --- a/crates/paimon/src/table/full_text_search_builder/scan.rs +++ b/crates/paimon/src/table/full_text_search_builder/scan.rs @@ -24,8 +24,6 @@ use crate::table::partition_filter::PartitionFilter; #[derive(Clone, PartialEq)] pub(super) struct PlanContext { - location: String, - branch: String, pub(super) column: String, pub(super) filter: Option, pub(super) include_row_ids: Option, @@ -69,8 +67,6 @@ impl PlanContext { }); } Ok(Self { - location: builder.table.location().trim_end_matches('/').into(), - branch: builder.table.branch().into(), column: column.into(), filter: builder.filter.clone(), include_row_ids: builder.include_row_ids.clone(), @@ -79,54 +75,85 @@ impl PlanContext { } } -/// One selected snapshot and its partition-pruned indexes or PK buckets. A plan remains -/// usable after its scan is dropped, and does not re-resolve the latest snapshot. -#[derive(Clone)] +/// A captured snapshot and its authoritative search sources. +/// +/// Plans contain transportable metadata only. The reader supplies the table, +/// filesystem, authorization, schema and query configuration. +#[derive(Clone, serde::Serialize)] pub struct FullTextScanPlan { - pub(super) context: PlanContext, - pub(super) table: Table, - pub(super) entries: Vec, - pub(super) next_row_id: Option, - pub(super) snapshot_id: Option, - pub(super) partition_filter: Option, - pub(super) row_filter: Option, - pub(super) primary_key_plan: - Option, + snapshot: Option, + splits: Vec, +} + +impl<'de> serde::Deserialize<'de> for FullTextScanPlan { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(serde::Deserialize)] + struct PlanData { + snapshot: Option, + splits: Vec, + } + let plan = PlanData::deserialize(deserializer)?; + Self::new(plan.snapshot, plan.splits).map_err(serde::de::Error::custom) + } } impl FullTextScanPlan { + pub fn new( + snapshot: Option, + splits: Vec, + ) -> crate::Result { + if snapshot.is_none() && !splits.is_empty() { + return Err(crate::Error::DataInvalid { + message: "Full-text sources require a captured snapshot".into(), + source: None, + }); + } + let primary_key = splits + .iter() + .any(|split| matches!(split, crate::table::SearchScanSplit::PrimaryKey { .. })); + if primary_key + && splits + .iter() + .any(|split| !matches!(split, crate::table::SearchScanSplit::PrimaryKey { .. })) + { + return Err(crate::Error::DataInvalid { + message: "Full-text plans cannot mix global and primary-key sources".into(), + source: None, + }); + } + if !primary_key { + super::super::search_scan_split::global_entries_from_splits(&splits)?; + } + if let Some(snapshot) = &snapshot { + for split in &splits { + if let crate::table::SearchScanSplit::PrimaryKey { data_split, .. } = split { + if data_split.snapshot_id() != 0 && data_split.snapshot_id() != snapshot.id() { + return Err(crate::Error::DataInvalid { + message: "Full-text source snapshot does not match its plan".into(), + source: None, + }); + } + } + } + } + Ok(Self { snapshot, splits }) + } + /// The snapshot captured by the scan; accessing it never resolves latest. pub fn snapshot(&self) -> Option<&crate::spec::Snapshot> { - self.table.travel_snapshot() + self.snapshot.as_ref() } /// Source metadata for inspection or transport to another search reader. - pub async fn splits(&self) -> crate::Result> { - if let Some(plan) = &self.primary_key_plan { - return Ok(plan - .splits - .iter() - .map(|split| crate::table::SearchScanSplit::PrimaryKey { - data_split: split.data_split.clone(), - payloads: split.current_payloads.clone(), - uncovered_data_files: split.uncovered_data_files.clone(), - row_ranges_by_file: Default::default(), - }) - .collect()); - } - super::super::search_scan_split::global_search_splits( - &self.table, - &self.entries, - &self.context.column, - self.context.filter.as_ref(), - &Default::default(), - true, - ) - .await + pub fn splits(&self) -> &[crate::table::SearchScanSplit] { + &self.splits } pub fn snapshot_id(&self) -> Option { - self.snapshot_id + self.snapshot.as_ref().map(crate::spec::Snapshot::id) } } @@ -147,87 +174,101 @@ impl FullTextScan { pub async fn scan(&self) -> crate::Result { CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; - let mut plan = FullTextScanPlan { - context: self.context.clone(), - table: self.table.clone(), - entries: Vec::new(), - next_row_id: None, - snapshot_id: None, - partition_filter: None, - row_filter: None, - primary_key_plan: None, - }; - let mut partition_pruner = None; - if let Some(filter) = &self.context.filter { - let (partition, data) = split_partition_and_data_predicates( - filter.clone(), - self.table.schema().fields(), - self.table.schema().partition_keys(), - ); - if let Some(partition) = partition { - let fields = self.table.schema().partition_fields(); - partition_pruner = - Some(PartitionFilter::from_predicate(partition.clone(), &fields)); - let mapping = self - .table - .schema() - .partition_keys() - .iter() - .map(|key| { - self.table - .schema() - .fields() - .iter() - .position(|field| field.name() == key) - }) - .collect::>(); - plan.partition_filter = - Some(partition.remap_field_index(&mapping).ok_or_else(|| { - crate::Error::DataInvalid { - message: - "Cannot bind full-text partition predicate to the table schema" - .into(), - source: None, - } - })?); - } - plan.row_filter = (!data.is_empty()).then(|| Predicate::and(data)); - } - if plan.row_filter.is_some() && !self.table.schema().core_options().data_evolution_enabled() - { + let (partition_filter, row_filter) = + bind_filters(&self.table, self.context.filter.as_ref())?; + if row_filter.is_some() && !self.table.schema().core_options().data_evolution_enabled() { return Err(crate::Error::Unsupported { message: "Full-text row filters require a data-evolution table".into(), }); } let Some(snapshot) = crate::table::time_travel::resolve_snapshot(&self.table).await? else { - return Ok(plan); + return FullTextScanPlan::new(None, vec![]); }; - plan.snapshot_id = Some(snapshot.id()); - plan.next_row_id = snapshot.next_row_id(); - // Pin the data without changing the schema already bound by the query. - // A schema-only rename does not create a new data snapshot. - plan.table = self.table.copy_with_pinned_snapshot(Some(&snapshot)); - if let Some(field_id) = self.context.primary_key_field_id { - plan.primary_key_plan = Some( - PrimaryKeyFullTextScan::new(&plan.table, field_id, self.context.filter.clone()) - .plan() - .await?, - ); - return Ok(plan); - } - if let Some(name) = snapshot.index_manifest() { - let path = self.table.snapshot_manager().manifest_path(name); - plan.entries = IndexManifest::read(self.table.file_io(), &path).await?; - } - if let Some(partition) = partition_pruner { - let mut entries = Vec::new(); - for entry in plan.entries { - if partition.matches_entry(&entry.partition)? { - entries.push(entry); + // Pin only the data: a schema-only rename does not create a snapshot. + let table = self.table.copy_with_pinned_snapshot(Some(&snapshot)); + let splits = if let Some(field_id) = self.context.primary_key_field_id { + PrimaryKeyFullTextScan::new(&table, field_id, self.context.filter.clone()) + .plan() + .await? + .splits + .into_iter() + .map(|split| crate::table::SearchScanSplit::PrimaryKey { + data_split: split.data_split, + payloads: split.current_payloads, + uncovered_data_files: split.uncovered_data_files, + row_ranges_by_file: Default::default(), + }) + .collect() + } else { + let mut entries = if let Some(name) = snapshot.index_manifest() { + let path = table.snapshot_manager().manifest_path(name); + IndexManifest::read(table.file_io(), &path).await? + } else { + vec![] + }; + if partition_filter.is_some() { + let (partition, _) = split_partition_and_data_predicates( + self.context.filter.clone().unwrap(), + table.schema().fields(), + table.schema().partition_keys(), + ); + let fields = table.schema().partition_fields(); + let pruner = PartitionFilter::from_predicate(partition.unwrap(), &fields); + let mut selected = Vec::new(); + for entry in entries { + if pruner.matches_entry(&entry.partition)? { + selected.push(entry); + } } + entries = selected; } - plan.entries = entries; - } - Ok(plan) + super::super::search_scan_split::global_search_splits( + &table, + &entries, + &self.context.column, + self.context.filter.as_ref(), + &Default::default(), + true, + ) + .await? + }; + FullTextScanPlan::new(Some(snapshot), splits) } } + +pub(super) fn bind_filters( + table: &Table, + filter: Option<&Predicate>, +) -> crate::Result<(Option, Option)> { + let Some(filter) = filter else { + return Ok((None, None)); + }; + let (partition, data) = split_partition_and_data_predicates( + filter.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ); + let partition = partition + .map(|partition| { + let mapping = table + .schema() + .partition_keys() + .iter() + .map(|key| { + table + .schema() + .fields() + .iter() + .position(|field| field.name() == key) + }) + .collect::>(); + partition + .remap_field_index(&mapping) + .ok_or_else(|| crate::Error::DataInvalid { + message: "Cannot bind full-text partition predicate to the table schema".into(), + source: None, + }) + }) + .transpose()?; + Ok((partition, (!data.is_empty()).then(|| Predicate::and(data)))) +} diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index ce57a9fb8..9a1c30610 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -321,7 +321,8 @@ impl<'a> HybridSearchBuilder<'a> { /// Schedule each route's shard work while preserving one shared snapshot /// and the configured core fusion ranker. - pub async fn execute_with_executor( + #[cfg(any(feature = "internal-api", test))] + pub(crate) async fn execute_with_executor( &self, executor: &dyn super::SearchTaskExecutor, ) -> crate::Result { diff --git a/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs index d3ad4acf4..1d542a95d 100644 --- a/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs +++ b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs @@ -195,6 +195,25 @@ async fn both_routes_filter_before_top_k_and_accumulate_predicates() { .with_filter(pb.equal("id", Datum::Int(6)).unwrap()) .with_partition_filter(pb.equal("pt", Datum::Int(1)).unwrap()) .unwrap(); + let pinned = builder.resolve_pinned_route_table().await.unwrap().unwrap(); + let vector = builder + .vector_route_builder(&pinned, &builder.routes[0]) + .execute() + .await + .unwrap(); + assert_eq!( + vector.row_ids().unwrap().row_ids, + vec![6], + "the vector route must filter before its top-k" + ); + let text = execute_full_text_route(&pinned, &builder.routes[1], builder.filter.as_ref(), None) + .await + .unwrap(); + assert_eq!( + text.row_ids, + vec![6], + "the text route must filter before its top-k" + ); let result = builder.execute_scored().await.unwrap(); assert_eq!(result.row_ids().unwrap().row_ids, vec![6]); assert!((result.row_ids().unwrap().scores[0] - 2.0 / 61.0).abs() < 1e-6); @@ -494,3 +513,57 @@ async fn scheduled_hybrid_shares_snapshot_and_preserves_all_core_rankers() { assert!(kinds.iter().any(|kind| kind == "fulltext-index")); } } + +#[tokio::test] +async fn explicit_raw_source_range_with_partition_filter_keeps_global_row_ids() { + use futures::TryStreamExt; + let table = table(); + append(&table, 0).await; + append(&table, 4).await; + let snapshot = table.resolve_read_snapshot().await.unwrap().unwrap(); + let plan = crate::table::VectorScanPlan::new( + Some(snapshot), + vec![crate::table::SearchScanSplit::Raw { + row_ranges: vec![RowRange::new(4, 7)], + scalar_index_files: vec![], + index_type: Some("ivf-flat".into()), + }], + ) + .unwrap(); + let predicates = PredicateBuilder::new(table.schema().fields()); + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![0.0, 0.0]) + .with_limit(1) + .with_filter(Predicate::and(vec![ + predicates.equal("pt", Datum::Int(1)).unwrap(), + predicates.equal("id", Datum::Int(6)).unwrap(), + ])); + let result = builder.new_read().unwrap().read(plan).await.unwrap(); + assert_eq!(result.row_ids().unwrap().row_ids, vec![6]); + let batches: Vec = result + .new_read_builder() + .with_projection(&["id"]) + .read() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let ids = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(ids, vec![6]); +} diff --git a/crates/paimon/src/table/internal.rs b/crates/paimon/src/table/internal.rs new file mode 100644 index 000000000..06b6a80c1 --- /dev/null +++ b/crates/paimon/src/table/internal.rs @@ -0,0 +1,67 @@ +// 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. + +//! Unstable scheduling bridge for the repository's language bindings. +//! +//! Source plans are the public input to search readers. Candidate and scoring +//! requests below are internal execution stages: Core owns their ordering and +//! global result coordination. This module is available only with `internal-api`. + +pub use super::search_task::{SearchTask, SearchTaskExecutor, SearchTaskResult}; +use super::{ + BatchVectorRead, GlobalIndexResult, HybridSearchBuilder, Table, VectorRead, VectorScanPlan, +}; + +pub async fn execute_search_task( + table: &Table, + task: &[u8], + context: &[u8], +) -> crate::Result> { + table.execute_search_task(task, context).await +} + +pub async fn read_vector_with_executor( + reader: &VectorRead, + plan: VectorScanPlan, + executor: &dyn SearchTaskExecutor, +) -> crate::Result { + reader.read_with_executor(plan, executor).await +} + +pub async fn read_batch_vector_with_executor( + reader: &BatchVectorRead, + plan: VectorScanPlan, + executor: &dyn SearchTaskExecutor, +) -> crate::Result> { + reader.read_with_executor(plan, executor).await +} + +#[cfg(feature = "fulltext")] +pub async fn read_full_text_with_executor( + reader: &super::FullTextRead, + plan: super::FullTextScanPlan, + executor: &dyn SearchTaskExecutor, +) -> crate::Result { + reader.read_with_executor(plan, executor).await +} + +pub async fn execute_hybrid_with_executor( + builder: &HybridSearchBuilder<'_>, + executor: &dyn SearchTaskExecutor, +) -> crate::Result { + builder.execute_with_executor(executor).await +} diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 0862146c7..7cbce5ba2 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -151,7 +151,14 @@ mod scan_stats; pub(crate) mod schema_manager; mod search_scan_split; mod search_task; -pub use search_task::{SearchTask, SearchTaskExecutor, SearchTaskResult}; +#[cfg(any(feature = "fulltext", test))] +pub(crate) use search_task::SearchTask; +pub(crate) use search_task::SearchTaskExecutor; +#[cfg(test)] +pub(crate) use search_task::SearchTaskResult; +#[cfg(feature = "internal-api")] +#[doc(hidden)] +pub mod internal; mod shard_table_updater; pub(crate) mod snapshot_commit; mod snapshot_manager; diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 8d3d033a8..7d40fa320 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -615,6 +615,200 @@ fn plan_from_bucket_splits( }) } +/// Restore authoritative engine/transport sources through the same ANN source +/// validation and bucket assembler used by the manifest planner. +pub(super) fn plan_from_search_splits( + table: &Table, + column: &str, + filter: Option<&Predicate>, + snapshot: Option, + sources: &[crate::table::SearchScanSplit], +) -> crate::Result { + let snapshot_id = snapshot.as_ref().map_or(0, Snapshot::id); + let core = table.schema().core_options(); + let index_type = core.primary_key_vector_index_type(column)?; + let field_id = crate::table::find_field_id_by_name(table.schema().fields(), column) + .ok_or_else(|| data_invalid("PK-vector column does not exist"))?; + let partition = filter + .and_then(|filter| { + split_partition_and_data_predicates( + filter.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ) + .0 + }) + .map(|filter| PartitionFilter::from_predicate(filter, &table.schema().partition_fields())); + let mut data = Vec::new(); + let mut entries = Vec::new(); + let mut payloads = HashMap::new(); + let mut restrictions = HashMap::new(); + let mut uncovered = HashMap::new(); + let mut seen = HashSet::new(); + for source in sources { + let crate::table::SearchScanSplit::PrimaryKey { + data_split, + payloads: files, + uncovered_data_files, + row_ranges_by_file, + } = source + else { + return Err(data_invalid( + "Primary-key vector reader cannot consume global row-ID sources", + )); + }; + if data_split.snapshot_id() != 0 && data_split.snapshot_id() != snapshot_id { + return Err(data_invalid( + "Primary-key source split belongs to another snapshot", + )); + } + if data_split.query_auth_required() { + return Err(crate::table::query_auth::unsupported( + "transported primary-key vector split carries no authorization", + )); + } + if data_split.row_ranges().is_some() { + return Err(data_invalid( + "Primary-key bucket sources must use row_ranges_by_file", + )); + } + let key = ( + data_split.partition().to_serialized_bytes(), + data_split.bucket(), + ); + if !seen.insert(key.clone()) { + return Err(data_invalid("Primary-key search sources repeat a bucket")); + } + if let Some(filter) = &partition { + if !filter.matches_entry(&key.0)? { + continue; + } + } + let names = data_split + .data_files() + .iter() + .map(|file| file.file_name.as_str()) + .collect::>(); + if uncovered_data_files + .iter() + .any(|name| !names.contains(name.as_str())) + || row_ranges_by_file + .keys() + .any(|name| !names.contains(name.as_str())) + { + return Err(data_invalid( + "Primary-key source selects an unknown data file", + )); + } + for (name, ranges) in row_ranges_by_file { + let rows = data_split + .data_files() + .iter() + .find(|file| file.file_name == *name) + .unwrap() + .row_count; + if ranges + .iter() + .any(|range| range.from() < 0 || range.to() < range.from() || range.to() >= rows) + { + return Err(data_invalid( + "Primary-key source row range exceeds its data file", + )); + } + } + for file in files { + let meta = file + .global_index_meta + .as_ref() + .ok_or_else(|| data_invalid("PK payload lacks global metadata"))?; + if file.index_type != index_type || meta.index_field_id != field_id { + return Err(data_invalid( + "PK payload does not match the reader's vector column or index type", + )); + } + entries.push(( + data_split.partition().clone(), + data_split.bucket(), + UnresolvedAnnSegment { + source_meta: PrimaryKeyIndexSourceMeta::from_global_index_meta(meta)?, + file_name: file.file_name.clone(), + external_path: file.external_path.clone(), + file_size: u64::try_from(file.file_size) + .map_err(|_| data_invalid("Index file size must not be negative"))?, + index_meta: meta.index_meta.clone().unwrap_or_default(), + }, + )); + } + let data_split = if data_split.snapshot_id() == 0 { + let mut builder = DataSplitBuilder::new() + .with_snapshot(snapshot_id) + .with_partition(data_split.partition().clone()) + .with_bucket(data_split.bucket()) + .with_bucket_path(data_split.bucket_path().to_string()) + .with_total_buckets(data_split.total_buckets()) + .with_data_files(data_split.data_files().to_vec()) + .with_raw_convertible(data_split.raw_convertible()) + .with_streaming(data_split.is_streaming()); + if let Some(dvs) = data_split.data_deletion_files() { + builder = builder.with_data_deletion_files(dvs.to_vec()); + } + builder.build()? + } else { + data_split.clone() + }; + data.push(data_split); + payloads.insert(key.clone(), files.clone()); + restrictions.insert( + key.clone(), + row_ranges_by_file + .iter() + .map(|(name, ranges)| (name.clone(), merge_row_ranges(ranges.clone()))) + .collect(), + ); + uncovered.insert( + key, + uncovered_data_files.iter().cloned().collect::>(), + ); + } + let mut splits = plan_from_inputs( + snapshot_id, + data, + entries, + table.location(), + core.index_file_in_data_file_dir(), + )?; + let physical_row_ranges_by_split = splits + .iter_mut() + .map(|split| { + let key = ( + split.data_split.partition().to_serialized_bytes(), + split.data_split.bucket(), + ); + // A source subset is authoritative even if it omitted an eligible raw + // file. The ANN source policy may validate segments, but never adds files + // to the caller's exact fallback set. + let names = &uncovered[&key]; + let covered = split + .ann_segments + .iter() + .flat_map(|segment| segment.source_meta.source_files()) + .map(|file| file.file_name()) + .collect::>(); + split.active_files.retain(|file| { + names.contains(&file.file_name) || covered.contains(file.file_name.as_str()) + }); + restrictions.remove(&key).unwrap_or_default() + }) + .collect(); + Ok(PkVectorScanPlan { + snapshot, + payloads, + snapshot_id, + splits, + physical_row_ranges_by_split: Some(physical_row_ranges_by_split), + }) +} + /// Pure planning core, drivable without a live snapshot: group ANN payloads and /// data splits by `(partition, bucket)`, then assemble one search split per /// bucket that has data. Index-only buckets are dropped, not errored. diff --git a/crates/paimon/src/table/search_scan_split.rs b/crates/paimon/src/table/search_scan_split.rs index 692ae15f8..42e8ffd04 100644 --- a/crates/paimon/src/table/search_scan_split.rs +++ b/crates/paimon/src/table/search_scan_split.rs @@ -26,11 +26,12 @@ use super::{find_field_id_by_name, merge_row_ranges, DataSplit, RowRange, Table} use crate::spec::{ CoreOptions, FileKind, GlobalIndexSearchMode, IndexFileMeta, IndexManifestEntry, Predicate, }; +use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, HashMap, HashSet}; /// Search source metadata, without live readers or filesystem handles. /// The global row-ID and primary-key physical address spaces remain distinct. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub enum SearchScanSplit { Index { column: String, @@ -90,7 +91,6 @@ pub(super) async fn global_search_splits( if full_text && selected.is_empty() { return Ok(vec![]); } - let mut filter_fields = Vec::new(); let row_filter = filter.and_then(|filter| { let (_, data) = split_partition_and_data_predicates( filter.clone(), @@ -99,26 +99,13 @@ pub(super) async fn global_search_splits( ); (!data.is_empty()).then(|| Predicate::and(data)) }); - if let Some(filter) = &row_filter { - crate::arrow::residual::collect_predicate_leaf_refs(filter, &mut filter_fields); - } - let filter_ids = filter_fields - .into_iter() - .filter_map(|(name, _)| find_field_id_by_name(table.schema().fields(), name)) - .collect::>(); let scalar_files = entries .iter() .filter(|entry| entry.kind == FileKind::Add) .map(|entry| &entry.index_file) .filter(|file| { normalize_queryable_global_index_type(&file.index_type).is_some() - && file.global_index_meta.as_ref().is_some_and(|meta| { - filter_ids.contains(&meta.index_field_id) - || meta - .extra_field_ids - .as_ref() - .is_some_and(|ids| ids.iter().any(|id| filter_ids.contains(id))) - }) + && file.global_index_meta.is_some() }) .cloned() .collect::>(); @@ -222,3 +209,190 @@ pub(super) async fn global_search_splits( } Ok(result) } + +fn invalid_source(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +/// Restore only the index payloads explicitly named by a source plan. The +/// partition/bucket fields are irrelevant to global payload addressing; scalar +/// and vector readers use the persisted global metadata instead. +pub(super) fn global_entries_from_splits( + splits: &[SearchScanSplit], +) -> crate::Result> { + let mut entries = Vec::new(); + let mut scalar_files = Vec::::new(); + for split in splits { + let scalar = match split { + SearchScanSplit::Index { + column, + row_range, + index_files, + scalar_index_files, + } => { + if column.is_empty() || row_range.from() < 0 || row_range.to() < row_range.from() { + return Err(invalid_source("Invalid global search index split")); + } + for file in index_files { + let meta = file.global_index_meta.as_ref().ok_or_else(|| { + invalid_source(format!( + "Search index '{}' has no global metadata", + file.file_name + )) + })?; + if meta.row_range_start != row_range.from() + || meta.row_range_end != row_range.to() + { + return Err(invalid_source( + "Search index split range disagrees with its payload", + )); + } + entries.push(IndexManifestEntry { + version: 1, + kind: FileKind::Add, + partition: crate::spec::BinaryRow::new(0).to_serialized_bytes(), + bucket: -1, + index_file: file.clone(), + }); + } + scalar_index_files + } + SearchScanSplit::Raw { + row_ranges, + scalar_index_files, + .. + } => { + if row_ranges + .iter() + .any(|range| range.from() < 0 || range.to() < range.from()) + { + return Err(invalid_source("Invalid global search raw row range")); + } + scalar_index_files + } + SearchScanSplit::PrimaryKey { .. } => { + return Err(invalid_source( + "Global search cannot consume primary-key physical source splits", + )) + } + }; + for file in scalar { + let meta = file.global_index_meta.as_ref().ok_or_else(|| { + invalid_source(format!( + "Scalar search index '{}' has no global metadata", + file.file_name + )) + })?; + if meta.row_range_start < 0 || meta.row_range_end < meta.row_range_start { + return Err(invalid_source("Invalid scalar search index range")); + } + if !scalar_files.contains(file) { + scalar_files.push(file.clone()); + } + } + } + entries.extend( + scalar_files + .into_iter() + .map(|index_file| IndexManifestEntry { + version: 1, + kind: FileKind::Add, + partition: crate::spec::BinaryRow::new(0).to_serialized_bytes(), + bucket: -1, + index_file, + }), + ); + Ok(entries) +} + +/// Read deletion vectors from the data scan of the supplied pinned snapshot. +/// A transported source plan need not copy bucket-local DV index entries into +/// every search split. Data-source planning resolves those vectors without +/// selecting a newer snapshot or extending the plan's search sources. +pub(super) async fn deleted_row_ranges_for_search_sources( + table: &Table, +) -> crate::Result> { + if !table.schema().core_options().deletion_vectors_enabled() { + return Ok(vec![]); + } + let plan = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan() + .await?; + let mut ranges = Vec::new(); + for split in plan.splits() { + let Some(dvs) = split.data_deletion_files() else { + continue; + }; + for (file, deletion) in split.data_files().iter().zip(dvs) { + let Some(deletion) = deletion else { + continue; + }; + let first = file.first_row_id.ok_or_else(|| { + invalid_source(format!( + "Deletion vector anchor '{}' has no first row id", + file.file_name + )) + })?; + for deleted in + crate::deletion_vector::DeletionVectorFactory::read(table.file_io(), deletion) + .await? + .iter() + { + let deleted = i64::try_from(deleted) + .map_err(|_| invalid_source("Deleted position exceeds i64::MAX"))?; + let row_id = first + .checked_add(deleted) + .ok_or_else(|| invalid_source("Deleted row id exceeds i64::MAX"))?; + ranges.push(RowRange::new(row_id, row_id)); + } + } + } + Ok(merge_row_ranges(ranges)) +} + +/// Resolve partition membership in the global row-ID address space from the +/// pinned data-file metadata. Partition predicates are constant per split; +/// reading Arrow rows would add work and couple this to read projections. +pub(super) async fn partition_row_ids_for_search_sources( + table: &Table, + partition: &Predicate, + source_ranges: Vec, +) -> crate::Result { + let source_ranges = merge_row_ranges(source_ranges); + let mut read = table.new_read_builder(); + read.with_filter(partition.clone()) + .with_row_ranges(source_ranges.clone()); + let plan = read.new_scan().plan().await?; + let mut file_ranges = Vec::new(); + for split in plan.splits() { + for file in split.data_files() { + if file.row_count <= 0 { + continue; + } + let first = file.first_row_id.ok_or_else(|| { + invalid_source(format!( + "Partition search source '{}' has no first row id", + file.file_name + )) + })?; + let last = first.checked_add(file.row_count - 1).ok_or_else(|| { + invalid_source("Partition search source row range exceeds i64::MAX") + })?; + if first < 0 { + return Err(invalid_source("Partition search source row ID is negative")); + } + file_ranges.push(RowRange::new(first, last)); + } + } + let mut selected = roaring::RoaringTreemap::new(); + for range in intersect_sorted_ranges(&merge_row_ranges(file_ranges), &source_ranges) { + selected.insert_range(range.from() as u64..=range.to() as u64); + } + Ok(selected) +} diff --git a/crates/paimon/src/table/search_task.rs b/crates/paimon/src/table/search_task.rs index fdd9866fd..eab59d920 100644 --- a/crates/paimon/src/table/search_task.rs +++ b/crates/paimon/src/table/search_task.rs @@ -18,7 +18,9 @@ //! Opaque, snapshot-scoped search work for local or remote schedulers. use super::Table; -use crate::spec::{CoreOptions, Snapshot, TableSchema}; +#[cfg(any(feature = "internal-api", test))] +use crate::spec::CoreOptions; +use crate::spec::{Snapshot, TableSchema}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; @@ -33,10 +35,12 @@ pub(super) fn invalid(message: impl Into) -> crate::Error { } } +#[cfg(any(feature = "internal-api", test))] pub(super) fn encode(value: &T) -> crate::Result> { serde_json::to_vec(value).map_err(|error| invalid(format!("Invalid search task: {error}"))) } +#[cfg(any(feature = "internal-api", test))] pub(super) fn decode(value: &[u8]) -> crate::Result { serde_json::from_slice(value).map_err(|error| invalid(format!("Invalid search task: {error}"))) } @@ -68,7 +72,10 @@ impl SearchTask { payload, }) } +} +#[cfg(any(feature = "internal-api", test))] +impl SearchTask { pub fn serialize(&self) -> crate::Result> { encode(self) } @@ -219,14 +226,6 @@ impl SearchTask { Ok(task) } - pub fn deserialize(bytes: &[u8]) -> crate::Result { - let task: Self = decode(bytes)?; - if task.version != VERSION { - return Err(invalid("Unsupported search task version")); - } - Ok(task) - } - pub async fn execute(&self, table: &Table) -> crate::Result { // The worker's runtime restrictions take precedence over transported options. CoreOptions::new(table.schema().options()).ensure_read_authorized()?; @@ -268,6 +267,7 @@ pub struct SearchTaskResult { payload: Value, } +#[cfg(any(feature = "internal-api", test))] impl SearchTaskResult { pub fn serialize(&self) -> crate::Result> { encode(self) @@ -317,12 +317,14 @@ pub(super) async fn execute_tasks( .collect() } +#[cfg(any(feature = "internal-api", test))] #[derive(Serialize, Deserialize)] struct SharedTaskContext { version: u8, id: String, values: Vec, } +#[cfg(any(feature = "internal-api", test))] #[derive(Serialize, Deserialize)] struct SharedTaskFrame { version: u8, @@ -331,9 +333,14 @@ struct SharedTaskFrame { references: Vec<(String, usize)>, } +#[cfg(any(feature = "internal-api", test))] impl Table { - /// Execute one opaque task and its shared inputs using this worker's credentials. - pub async fn execute_search_task(&self, task: &[u8], context: &[u8]) -> crate::Result> { + /// Internal worker bridge; ordinary callers execute a constructed scan plan. + pub(crate) async fn execute_search_task( + &self, + task: &[u8], + context: &[u8], + ) -> crate::Result> { SearchTask::deserialize_with_context(task, context)? .execute(self) .await? diff --git a/crates/paimon/src/table/vector_read.rs b/crates/paimon/src/table/vector_read.rs index 294e4899b..ebc7a64e7 100644 --- a/crates/paimon/src/table/vector_read.rs +++ b/crates/paimon/src/table/vector_read.rs @@ -21,14 +21,15 @@ use crate::spec::{CoreOptions, Predicate}; use crate::table::de_vector_read::DeVectorRead; use crate::table::pk_vector_read::PkVectorRead; use crate::table::pk_vector_search_params::PkVectorSearchParams; -use crate::table::vector_scan::{PlanContext, VectorScanPlan, VectorScanWork}; +use crate::table::vector_scan::{PlanContext, VectorScanPlan}; use crate::table::vector_search_common::{take_only_result, targets_primary_key_column}; use crate::table::{SearchTaskExecutor, Table}; use crate::vector_search::GlobalIndexResult; use std::collections::HashMap; use std::future::Future; -/// Execute a resolved plan without resolving another snapshot or manifest. +/// Execute a source plan without resolving another snapshot or discovering +/// additional search sources. /// /// Both implementations return snapshot-scoped search results. The associated /// plan type prevents passing a DE plan to a PK reader or vice versa. @@ -50,7 +51,7 @@ impl VectorRead { pub async fn read(&self, plan: VectorScanPlan) -> crate::Result { take_only_result(self.batch.read(plan).await?, "vector search") } - pub async fn read_with_executor( + pub(crate) async fn read_with_executor( &self, plan: VectorScanPlan, executor: &dyn SearchTaskExecutor, @@ -66,6 +67,7 @@ impl VectorRead { /// Owns its configuration so it can outlive the builder that created it. pub struct BatchVectorRead { context: PlanContext, + table: Table, reader: VectorReadKind, } @@ -95,13 +97,17 @@ impl BatchVectorRead { } else { VectorReadKind::DataEvolution(DeVectorRead::new(column, queries, limit, options)?) }; - Ok(Self { context, reader }) + Ok(Self { + context, + table: table.clone(), + reader, + }) } pub async fn read(&self, plan: VectorScanPlan) -> crate::Result> { self.read_impl(plan, None).await } - pub async fn read_with_executor( + pub(crate) async fn read_with_executor( &self, plan: VectorScanPlan, executor: &dyn SearchTaskExecutor, @@ -113,29 +119,76 @@ impl BatchVectorRead { plan: VectorScanPlan, executor: Option<&dyn SearchTaskExecutor>, ) -> crate::Result> { - if self.context != plan.context { - return Err(crate::Error::DataInvalid { - message: "vector plan and reader must use the same table, column and pre-filter" - .to_string(), - source: None, - }); + CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + if let Some(snapshot) = &self.context.prepared_snapshot { + if snapshot.as_ref() != plan.snapshot() { + return Err(crate::Error::DataInvalid { + message: + "Prepared vector filter and source plan must use the same captured snapshot" + .into(), + source: None, + }); + } } - match (&self.reader, plan.work) { - (VectorReadKind::DataEvolution(reader), VectorScanWork::DataEvolution(plan)) => { - reader.read_with_executor(*plan, executor).await + let table = self.table.copy_with_pinned_snapshot(plan.snapshot()); + match &self.reader { + VectorReadKind::DataEvolution(reader) => { + for source in plan.splits() { + if let crate::table::SearchScanSplit::Index { index_files, .. } = source { + let field_id = crate::table::find_field_id_by_name( + table.schema().fields(), + &self.context.column, + ); + if index_files.iter().any(|file| { + super::vector_search_common::VectorIndexBackend::from_index_type( + &file.index_type, + ) + .is_none() + || file + .global_index_meta + .as_ref() + .map(|meta| meta.index_field_id) + != field_id + }) { + return Err(crate::Error::DataInvalid { + message: + "Vector source index does not match the reader's vector column" + .into(), + source: None, + }); + } + } + } + let index_entries = + super::search_scan_split::global_entries_from_splits(plan.splits())?; + let work = super::de_vector_scan::DeVectorScanPlan { + next_row_id: plan.snapshot().and_then(|snapshot| snapshot.next_row_id()), + skip_search: plan.snapshot().is_none() || plan.splits().is_empty(), + table, + index_entries, + filter: self.context.filter.clone(), + include_row_ids: self.context.include_row_ids.clone(), + timing: None, + source_splits: Some(plan.splits().to_vec()), + }; + reader.read_with_executor(work, executor).await } - (VectorReadKind::PrimaryKey(reader), VectorScanWork::PrimaryKey(plan)) => { + VectorReadKind::PrimaryKey(reader) => { if executor.is_some() { return Err(crate::Error::Unsupported { message: "Distributed primary-key vector search is not supported".into(), }); } - reader.read(*plan).await + reader + .read(super::pk_vector_scan::plan_from_search_splits( + &table, + &self.context.column, + self.context.filter.as_ref(), + plan.snapshot().cloned(), + plan.splits(), + )?) + .await } - _ => Err(crate::Error::DataInvalid { - message: "vector plan and reader use different DE/PK address spaces".to_string(), - source: None, - }), } } } diff --git a/crates/paimon/src/table/vector_scan.rs b/crates/paimon/src/table/vector_scan.rs index f8411fb0e..77ef44015 100644 --- a/crates/paimon/src/table/vector_scan.rs +++ b/crates/paimon/src/table/vector_scan.rs @@ -23,7 +23,7 @@ use std::sync::Arc; use roaring::RoaringTreemap; use crate::spec::{CoreOptions, Predicate, Snapshot}; -use crate::table::de_vector_scan::{DeVectorScan, DeVectorScanPlan, PreparedVectorSearchFilter}; +use crate::table::de_vector_scan::{DeVectorScan, PreparedVectorSearchFilter}; use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan}; use crate::table::vector_search_common::{ targets_primary_key_column, validate_global_vector_index_types, @@ -42,146 +42,156 @@ pub(super) trait Scan: Sync { fn plan(&self) -> impl Future> + Send; } -/// Snapshot and search work shared by the single-query and batch readers. -/// Plans own their source context and remain usable after the scan is dropped. -#[derive(Clone)] +/// Snapshot and source work shared by single-query and batch readers. +/// +/// A plan is ordinary transportable data. Query vectors, filters, options, +/// authorization and filesystem handles belong to the reader that consumes it. +#[derive(Clone, Debug, serde::Serialize)] pub struct VectorScanPlan { - pub(super) context: PlanContext, - pub(super) work: VectorScanWork, - table: Table, -} - -#[derive(Clone)] -pub(super) enum VectorScanWork { - DataEvolution(Box), - PrimaryKey(Box), + snapshot: Option, + splits: Vec, } impl VectorScanPlan { - /// The snapshot captured by planning, including retained tag snapshots. - pub fn snapshot(&self) -> Option<&Snapshot> { - match &self.work { - VectorScanWork::DataEvolution(plan) => plan.table.travel_snapshot(), - VectorScanWork::PrimaryKey(plan) => plan.snapshot.as_ref(), + /// Construct a plan from authoritative source work. No table, snapshot + /// manager or index manifest is consulted. `None` denotes an empty table. + pub fn new(snapshot: Option, splits: Vec) -> crate::Result { + if snapshot.is_none() && !splits.is_empty() { + return Err(crate::Error::DataInvalid { + message: "A vector source plan without a snapshot must be empty".into(), + source: None, + }); } - } - - /// Inspect source work using the scan's effective search options. This does - /// not resolve a newer snapshot or reread its index manifest. Query readers - /// may still use different query vectors, limits and search options. - pub async fn splits( - &self, - options: &HashMap, - ) -> crate::Result> { - match &self.work { - VectorScanWork::DataEvolution(plan) => { - if plan.skip_search { - return Ok(vec![]); - } - super::search_scan_split::global_search_splits( - &plan.table, - &plan.index_entries, - &self.context.column, - self.context.filter.as_ref(), - options, - false, - ) - .await - } - VectorScanWork::PrimaryKey(plan) => { - if plan.snapshot.is_none() - && plan - .splits - .iter() - .any(|split| !split.ann_segments.is_empty()) - { - return Err(crate::Error::Unsupported { - message: - "Source metadata export requires a manifest-planned primary-key search" - .into(), - }); + let primary_key = splits + .iter() + .any(|source| matches!(source, SearchScanSplit::PrimaryKey { .. })); + if primary_key + && splits + .iter() + .any(|source| !matches!(source, SearchScanSplit::PrimaryKey { .. })) + { + return Err(crate::Error::DataInvalid { + message: "A vector source plan cannot mix global row IDs and primary-key positions" + .into(), + source: None, + }); + } + if !primary_key { + super::search_scan_split::global_entries_from_splits(&splits)?; + } + if let Some(snapshot) = &snapshot { + for split in &splits { + if let SearchScanSplit::PrimaryKey { data_split, .. } = split { + if data_split.snapshot_id() != 0 && data_split.snapshot_id() != snapshot.id() { + return Err(crate::Error::DataInvalid { + message: "Primary-key vector split belongs to another snapshot".into(), + source: None, + }); + } } - let selected = super::pk_vector_read::filtered_row_ranges( - &self.table, - self.context.filter.as_ref(), - plan, - ) - .await?; - plan.splits - .iter() - .enumerate() - .map(|(index, split)| { - let payloads = plan - .payloads - .get(&( - split.data_split.partition().to_serialized_bytes(), - split.data_split.bucket(), - )) - .into_iter() - .flatten() - .filter(|payload| { - split.ann_segments.iter().any(|segment| { - let location = - super::index_file_path::IndexFileLocation::BucketLocal { - table_path: &self.context.location, - bucket_path: split.data_split.bucket_path(), - index_file_in_data_file_dir: false, - }; - segment.path - == location.resolve( - &payload.file_name, - payload.external_path.as_deref(), - ) - || segment - .path - .ends_with(&format!("/{}", payload.file_name)) - }) - }) - .cloned() - .collect::>(); - let covered = split - .ann_segments - .iter() - .flat_map(|segment| segment.source_meta.source_files()) - .map(|file| file.file_name()) - .collect::>(); - Ok(SearchScanSplit::PrimaryKey { - data_split: split.data_split.clone(), - payloads, - uncovered_data_files: split - .active_files - .iter() - .filter(|file| !covered.contains(file.file_name.as_str())) - .map(|file| file.file_name.clone()) - .collect(), - row_ranges_by_file: selected - .as_ref() - .map(|ranges| ranges[index].clone()) - .unwrap_or_default(), - }) - }) - .collect() } } + Ok(Self { snapshot, splits }) + } + + /// The complete snapshot captured by planning, including retained tags. + pub fn snapshot(&self) -> Option<&Snapshot> { + self.snapshot.as_ref() + } + + /// The exact sources selected by the planner. An empty or restricted list + /// never causes the reader to discover additional index or raw sources. + pub fn splits(&self) -> &[SearchScanSplit] { + &self.splits } pub fn snapshot_id(&self) -> Option { - match &self.work { - VectorScanWork::DataEvolution(plan) => plan.table.travel_snapshot().map(|s| s.id()), - VectorScanWork::PrimaryKey(plan) => (plan.snapshot_id != 0).then_some(plan.snapshot_id), + self.snapshot.as_ref().map(Snapshot::id) + } +} + +impl<'de> serde::Deserialize<'de> for VectorScanPlan { + fn deserialize>( + deserializer: D, + ) -> std::result::Result { + #[derive(serde::Deserialize)] + struct Sources { + snapshot: Option, + splits: Vec, } + let Sources { snapshot, splits } = + ::deserialize(deserializer)?; + Self::new(snapshot, splits).map_err(serde::de::Error::custom) } } -/// A plan must use the reader's table, column and pre-filter. Query vectors, -/// limits and index search options may differ when a plan is reused. -#[derive(Clone, PartialEq)] +async fn primary_key_sources( + table: &Table, + context: &PlanContext, + plan: &PkVectorScanPlan, +) -> crate::Result> { + let selected = + super::pk_vector_read::filtered_row_ranges(table, context.filter.as_ref(), plan).await?; + plan.splits + .iter() + .enumerate() + .map(|(index, split)| { + let payloads = plan + .payloads + .get(&( + split.data_split.partition().to_serialized_bytes(), + split.data_split.bucket(), + )) + .into_iter() + .flatten() + .filter(|payload| { + split.ann_segments.iter().any(|segment| { + let location = super::index_file_path::IndexFileLocation::BucketLocal { + table_path: table.location(), + bucket_path: split.data_split.bucket_path(), + index_file_in_data_file_dir: table + .schema() + .core_options() + .index_file_in_data_file_dir(), + }; + segment.path + == location + .resolve(&payload.file_name, payload.external_path.as_deref()) + }) + }) + .cloned() + .collect(); + let covered = split + .ann_segments + .iter() + .flat_map(|segment| segment.source_meta.source_files()) + .map(|file| file.file_name()) + .collect::>(); + Ok(SearchScanSplit::PrimaryKey { + data_split: split.data_split.clone(), + payloads, + uncovered_data_files: split + .active_files + .iter() + .filter(|file| !covered.contains(file.file_name.as_str())) + .map(|file| file.file_name.clone()) + .collect(), + row_ranges_by_file: selected + .as_ref() + .map(|ranges| ranges[index].clone()) + .unwrap_or_default(), + }) + }) + .collect() +} + +/// Validated reader/scan configuration. This is never part of a source plan. +#[derive(Clone)] pub(super) struct PlanContext { - location: String, - branch: String, - column: String, - filter: Option, - include_row_ids: Option>, + pub(super) column: String, + pub(super) filter: Option, + pub(super) include_row_ids: Option>, + pub(super) prepared_snapshot: Option>, } impl PlanContext { @@ -220,14 +230,13 @@ impl PlanContext { }); } Ok(Self { - location: table.location().trim_end_matches('/').to_string(), - branch: table.branch().to_string(), column: column.to_string(), filter: filter.cloned(), include_row_ids: prepared .map(PreparedVectorSearchFilter::include_row_ids) .or(include_row_ids) .cloned(), + prepared_snapshot: prepared.map(|filter| filter.table().travel_snapshot().cloned()), }) } } @@ -237,6 +246,7 @@ pub struct VectorScan { table: Table, context: PlanContext, scan: VectorScanKind, + options: HashMap, } enum VectorScanKind { @@ -281,11 +291,17 @@ impl VectorScan { table: table.clone(), context, scan, + options: HashMap::new(), }) } + pub(super) fn with_options(mut self, options: &HashMap) -> Self { + self.options = options.clone(); + self + } + pub async fn plan(&self) -> crate::Result { - let work = match &self.scan { + match &self.scan { VectorScanKind::DataEvolution(scan) => { let plan = scan.plan().await?; if let Some(field_id) = @@ -297,36 +313,92 @@ impl VectorScan { &self.context.column, )?; } - VectorScanWork::DataEvolution(Box::new(plan)) + let snapshot = plan.table.travel_snapshot().cloned(); + let splits = if plan.skip_search { + vec![] + } else { + super::search_scan_split::global_search_splits( + &plan.table, + &plan.index_entries, + &self.context.column, + self.context.filter.as_ref(), + &self.options, + false, + ) + .await? + }; + VectorScanPlan::new(snapshot, splits) } VectorScanKind::PrimaryKey(scan) => { - VectorScanWork::PrimaryKey(Box::new(scan.plan().await?)) + let plan = scan.plan().await?; + let splits = primary_key_sources(&self.table, &self.context, &plan).await?; + VectorScanPlan::new(plan.snapshot, splits) } - }; - Ok(VectorScanPlan { - table: self.table.clone(), - context: self.context.clone(), - work, - }) + } } - /// Adapt already-decoded Java PK bucket work into a common read plan. - /// The supplied snapshot, files and physical row ranges are authoritative; - /// this does not consult the table's snapshot or index manifests. - pub fn plan_from_bucket_splits( + /// Adapt Java bucket work to the same transportable source plan. A positive + /// bucket snapshot ID is loaded directly. Unversioned buckets capture the + /// scan's configured snapshot once here; the returned plan never follows latest. + pub async fn plan_from_bucket_splits( &self, splits: Vec, ) -> crate::Result { let VectorScanKind::PrimaryKey(scan) = &self.scan else { return Err(crate::Error::DataInvalid { - message: "bucket splits require a primary-key vector scan".to_string(), + message: "bucket splits require a primary-key vector scan".into(), source: None, }); }; - Ok(VectorScanPlan { - table: self.table.clone(), - context: self.context.clone(), - work: VectorScanWork::PrimaryKey(Box::new(scan.plan_for_bucket_vector_splits(splits)?)), - }) + // Preserve manifest-form payload metadata beside the canonical bucket + // assembler's validated segments, so the resulting plan is transportable. + let payloads = splits + .iter() + .map(|split| { + let key = ( + split.data_split().partition().to_serialized_bytes(), + split.data_split().bucket(), + ); + let files = split + .payload_files() + .iter() + .map(|payload| crate::spec::IndexFileMeta { + index_type: payload.index_type().to_string(), + file_name: payload.file_name().to_string(), + file_size: payload.file_size(), + row_count: payload.row_count(), + deletion_vectors_ranges: payload.deletion_vectors_ranges().cloned(), + external_path: payload.external_path().map(str::to_string), + global_index_meta: Some(payload.global_index_meta().clone()), + }) + .collect(); + (key, files) + }) + .collect(); + let mut plan = scan.plan_for_bucket_vector_splits(splits)?; + plan.payloads = payloads; + let snapshot = match self + .table + .travel_snapshot() + .filter(|snapshot| snapshot.id() == plan.snapshot_id) + { + Some(snapshot) => snapshot.clone(), + None if plan.snapshot_id == 0 => { + self.table.resolve_read_snapshot().await?.ok_or_else(|| { + crate::Error::DataInvalid { + message: "Bucket search sources require a captured table snapshot".into(), + source: None, + } + })? + } + None => { + self.table + .snapshot_manager() + .get_snapshot(plan.snapshot_id) + .await? + } + }; + let sources = primary_key_sources(&self.table, &self.context, &plan).await?; + VectorScanPlan::new(Some(snapshot), sources) } } diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index d8fc37f8a..eef281f29 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -98,6 +98,7 @@ impl<'a> VectorSearchBuilder<'a> { message: "Vector column must be set via with_vector_column()".to_string(), })?; VectorScan::new(self.table, column, self.filter.as_ref(), None, None) + .map(|scan| scan.with_options(&self.options)) } /// Create an owned reader; query errors are reported before planning. diff --git a/crates/paimon/tests/pk_vector_bucket_split_read_test.rs b/crates/paimon/tests/pk_vector_bucket_split_read_test.rs index 25edf9d95..6311b1819 100644 --- a/crates/paimon/tests/pk_vector_bucket_split_read_test.rs +++ b/crates/paimon/tests/pk_vector_bucket_split_read_test.rs @@ -194,7 +194,7 @@ async fn read_over_splits(table: &Table, splits: &[Vec], limit: usize) -> Ve .with_query_vector(vec![0.0, 0.0]) .with_limit(limit); async { - let plan = builder.new_scan()?.plan_from_bucket_splits(splits)?; + let plan = builder.new_scan()?.plan_from_bucket_splits(splits).await?; builder .new_read()? .read(plan) @@ -298,7 +298,10 @@ async fn rejects_an_empty_split_list() { .with_query_vector(vec![0.0, 0.0]) .with_limit(3); let error = match async { - let plan = builder.new_scan()?.plan_from_bucket_splits(Vec::new())?; + let plan = builder + .new_scan()? + .plan_from_bucket_splits(Vec::new()) + .await?; builder .new_read()? .read(plan) @@ -404,7 +407,7 @@ async fn typed_bucket_plan_is_reusable_without_the_index_manifest() { let mut builder = table.new_vector_search_builder(); builder.with_vector_column(VECTOR_COLUMN); let scan = builder.new_scan().unwrap(); - let plan = scan.plan_from_bucket_splits(splits).unwrap(); + let plan = scan.plan_from_bucket_splits(splits).await.unwrap(); let snapshot_id = plan.snapshot_id().unwrap(); let manager = table.snapshot_manager(); let snapshot = manager.get_snapshot(snapshot_id).await.unwrap(); From 4505d08cc48122aeb7efc3a5ca358180f2dbb256 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 22:47:08 +0800 Subject: [PATCH 3/6] Expose typed search plan and split objects in Python bindings --- .../python/pypaimon_rust/datafusion.pyi | 82 ++- bindings/python/src/context.rs | 7 + bindings/python/src/full_text_search.rs | 36 +- bindings/python/src/hybrid_search.rs | 85 ++- bindings/python/src/lib.rs | 1 + bindings/python/src/search_plan.rs | 644 ++++++++++++++++++ bindings/python/src/snapshot.rs | 21 +- bindings/python/src/vector_search.rs | 446 ++++++------ crates/paimon/src/table/search_scan_split.rs | 4 +- 9 files changed, 1047 insertions(+), 279 deletions(-) create mode 100644 bindings/python/src/search_plan.rs diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 391c5c9de..3dc875da6 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -941,12 +941,46 @@ class BatchVectorSearchBuilder: class VectorScan: def scan(self) -> "VectorScanPlan": ... +class IndexFileMeta: + """Immutable index metadata obtained from search source splits.""" + def index_type(self) -> str: ... + def file_name(self) -> str: ... + def file_size(self) -> int: ... + def row_count(self) -> int: ... + +class IndexVectorSearchSplit: + def __init__(self, row_range_start: int, row_range_end: int, + vector_index_files: Sequence[IndexFileMeta], + scalar_index_files: Sequence[IndexFileMeta] = ()) -> None: ... + def row_range_start(self) -> int: ... + def row_range_end(self) -> int: ... + def vector_index_files(self) -> List[IndexFileMeta]: ... + def scalar_index_files(self) -> List[IndexFileMeta]: ... + +class RawVectorSearchSplit: + def __init__(self, row_ranges: Sequence[Tuple[int, int]], + scalar_index_files: Sequence[IndexFileMeta] = (), + index_type: Optional[str] = None) -> None: ... + def row_ranges(self) -> List[Tuple[int, int]]: ... + def scalar_index_files(self) -> List[IndexFileMeta]: ... + def index_type(self) -> Optional[str]: ... + +class PrimaryKeyVectorSearchSplit: + def __init__(self, data_split: Split, payloads: Sequence[IndexFileMeta], + uncovered_data_files: Sequence[str], + row_ranges_by_file: Dict[str, Sequence[Tuple[int, int]]]) -> None: ... + def data_split(self) -> Split: ... + def payloads(self) -> List[IndexFileMeta]: ... + def uncovered_data_files(self) -> List[str]: ... + def row_ranges_by_file(self) -> Dict[str, List[Tuple[int, int]]]: ... + class VectorScanPlan: - def __init__(self, snapshot_json: Optional[str], splits_json: str) -> None: - """Capture resolved Java Snapshot JSON and exported search source splits; None pins empty.""" - ... - def metadata(self) -> str: ... - def snapshot_id(self) -> Optional[int]: ... + def __init__(self, splits: Sequence[Union[IndexVectorSearchSplit, RawVectorSearchSplit, + PrimaryKeyVectorSearchSplit]], + snapshot: Optional[Snapshot] = None) -> None: ... + def snapshot(self) -> Optional[Snapshot]: ... + def splits(self) -> List[Union[IndexVectorSearchSplit, RawVectorSearchSplit, + PrimaryKeyVectorSearchSplit]]: ... class VectorRead: def read_plan(self, plan: VectorScanPlan) -> "GlobalIndexResult": ... @@ -962,8 +996,8 @@ class GlobalIndexResultReadBuilder: class HybridSearchBuilder: def with_limit(self, limit: int) -> HybridSearchBuilder: ... - def with_snapshot(self, snapshot_json: Optional[str]) -> HybridSearchBuilder: - """Reuse resolved Java Snapshot JSON; None pins an empty read view.""" + def with_snapshot(self, snapshot: Optional[Snapshot]) -> HybridSearchBuilder: + """Pin the resolved snapshot object; None pins an empty read view.""" ... def with_ranker(self, ranker: Optional[str]) -> HybridSearchBuilder: ... def with_rrf_ranker(self) -> HybridSearchBuilder: ... @@ -989,12 +1023,36 @@ class FullTextSearchBuilder: class FullTextScan: def scan(self) -> FullTextScanPlan: ... +class IndexFullTextSearchSplit: + def __init__(self, column_name: str, row_range_start: int, row_range_end: int, + full_text_index_files: Sequence[IndexFileMeta], + scalar_index_files: Sequence[IndexFileMeta] = ()) -> None: ... + def column_name(self) -> str: ... + def row_range_start(self) -> int: ... + def row_range_end(self) -> int: ... + def full_text_index_files(self) -> List[IndexFileMeta]: ... + def scalar_index_files(self) -> List[IndexFileMeta]: ... + +class RawFullTextSearchSplit: + def __init__(self, row_ranges: Sequence[Tuple[int, int]]) -> None: ... + def row_ranges(self) -> List[Tuple[int, int]]: ... + +class PrimaryKeyFullTextSearchSplit: + def __init__(self, data_split: Split, payloads: Sequence[IndexFileMeta], + uncovered_data_files: Sequence[str], + row_ranges_by_file: Dict[str, Sequence[Tuple[int, int]]]) -> None: ... + def data_split(self) -> Split: ... + def payloads(self) -> List[IndexFileMeta]: ... + def uncovered_data_files(self) -> List[str]: ... + def row_ranges_by_file(self) -> Dict[str, List[Tuple[int, int]]]: ... + class FullTextScanPlan: - def __init__(self, snapshot_json: Optional[str], splits_json: str) -> None: - """Capture resolved Java Snapshot JSON and exported search source splits; None pins empty.""" - ... - def metadata(self) -> str: ... - def snapshot_id(self) -> Optional[int]: ... + def __init__(self, splits: Sequence[Union[IndexFullTextSearchSplit, RawFullTextSearchSplit, + PrimaryKeyFullTextSearchSplit]], + snapshot: Optional[Snapshot] = None) -> None: ... + def snapshot(self) -> Optional[Snapshot]: ... + def splits(self) -> List[Union[IndexFullTextSearchSplit, RawFullTextSearchSplit, + PrimaryKeyFullTextSearchSplit]]: ... class FullTextRead: def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult: ... diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 170b3eabd..102a3a4ae 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -1004,6 +1004,13 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> this.add_class::()?; this.add_class::()?; this.add_class::()?; + this.add_class::()?; + this.add_class::()?; + this.add_class::()?; + this.add_class::()?; + this.add_class::()?; + this.add_class::()?; + this.add_class::()?; this.add_class::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/full_text_search.rs b/bindings/python/src/full_text_search.rs index 383fb49e7..8889ee7c8 100644 --- a/bindings/python/src/full_text_search.rs +++ b/bindings/python/src/full_text_search.rs @@ -20,11 +20,13 @@ use crate::error::to_py_err; use crate::global_index_result::PyGlobalIndexResult; use crate::predicate::dict_to_table_predicate; +use crate::search_plan::{full_text_splits_from_python, full_text_splits_to_python}; +use crate::snapshot::PySnapshot; use paimon::spec::Predicate; use paimon::table::{FullTextRead, FullTextScan, FullTextScanPlan, Table}; use paimon_datafusion::runtime::runtime; use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyTuple}; use std::sync::Arc; #[derive(Default)] @@ -151,25 +153,31 @@ pub struct PyFullTextScanPlan { #[pymethods] impl PyFullTextScanPlan { #[new] - #[pyo3(signature = (snapshot_json, splits_json))] - fn new(snapshot_json: Option<&str>, splits_json: &str) -> PyResult { - let (snapshot, splits) = - crate::vector_search::search_plan_from_json(snapshot_json, splits_json)?; + #[pyo3(signature = (splits, snapshot=None))] + fn new(splits: &Bound<'_, PyAny>, snapshot: Option>) -> PyResult { Ok(Self { - inner: FullTextScanPlan::new(snapshot, splits).map_err(to_py_err)?, + inner: FullTextScanPlan::new( + snapshot.map(|snapshot| snapshot.inner.clone()), + full_text_splits_from_python(splits)?, + ) + .map_err(to_py_err)?, }) } - fn metadata(&self) -> PyResult { - crate::vector_search::search_plan_metadata_json( - self.inner.snapshot(), - self.inner.snapshot_id(), - self.inner.splits(), - ) + fn snapshot(&self) -> Option { + self.inner.snapshot().cloned().map(PySnapshot::new) } - fn snapshot_id(&self) -> Option { - self.inner.snapshot_id() + fn splits(&self, py: Python<'_>) -> PyResult>> { + full_text_splits_to_python(py, self.inner.splits()) + } + + fn __reduce__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<(Py, Py)> { + let plan = slf.borrow(); + let args = (plan.splits(py)?, plan.snapshot()) + .into_pyobject(py)? + .unbind(); + Ok((slf.get_type().unbind().into_any(), args)) } } diff --git a/bindings/python/src/hybrid_search.rs b/bindings/python/src/hybrid_search.rs index 374e5ac54..b946f5401 100644 --- a/bindings/python/src/hybrid_search.rs +++ b/bindings/python/src/hybrid_search.rs @@ -20,6 +20,7 @@ use crate::error::to_py_err; use crate::global_index_result::PyGlobalIndexResult; use crate::predicate::dict_to_table_predicate; +use crate::snapshot::PySnapshot; use paimon::spec::{Predicate, Snapshot}; use paimon::table::{HybridSearchRanker, HybridSearchRoute, Table}; use paimon_datafusion::runtime::runtime; @@ -122,19 +123,13 @@ impl PyHybridSearchBuilder { slf } - #[pyo3(signature = (snapshot_json))] + #[pyo3(signature = (snapshot))] fn with_snapshot<'py>( mut slf: PyRefMut<'py, Self>, - snapshot_json: Option<&str>, - ) -> PyResult> { - let snapshot = snapshot_json - .map(serde_json::from_str) - .transpose() - .map_err(|error| { - pyo3::exceptions::PyValueError::new_err(format!("Invalid snapshot JSON: {error}")) - })?; - slf.pinned_snapshot = Some(snapshot); - Ok(slf) + snapshot: Option>, + ) -> PyRefMut<'py, Self> { + slf.pinned_snapshot = Some(snapshot.map(|snapshot| snapshot.inner.clone())); + slf } #[pyo3(signature = (ranker))] @@ -206,3 +201,71 @@ impl PyHybridSearchBuilder { self.execute(py) } } + +#[cfg(test)] +mod tests { + use super::*; + use paimon::catalog::Identifier; + use paimon::io::FileIOBuilder; + use paimon::spec::{CommitKind, DataType, IntType, Schema, TableSchema}; + + #[test] + fn hybrid_snapshot_uses_typed_metadata_and_distinguishes_empty_capture() { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .build() + .unwrap(); + let table = Arc::new(Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "hybrid_typed_snapshot"), + "memory:/hybrid_typed_snapshot".into(), + TableSchema::new(0, &schema), + None, + )); + let snapshot = Snapshot::builder() + .version(3) + .id(7) + .schema_id(0) + .base_manifest_list(String::new()) + .delta_manifest_list(String::new()) + .commit_user("captured".into()) + .commit_identifier(1) + .commit_kind(CommitKind::APPEND) + .time_millis(0) + .build(); + Python::attach(|py| { + let builder = Py::new(py, PyHybridSearchBuilder::new(table)).unwrap(); + assert!(builder.borrow(py).pinned_snapshot.is_none()); + let captured = Py::new(py, PySnapshot::new(snapshot.clone())).unwrap(); + let returned = builder + .bind(py) + .call_method1("with_snapshot", (captured,)) + .unwrap(); + assert!(returned.is(builder.bind(py))); + assert_eq!( + builder + .borrow(py) + .pinned_snapshot + .as_ref() + .unwrap() + .as_ref(), + Some(&snapshot), + ); + let error = builder + .bind(py) + .call_method1("with_snapshot", ("{}",)) + .unwrap_err(); + assert!(error.is_instance_of::(py)); + builder + .bind(py) + .call_method1("with_snapshot", (py.None(),)) + .unwrap(); + assert!(builder + .borrow(py) + .pinned_snapshot + .as_ref() + .unwrap() + .is_none()); + }); + } +} diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 0b5ee863b..49b694209 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -34,6 +34,7 @@ mod oss_cpp; mod predicate; mod read; mod schema; +mod search_plan; mod search_task; mod shard_update; mod system_table; diff --git a/bindings/python/src/search_plan.rs b/bindings/python/src/search_plan.rs new file mode 100644 index 000000000..85badc6a0 --- /dev/null +++ b/bindings/python/src/search_plan.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. + +//! Typed, transportable source metadata for search plans. + +use std::collections::HashMap; + +use paimon::spec::IndexFileMeta; +use paimon::table::{RowRange, SearchScanSplit}; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyByteArray, PyBytes, PySequence, PyString}; + +use crate::error::to_py_err; +use crate::read::PySplit; + +#[pyclass(name = "IndexFileMeta", module = "pypaimon_rust.datafusion", frozen)] +pub struct PyIndexFileMeta { + pub(crate) inner: IndexFileMeta, +} + +#[pymethods] +impl PyIndexFileMeta { + fn index_type(&self) -> String { + self.inner.index_type.clone() + } + + fn file_name(&self) -> String { + self.inner.file_name.clone() + } + + fn file_size(&self) -> i64 { + self.inner.file_size + } + + fn row_count(&self) -> i64 { + self.inner.row_count + } + + fn _serialize<'py>(&self, py: Python<'py>) -> PyResult> { + let bytes = self.inner.to_serialized_row_data().map_err(to_py_err)?; + Ok(PyBytes::new(py, &bytes)) + } + + #[staticmethod] + fn _deserialize(state: &Bound<'_, PyBytes>) -> PyResult { + Ok(Self { + inner: IndexFileMeta::from_serialized_row_data(state.as_bytes()).map_err(to_py_err)?, + }) + } + + fn __reduce__<'py>( + slf: &Bound<'py, Self>, + py: Python<'py>, + ) -> PyResult<(Py, (Py,))> { + let bytes = slf.borrow()._serialize(py)?.unbind(); + Ok((slf.get_type().getattr("_deserialize")?.unbind(), (bytes,))) + } +} + +fn range(from: i64, to: i64) -> PyResult { + if from < 0 || from > to { + return Err(PyValueError::new_err( + "Search row ranges must be nonnegative and ordered", + )); + } + Ok(RowRange::new(from, to)) +} + +fn ranges(values: Vec<(i64, i64)>) -> PyResult> { + values + .into_iter() + .map(|(from, to)| range(from, to)) + .collect() +} + +fn range_values(values: &[RowRange]) -> Vec<(i64, i64)> { + values + .iter() + .map(|value| (value.from(), value.to())) + .collect() +} + +fn index_files(py: Python<'_>, values: Vec>) -> Vec { + values + .into_iter() + .map(|value| value.borrow(py).inner.clone()) + .collect() +} + +fn file_values(values: &[IndexFileMeta]) -> Vec { + values + .iter() + .cloned() + .map(|inner| PyIndexFileMeta { inner }) + .collect() +} + +fn position_ranges( + values: HashMap>, +) -> PyResult>> { + values + .into_iter() + .map(|(name, values)| Ok((name, ranges(values)?))) + .collect() +} + +fn position_values(values: &HashMap>) -> HashMap> { + values + .iter() + .map(|(name, values)| (name.clone(), range_values(values))) + .collect() +} + +fn decode_split(state: &Bound<'_, PyBytes>) -> PyResult { + let split: SearchScanSplit = serde_json::from_slice(state.as_bytes()) + .map_err(|error| PyValueError::new_err(format!("Invalid search split state: {error}")))?; + // Pickle is an internal encoding, but reconstruction still preserves the + // public constructors' range invariants. + let check = |value: &RowRange| range(value.from(), value.to()).map(|_| ()); + match &split { + SearchScanSplit::Index { row_range, .. } => check(row_range)?, + SearchScanSplit::Raw { row_ranges, .. } => { + for value in row_ranges { + check(value)?; + } + } + SearchScanSplit::PrimaryKey { + row_ranges_by_file, .. + } => { + for value in row_ranges_by_file.values().flatten() { + check(value)?; + } + } + } + Ok(split) +} + +// Pickle only uses private factories; the public constructors and getters +// remain typed domain APIs and carry no scheduling protocol. +macro_rules! split_methods { + ($class:ident, $kind:pat, { $($methods:item)* }) => { + #[pymethods] + impl $class { + $($methods)* + + fn _serialize<'py>(&self, py: Python<'py>) -> PyResult> { + let bytes = serde_json::to_vec(&self.inner) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok(PyBytes::new(py, &bytes)) + } + + #[staticmethod] + fn _deserialize(state: &Bound<'_, PyBytes>) -> PyResult { + let inner = decode_split(state)?; + if !matches!(&inner, $kind) { + return Err(PyValueError::new_err("Search split state has the wrong source kind")); + } + Ok(Self { inner }) + } + + fn __reduce__<'py>( + slf: &Bound<'py, Self>, + py: Python<'py>, + ) -> PyResult<(Py, (Py,))> { + let bytes = slf.borrow()._serialize(py)?.unbind(); + Ok((slf.get_type().getattr("_deserialize")?.unbind(), (bytes,))) + } + } + }; +} + +#[pyclass( + name = "IndexVectorSearchSplit", + module = "pypaimon_rust.datafusion", + frozen +)] +pub struct PyIndexVectorSearchSplit { + pub(crate) inner: SearchScanSplit, +} + +split_methods!(PyIndexVectorSearchSplit, SearchScanSplit::Index { .. }, { + #[new] + #[pyo3(signature = (row_range_start, row_range_end, vector_index_files, scalar_index_files=Vec::new()))] + fn new( + py: Python<'_>, + row_range_start: i64, + row_range_end: i64, + vector_index_files: Vec>, + scalar_index_files: Vec>, + ) -> PyResult { + Ok(Self { + inner: SearchScanSplit::Index { + column: String::new(), + row_range: range(row_range_start, row_range_end)?, + index_files: index_files(py, vector_index_files), + scalar_index_files: index_files(py, scalar_index_files), + }, + }) + } + fn row_range_start(&self) -> i64 { + let SearchScanSplit::Index { row_range, .. } = &self.inner else { + unreachable!() + }; + row_range.from() + } + fn row_range_end(&self) -> i64 { + let SearchScanSplit::Index { row_range, .. } = &self.inner else { + unreachable!() + }; + row_range.to() + } + fn vector_index_files(&self) -> Vec { + let SearchScanSplit::Index { index_files, .. } = &self.inner else { + unreachable!() + }; + file_values(index_files) + } + fn scalar_index_files(&self) -> Vec { + let SearchScanSplit::Index { + scalar_index_files, .. + } = &self.inner + else { + unreachable!() + }; + file_values(scalar_index_files) + } +}); + +#[pyclass( + name = "RawVectorSearchSplit", + module = "pypaimon_rust.datafusion", + frozen +)] +pub struct PyRawVectorSearchSplit { + pub(crate) inner: SearchScanSplit, +} + +split_methods!(PyRawVectorSearchSplit, SearchScanSplit::Raw { .. }, { + #[new] + #[pyo3(signature = (row_ranges, scalar_index_files=Vec::new(), index_type=None))] + fn new( + py: Python<'_>, + row_ranges: Vec<(i64, i64)>, + scalar_index_files: Vec>, + index_type: Option, + ) -> PyResult { + Ok(Self { + inner: SearchScanSplit::Raw { + row_ranges: ranges(row_ranges)?, + scalar_index_files: index_files(py, scalar_index_files), + index_type, + }, + }) + } + fn row_ranges(&self) -> Vec<(i64, i64)> { + let SearchScanSplit::Raw { row_ranges, .. } = &self.inner else { + unreachable!() + }; + range_values(row_ranges) + } + fn scalar_index_files(&self) -> Vec { + let SearchScanSplit::Raw { + scalar_index_files, .. + } = &self.inner + else { + unreachable!() + }; + file_values(scalar_index_files) + } + fn index_type(&self) -> Option { + let SearchScanSplit::Raw { index_type, .. } = &self.inner else { + unreachable!() + }; + index_type.clone() + } +}); + +#[pyclass( + name = "PrimaryKeyVectorSearchSplit", + module = "pypaimon_rust.datafusion", + frozen +)] +pub struct PyPrimaryKeyVectorSearchSplit { + pub(crate) inner: SearchScanSplit, +} + +split_methods!( + PyPrimaryKeyVectorSearchSplit, + SearchScanSplit::PrimaryKey { .. }, + { + #[new] + fn new( + py: Python<'_>, + data_split: &PySplit, + payloads: Vec>, + uncovered_data_files: Vec, + row_ranges_by_file: HashMap>, + ) -> PyResult { + Ok(Self { + inner: SearchScanSplit::PrimaryKey { + data_split: data_split.inner.clone(), + payloads: index_files(py, payloads), + uncovered_data_files, + row_ranges_by_file: position_ranges(row_ranges_by_file)?, + }, + }) + } + fn data_split(&self) -> PySplit { + let SearchScanSplit::PrimaryKey { data_split, .. } = &self.inner else { + unreachable!() + }; + PySplit { + inner: data_split.clone(), + } + } + fn payloads(&self) -> Vec { + let SearchScanSplit::PrimaryKey { payloads, .. } = &self.inner else { + unreachable!() + }; + file_values(payloads) + } + fn uncovered_data_files(&self) -> Vec { + let SearchScanSplit::PrimaryKey { + uncovered_data_files, + .. + } = &self.inner + else { + unreachable!() + }; + uncovered_data_files.clone() + } + fn row_ranges_by_file(&self) -> HashMap> { + let SearchScanSplit::PrimaryKey { + row_ranges_by_file, .. + } = &self.inner + else { + unreachable!() + }; + position_values(row_ranges_by_file) + } + } +); + +#[pyclass( + name = "IndexFullTextSearchSplit", + module = "pypaimon_rust.datafusion", + frozen +)] +pub struct PyIndexFullTextSearchSplit { + pub(crate) inner: SearchScanSplit, +} + +split_methods!(PyIndexFullTextSearchSplit, SearchScanSplit::Index { .. }, { + #[new] + #[pyo3(signature = (column_name, row_range_start, row_range_end, full_text_index_files, scalar_index_files=Vec::new()))] + fn new( + py: Python<'_>, + column_name: String, + row_range_start: i64, + row_range_end: i64, + full_text_index_files: Vec>, + scalar_index_files: Vec>, + ) -> PyResult { + if column_name.is_empty() { + return Err(PyValueError::new_err( + "Full-text source column must not be empty", + )); + } + Ok(Self { + inner: SearchScanSplit::Index { + column: column_name, + row_range: range(row_range_start, row_range_end)?, + index_files: index_files(py, full_text_index_files), + scalar_index_files: index_files(py, scalar_index_files), + }, + }) + } + fn column_name(&self) -> String { + let SearchScanSplit::Index { column, .. } = &self.inner else { + unreachable!() + }; + column.clone() + } + fn row_range_start(&self) -> i64 { + let SearchScanSplit::Index { row_range, .. } = &self.inner else { + unreachable!() + }; + row_range.from() + } + fn row_range_end(&self) -> i64 { + let SearchScanSplit::Index { row_range, .. } = &self.inner else { + unreachable!() + }; + row_range.to() + } + fn full_text_index_files(&self) -> Vec { + let SearchScanSplit::Index { index_files, .. } = &self.inner else { + unreachable!() + }; + file_values(index_files) + } + fn scalar_index_files(&self) -> Vec { + let SearchScanSplit::Index { + scalar_index_files, .. + } = &self.inner + else { + unreachable!() + }; + file_values(scalar_index_files) + } +}); + +#[pyclass( + name = "RawFullTextSearchSplit", + module = "pypaimon_rust.datafusion", + frozen +)] +pub struct PyRawFullTextSearchSplit { + pub(crate) inner: SearchScanSplit, +} + +split_methods!(PyRawFullTextSearchSplit, SearchScanSplit::Raw { .. }, { + #[new] + fn new(row_ranges: Vec<(i64, i64)>) -> PyResult { + Ok(Self { + inner: SearchScanSplit::Raw { + row_ranges: ranges(row_ranges)?, + scalar_index_files: vec![], + index_type: None, + }, + }) + } + fn row_ranges(&self) -> Vec<(i64, i64)> { + let SearchScanSplit::Raw { row_ranges, .. } = &self.inner else { + unreachable!() + }; + range_values(row_ranges) + } +}); + +#[pyclass( + name = "PrimaryKeyFullTextSearchSplit", + module = "pypaimon_rust.datafusion", + frozen +)] +pub struct PyPrimaryKeyFullTextSearchSplit { + pub(crate) inner: SearchScanSplit, +} + +split_methods!( + PyPrimaryKeyFullTextSearchSplit, + SearchScanSplit::PrimaryKey { .. }, + { + #[new] + fn new( + py: Python<'_>, + data_split: &PySplit, + payloads: Vec>, + uncovered_data_files: Vec, + row_ranges_by_file: HashMap>, + ) -> PyResult { + Ok(Self { + inner: SearchScanSplit::PrimaryKey { + data_split: data_split.inner.clone(), + payloads: index_files(py, payloads), + uncovered_data_files, + row_ranges_by_file: position_ranges(row_ranges_by_file)?, + }, + }) + } + fn data_split(&self) -> PySplit { + let SearchScanSplit::PrimaryKey { data_split, .. } = &self.inner else { + unreachable!() + }; + PySplit { + inner: data_split.clone(), + } + } + fn payloads(&self) -> Vec { + let SearchScanSplit::PrimaryKey { payloads, .. } = &self.inner else { + unreachable!() + }; + file_values(payloads) + } + fn uncovered_data_files(&self) -> Vec { + let SearchScanSplit::PrimaryKey { + uncovered_data_files, + .. + } = &self.inner + else { + unreachable!() + }; + uncovered_data_files.clone() + } + fn row_ranges_by_file(&self) -> HashMap> { + let SearchScanSplit::PrimaryKey { + row_ranges_by_file, .. + } = &self.inner + else { + unreachable!() + }; + position_values(row_ranges_by_file) + } + } +); + +fn split_sequence<'a, 'py>(values: &'a Bound<'py, PyAny>) -> PyResult<&'a Bound<'py, PySequence>> { + if values.is_instance_of::() + || values.is_instance_of::() + || values.is_instance_of::() + { + return Err(PyTypeError::new_err( + "Search plans require a sequence of typed search source splits", + )); + } + values.cast::().map_err(|_| { + PyTypeError::new_err("Search plans require a sequence of typed search source splits") + }) +} + +pub(crate) fn vector_splits_from_python( + values: &Bound<'_, PyAny>, +) -> PyResult> { + split_sequence(values)? + .try_iter()? + .map(|value| { + let value = value?; + if let Ok(value) = value.extract::>() { + return Ok(value.inner.clone()); + } + if let Ok(value) = value.extract::>() { + return Ok(value.inner.clone()); + } + if let Ok(value) = value.extract::>() { + return Ok(value.inner.clone()); + } + Err(PyTypeError::new_err( + "Vector plans require typed vector search source splits", + )) + }) + .collect() +} + +pub(crate) fn full_text_splits_from_python( + values: &Bound<'_, PyAny>, +) -> PyResult> { + split_sequence(values)? + .try_iter()? + .map(|value| { + let value = value?; + if let Ok(value) = value.extract::>() { + return Ok(value.inner.clone()); + } + if let Ok(value) = value.extract::>() { + return Ok(value.inner.clone()); + } + if let Ok(value) = value.extract::>() { + return Ok(value.inner.clone()); + } + Err(PyTypeError::new_err( + "Full-text plans require typed full-text search source splits", + )) + }) + .collect() +} + +pub(crate) fn vector_splits_to_python( + py: Python<'_>, + sources: &[SearchScanSplit], +) -> PyResult>> { + sources + .iter() + .map(|source| match source { + SearchScanSplit::Index { .. } => Py::new( + py, + PyIndexVectorSearchSplit { + inner: source.clone(), + }, + ) + .map(Py::into_any), + SearchScanSplit::Raw { .. } => Py::new( + py, + PyRawVectorSearchSplit { + inner: source.clone(), + }, + ) + .map(Py::into_any), + SearchScanSplit::PrimaryKey { .. } => Py::new( + py, + PyPrimaryKeyVectorSearchSplit { + inner: source.clone(), + }, + ) + .map(Py::into_any), + }) + .collect() +} + +pub(crate) fn full_text_splits_to_python( + py: Python<'_>, + sources: &[SearchScanSplit], +) -> PyResult>> { + sources + .iter() + .map(|source| match source { + SearchScanSplit::Index { .. } => Py::new( + py, + PyIndexFullTextSearchSplit { + inner: source.clone(), + }, + ) + .map(Py::into_any), + SearchScanSplit::Raw { .. } => Py::new( + py, + PyRawFullTextSearchSplit { + inner: source.clone(), + }, + ) + .map(Py::into_any), + SearchScanSplit::PrimaryKey { .. } => Py::new( + py, + PyPrimaryKeyFullTextSearchSplit { + inner: source.clone(), + }, + ) + .map(Py::into_any), + }) + .collect() +} diff --git a/bindings/python/src/snapshot.rs b/bindings/python/src/snapshot.rs index 11c40c98a..825e32a5b 100644 --- a/bindings/python/src/snapshot.rs +++ b/bindings/python/src/snapshot.rs @@ -18,11 +18,12 @@ use std::collections::HashMap; use paimon::spec::Snapshot; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; #[pyclass(name = "Snapshot", module = "pypaimon_rust.datafusion")] pub struct PySnapshot { - inner: Snapshot, + pub(crate) inner: Snapshot, } impl PySnapshot { @@ -33,6 +34,24 @@ impl PySnapshot { #[pymethods] impl PySnapshot { + #[staticmethod] + fn _from_json(value: &str) -> PyResult { + serde_json::from_str(value) + .map(Self::new) + .map_err(|err| PyValueError::new_err(format!("Invalid snapshot metadata: {err}"))) + } + + fn _to_json(&self) -> PyResult { + serde_json::to_string(&self.inner) + .map_err(|err| PyValueError::new_err(format!("Cannot serialize snapshot: {err}"))) + } + + fn __reduce__(slf: &Bound<'_, Self>) -> PyResult<(Py, (String,))> { + let value = slf.borrow()._to_json()?; + let restore = slf.get_type().getattr("_from_json")?.unbind(); + Ok((restore, (value,))) + } + fn id(&self) -> i64 { self.inner.id() } diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index 5e91f7cc6..ffb864eb6 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -19,148 +19,16 @@ use std::collections::HashMap; use std::sync::Arc; use paimon::spec::Predicate; -use paimon::table::{ - BatchVectorRead, DataSplit, RowRange, SearchScanSplit, Table, VectorRead, VectorScan, - VectorScanPlan, -}; +use paimon::table::{BatchVectorRead, Table, VectorRead, VectorScan, VectorScanPlan}; use paimon_datafusion::runtime::runtime; use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyTuple}; use crate::error::to_py_err; use crate::global_index_result::PyGlobalIndexResult; use crate::predicate::dict_to_table_predicate; - -pub(crate) fn search_plan_metadata_json( - snapshot: Option<&paimon::spec::Snapshot>, - snapshot_id: Option, - splits: &[SearchScanSplit], -) -> PyResult { - use serde_json::json; - let files = |files: &[paimon::spec::IndexFileMeta]| { - files - .iter() - .map(|file| file.to_serialized_row_data().map_err(to_py_err)) - .collect::>>() - }; - let splits = splits.iter().map(|split| Ok(match split { - SearchScanSplit::Index { column, row_range, index_files, scalar_index_files } => json!({ - "kind": "index", "column": column, "row_range": row_range, - "index_files": files(index_files)?, "scalar_index_files": files(scalar_index_files)?, - }), - SearchScanSplit::Raw { row_ranges, scalar_index_files, index_type } => json!({ - "kind": "raw", "row_ranges": row_ranges, "scalar_index_files": files(scalar_index_files)?, "index_type": index_type, - }), - SearchScanSplit::PrimaryKey { data_split, payloads, uncovered_data_files, row_ranges_by_file } => json!({ - "kind": "primary-key", "data_split": data_split.serialize_split_v1().map_err(to_py_err)?, - "payloads": files(payloads)?, "uncovered_data_files": uncovered_data_files, - "row_ranges_by_file": row_ranges_by_file, - }), - })).collect::>>()?; - serde_json::to_string( - &json!({"snapshot": snapshot, "snapshot_id": snapshot_id, "splits": splits}), - ) - .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) -} - -/// Decode the same Java metadata codecs used by source-plan export. The binding -/// does not infer work from the snapshot or bind a plan to a live reader. -pub(crate) fn search_plan_from_json( - snapshot_json: Option<&str>, - splits_json: &str, -) -> PyResult<(Option, Vec)> { - use pyo3::exceptions::PyValueError; - use serde_json::Value; - - let invalid = |error: serde_json::Error| PyValueError::new_err(error.to_string()); - let snapshot = snapshot_json - .map(serde_json::from_str) - .transpose() - .map_err(invalid)?; - let values: Vec = serde_json::from_str(splits_json).map_err(invalid)?; - let field = |value: &Value, name: &str| -> PyResult { - value - .get(name) - .cloned() - .ok_or_else(|| PyValueError::new_err(format!("Search split is missing '{name}'"))) - }; - let files = |value: Value| -> PyResult> { - let bytes: Vec> = serde_json::from_value(value).map_err(invalid)?; - bytes - .iter() - .map(|bytes| { - paimon::spec::IndexFileMeta::from_serialized_row_data(bytes).map_err(to_py_err) - }) - .collect() - }; - let validate_ranges = |ranges: &[RowRange]| -> PyResult<()> { - if ranges - .iter() - .any(|range| range.from() < 0 || range.from() > range.to()) - { - return Err(PyValueError::new_err( - "Search row ranges must be nonnegative and ordered", - )); - } - Ok(()) - }; - let splits = values - .iter() - .map(|value| -> PyResult { - let kind: String = serde_json::from_value(field(value, "kind")?).map_err(invalid)?; - Ok(match kind.as_str() { - "index" => { - let row_range: RowRange = - serde_json::from_value(field(value, "row_range")?).map_err(invalid)?; - validate_ranges(std::slice::from_ref(&row_range))?; - SearchScanSplit::Index { - column: serde_json::from_value(field(value, "column")?).map_err(invalid)?, - row_range, - index_files: files(field(value, "index_files")?)?, - scalar_index_files: files(field(value, "scalar_index_files")?)?, - } - } - "raw" => { - let row_ranges: Vec = - serde_json::from_value(field(value, "row_ranges")?).map_err(invalid)?; - validate_ranges(&row_ranges)?; - SearchScanSplit::Raw { - row_ranges, - scalar_index_files: files(field(value, "scalar_index_files")?)?, - index_type: serde_json::from_value(field(value, "index_type")?) - .map_err(invalid)?, - } - } - "primary-key" => { - let bytes: Vec = - serde_json::from_value(field(value, "data_split")?).map_err(invalid)?; - let row_ranges_by_file: HashMap> = - serde_json::from_value(field(value, "row_ranges_by_file")?) - .map_err(invalid)?; - for ranges in row_ranges_by_file.values() { - validate_ranges(ranges)?; - } - SearchScanSplit::PrimaryKey { - data_split: DataSplit::deserialize_split_v1(&bytes).map_err(to_py_err)?, - payloads: files(field(value, "payloads")?)?, - uncovered_data_files: serde_json::from_value(field( - value, - "uncovered_data_files", - )?) - .map_err(invalid)?, - row_ranges_by_file, - } - } - _ => { - return Err(PyValueError::new_err(format!( - "Unknown search split kind '{kind}'" - ))) - } - }) - }) - .collect::>>()?; - Ok((snapshot, splits)) -} +use crate::search_plan::{vector_splits_from_python, vector_splits_to_python}; +use crate::snapshot::PySnapshot; #[derive(Clone, Default)] struct SearchConfig { @@ -429,24 +297,31 @@ pub struct PyVectorScanPlan { #[pymethods] impl PyVectorScanPlan { #[new] - #[pyo3(signature = (snapshot_json, splits_json))] - fn new(snapshot_json: Option<&str>, splits_json: &str) -> PyResult { - let (snapshot, splits) = search_plan_from_json(snapshot_json, splits_json)?; + #[pyo3(signature = (splits, snapshot=None))] + fn new(splits: &Bound<'_, PyAny>, snapshot: Option>) -> PyResult { Ok(Self { - inner: VectorScanPlan::new(snapshot, splits).map_err(to_py_err)?, + inner: VectorScanPlan::new( + snapshot.map(|snapshot| snapshot.inner.clone()), + vector_splits_from_python(splits)?, + ) + .map_err(to_py_err)?, }) } - fn metadata(&self) -> PyResult { - search_plan_metadata_json( - self.inner.snapshot(), - self.inner.snapshot_id(), - self.inner.splits(), - ) + fn snapshot(&self) -> Option { + self.inner.snapshot().cloned().map(PySnapshot::new) + } + + fn splits(&self, py: Python<'_>) -> PyResult>> { + vector_splits_to_python(py, self.inner.splits()) } - fn snapshot_id(&self) -> Option { - self.inner.snapshot_id() + fn __reduce__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<(Py, Py)> { + let plan = slf.borrow(); + let args = (plan.splits(py)?, plan.snapshot()) + .into_pyobject(py)? + .unbind(); + Ok((slf.get_type().unbind().into_any(), args)) } } @@ -527,27 +402,24 @@ impl PyBatchVectorRead { #[cfg(test)] mod source_plan_tests { use super::*; - use crate::full_text_search::PyFullTextScanPlan; - use paimon::spec::{BinaryRow, CommitKind, IndexFileMeta, Snapshot}; - use pyo3::types::PyModule; - use serde_json::{json, Value}; - - fn snapshot_json() -> String { - serde_json::to_string( - &Snapshot::builder() - .version(3) - .id(7) - .schema_id(2) - .base_manifest_list("base".into()) - .delta_manifest_list("delta".into()) - .commit_user("test".into()) - .commit_identifier(1) - .commit_kind(CommitKind::APPEND) - .time_millis(0) - .next_row_id(Some(10)) - .build(), - ) - .unwrap() + use crate::search_plan::*; + use paimon::spec::{BinaryRow, CommitKind, GlobalIndexMeta, IndexFileMeta, Snapshot}; + use paimon::table::{DataSplit, RowRange, SearchScanSplit}; + use pyo3::types::{PyBytes, PyList, PyModule}; + + fn snapshot() -> Snapshot { + Snapshot::builder() + .version(3) + .id(7) + .schema_id(2) + .base_manifest_list("base".into()) + .delta_manifest_list("delta".into()) + .commit_user("test".into()) + .commit_identifier(1) + .commit_kind(CommitKind::APPEND) + .time_millis(0) + .next_row_id(Some(10)) + .build() } fn index_file() -> IndexFileMeta { @@ -558,14 +430,37 @@ mod source_plan_tests { row_count: 10, deletion_vectors_ranges: None, external_path: Some("s3://bucket/index-1".into()), - global_index_meta: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: 0, + row_range_end: 9, + index_field_id: 2, + extra_field_ids: Some(vec![3]), + index_meta: Some(vec![1, 2, 3]), + source_meta: None, + }), } } + fn module(py: Python<'_>) -> Bound<'_, PyModule> { + // Initialize the actual binding module in this embedded interpreter. + // This exercises class registration and pickle names without installing + // an extension into the PyPaimon test environment. + let package = PyModule::new(py, "pypaimon_rust").unwrap(); + crate::pypaimon_rust(py, &package).unwrap(); + py.import("sys") + .unwrap() + .getattr("modules") + .unwrap() + .set_item("pypaimon_rust", &package) + .unwrap(); + package.getattr("datafusion").unwrap().cast_into().unwrap() + } + #[test] fn source_metadata_reuses_java_index_and_split_codecs() { - Python::attach(|_| { - let split = DataSplit::builder() + Python::attach(|py| { + let module = module(py); + let data_split = DataSplit::builder() .with_snapshot(7) .with_partition(BinaryRow::new(0)) .with_bucket(0) @@ -587,7 +482,7 @@ mod source_plan_tests { index_type: Some("lucene-vector".into()), }, SearchScanSplit::PrimaryKey { - data_split: split, + data_split: data_split.clone(), payloads: vec![index_file()], uncovered_data_files: vec!["data-1.parquet".into()], row_ranges_by_file: HashMap::from([( @@ -596,82 +491,155 @@ mod source_plan_tests { )]), }, ]; - let snapshot: Snapshot = serde_json::from_str(&snapshot_json()).unwrap(); - let metadata = search_plan_metadata_json(Some(&snapshot), Some(7), &sources).unwrap(); - let value: Value = serde_json::from_str(&metadata).unwrap(); - let (restored_snapshot, restored) = - search_plan_from_json(Some(&snapshot_json()), &value["splits"].to_string()) - .unwrap(); - assert_eq!(restored_snapshot.as_ref(), Some(&snapshot)); - assert_eq!( - serde_json::from_str::( - &search_plan_metadata_json(restored_snapshot.as_ref(), Some(7), &restored) - .unwrap() + let locals = PyDict::new(py); + locals.set_item("sdk", &module).unwrap(); + locals + .set_item("snapshot", PySnapshot::new(snapshot())) + .unwrap(); + locals + .set_item( + "index_bytes", + PyBytes::new(py, &index_file().to_serialized_row_data().unwrap()), + ) + .unwrap(); + locals + .set_item( + "split_bytes", + PyBytes::new(py, &data_split.serialize_split_v1().unwrap()), + ) + .unwrap(); + for full_text in [false, true] { + let objects = if full_text { + full_text_splits_to_python(py, &sources).unwrap() + } else { + vector_splits_to_python(py, &sources).unwrap() + }; + let values = PyList::new(py, &objects).unwrap(); + let restored = if full_text { + full_text_splits_from_python(values.as_any()).unwrap() + } else { + vector_splits_from_python(values.as_any()).unwrap() + }; + assert_eq!( + serde_json::to_value(restored).unwrap(), + serde_json::to_value(&sources).unwrap() + ); + locals.set_item("sources", values).unwrap(); + locals.set_item("full_text", full_text).unwrap(); + py.run( + cr#" +import pickle +cls = sdk.FullTextScanPlan if full_text else sdk.VectorScanPlan +pk_cls = sdk.PrimaryKeyFullTextSearchSplit if full_text else sdk.PrimaryKeyVectorSearchSplit +for source in sources: + copied = pickle.loads(pickle.dumps(source)) + assert type(copied) is type(source) + assert copied._serialize() == source._serialize() +meta = sdk.IndexFileMeta._deserialize(index_bytes) +assert meta.file_name() == 'index-1' +assert pickle.loads(pickle.dumps(meta))._serialize() == index_bytes +pk = sources[2] +assert pk.data_split().serialize() == split_bytes +rebuilt = pk_cls(pk.data_split(), pk.payloads(), pk.uncovered_data_files(), pk.row_ranges_by_file()) +plan = cls([rebuilt], snapshot) +restored = pickle.loads(pickle.dumps(plan)) +assert restored.snapshot().id() == 7 +assert restored.splits()[0].data_split().serialize() == split_bytes +assert restored.splits()[0].payloads()[0]._serialize() == index_bytes +assert restored.splits()[0].row_ranges_by_file() == {'data-1.parquet': [(1, 2)]} +"#, + Some(&locals), + Some(&locals), ) - .unwrap(), - value, - ); - assert_eq!( - value["splits"][0]["index_files"][0], - json!(index_file().to_serialized_row_data().unwrap()) - ); + .unwrap(); + } }); } #[test] - fn python_source_plan_constructors_capture_snapshot_and_splits() { + fn python_source_plan_constructors_capture_typed_snapshot_and_splits() { Python::attach(|py| { - let module = PyModule::new(py, "plans").unwrap(); - module.add_class::().unwrap(); - module.add_class::().unwrap(); - let sources = json!([{ - "kind": "raw", "row_ranges": [{"from": 2, "to": 4}], - "scalar_index_files": [], "index_type": null, - }]) - .to_string(); - for name in ["VectorScanPlan", "FullTextScanPlan"] { - let class = module.getattr(name).unwrap(); - let plan = class.call1((snapshot_json(), &sources)).unwrap(); - assert_eq!( - plan.call_method0("snapshot_id") - .unwrap() - .extract::() - .unwrap(), - 7 - ); - let metadata: String = plan.call_method0("metadata").unwrap().extract().unwrap(); - let value: Value = serde_json::from_str(&metadata).unwrap(); - assert_eq!(value["splits"].to_string(), sources); - assert_eq!(value["snapshot"]["id"], json!(7)); - let copied = class - .call1((value["snapshot"].to_string(), value["splits"].to_string())) - .unwrap(); - let copied_metadata: String = - copied.call_method0("metadata").unwrap().extract().unwrap(); - assert_eq!(metadata, copied_metadata); - let empty = class.call1((py.None(), "[]")).unwrap(); - assert!(empty.call_method0("snapshot_id").unwrap().is_none()); - assert!(class.call1((py.None(), &sources)).is_err()); - } + let locals = PyDict::new(py); + locals.set_item("sdk", module(py)).unwrap(); + locals + .set_item("snapshot", PySnapshot::new(snapshot())) + .unwrap(); + locals + .set_item( + "meta", + PyIndexFileMeta { + inner: index_file(), + }, + ) + .unwrap(); + py.run(cr#" +import pickle +for cls, raw_cls, indexed in ( + (sdk.VectorScanPlan, sdk.RawVectorSearchSplit, sdk.IndexVectorSearchSplit(0, 9, [meta])), + (sdk.FullTextScanPlan, sdk.RawFullTextSearchSplit, sdk.IndexFullTextSearchSplit('body', 0, 9, [meta])), +): + ranges = [(2, 4)] + raw = raw_cls(ranges) + sources = [indexed, raw] + plan = cls(sources, snapshot) + sources.clear() + ranges.clear() + assert plan.snapshot().id() == 7 + assert len(plan.splits()) == 2 + assert plan.splits()[1].row_ranges() == [(2, 4)] + plan.splits().clear() + assert len(plan.splits()) == 2 + assert len(cls(tuple(plan.splits()), snapshot).splits()) == 2 + restored = pickle.loads(pickle.dumps(plan)) + assert restored.snapshot()._to_json() == snapshot._to_json() + assert restored.splits()[0].row_range_start() == 0 + assert restored.splits()[1].row_ranges() == [(2, 4)] + assert not hasattr(plan, 'metadata') + empty = cls([]) + assert empty.snapshot() is None and empty.splits() == [] + assert pickle.loads(pickle.dumps(empty)).snapshot() is None +"#, Some(&locals), Some(&locals)).unwrap(); }); } #[test] - fn source_plan_decode_rejects_corrupt_codecs_and_ranges() { - Python::attach(|_| { - for value in [ - json!([{"kind": "unknown"}]), - json!([{"kind": "index", "column": "v", "row_range": {"from": 2, "to": 1}, "index_files": [], "scalar_index_files": []}]), - json!([{"kind": "raw", "row_ranges": [{"from": -1, "to": 1}], "scalar_index_files": [], "index_type": null}]), - json!([{"kind": "index", "column": "v", "row_range": {"from": 0, "to": 1}, "index_files": [[1, 2, 3]], "scalar_index_files": []}]), - json!([{"kind": "primary-key", "data_split": [1, 2, 3], "payloads": [], "uncovered_data_files": [], "row_ranges_by_file": {}}]), - ] { - assert!( - search_plan_from_json(Some(&snapshot_json()), &value.to_string()).is_err(), - "{value}" - ); - } - assert!(search_plan_from_json(Some("invalid snapshot"), "[]").is_err()); + fn source_plan_constructors_reject_wire_values_and_wrong_source_types() { + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("sdk", module(py)).unwrap(); + locals + .set_item("snapshot", PySnapshot::new(snapshot())) + .unwrap(); + py.run( + cr#" +def rejects(error, fn, *args): + try: + fn(*args) + except error: + return + raise AssertionError('Invalid source input was accepted') +for cls, raw_cls, wrong_cls in ( + (sdk.VectorScanPlan, sdk.RawVectorSearchSplit, sdk.RawFullTextSearchSplit), + (sdk.FullTextScanPlan, sdk.RawFullTextSearchSplit, sdk.RawVectorSearchSplit), +): + rejects(TypeError, cls, '[]', snapshot) + rejects(TypeError, cls, '', snapshot) + rejects(TypeError, cls, {}, snapshot) + rejects(TypeError, cls, b'', snapshot) + rejects(TypeError, cls, bytearray(), snapshot) + rejects(TypeError, cls, [{'kind': 'raw'}], snapshot) + rejects(TypeError, cls, [wrong_cls([(0, 1)])], snapshot) + rejects(TypeError, cls, [], '{}') + rejects(ValueError, cls, [raw_cls([(0, 1)])]) + rejects(ValueError, raw_cls, [(-1, 2)]) + rejects(ValueError, raw_cls, [(2, 1)]) +rejects(ValueError, sdk.IndexFileMeta._deserialize, b'bad') +rejects(ValueError, sdk.Split.deserialize, b'bad') +"#, + Some(&locals), + Some(&locals), + ) + .unwrap(); }); } } diff --git a/crates/paimon/src/table/search_scan_split.rs b/crates/paimon/src/table/search_scan_split.rs index 42e8ffd04..cf27c7cab 100644 --- a/crates/paimon/src/table/search_scan_split.rs +++ b/crates/paimon/src/table/search_scan_split.rs @@ -228,12 +228,12 @@ pub(super) fn global_entries_from_splits( for split in splits { let scalar = match split { SearchScanSplit::Index { - column, row_range, index_files, scalar_index_files, + .. } => { - if column.is_empty() || row_range.from() < 0 || row_range.to() < row_range.from() { + if row_range.from() < 0 || row_range.to() < row_range.from() { return Err(invalid_source("Invalid global search index split")); } for file in index_files { From 264806286629220d3ec4ab38aab60fafcdbd2652 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 23:03:47 +0800 Subject: [PATCH 4/6] Unify native search builder execution under execute --- .../python/pypaimon_rust/datafusion.pyi | 7 +- bindings/python/src/full_text_search.rs | 2 +- bindings/python/src/hybrid_search.rs | 4 - bindings/python/src/vector_search.rs | 77 ++++++++++++++++++- bindings/python/tests/test_vector_search.py | 26 ++++--- 5 files changed, 94 insertions(+), 22 deletions(-) diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 3dc875da6..81770d157 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -924,7 +924,7 @@ class VectorSearchBuilder: ... def new_vector_search_scan(self) -> "VectorScan": ... def new_vector_search_read(self) -> "VectorRead": ... - def execute_local(self) -> "GlobalIndexResult": ... + def execute(self) -> "GlobalIndexResult": ... class BatchVectorSearchBuilder: def with_vector_column(self, name: str) -> "BatchVectorSearchBuilder": ... @@ -936,7 +936,7 @@ class BatchVectorSearchBuilder: def with_partition_filter(self, predicate: Dict[str, Any]) -> "BatchVectorSearchBuilder": ... def new_vector_search_scan(self) -> "VectorScan": ... def new_batch_vector_search_read(self) -> "BatchVectorRead": ... - def execute_batch_local(self) -> List["GlobalIndexResult"]: ... + def execute(self) -> List["GlobalIndexResult"]: ... class VectorScan: def scan(self) -> "VectorScanPlan": ... @@ -1008,7 +1008,6 @@ class HybridSearchBuilder: weight: float = 1.0, options: Optional[Dict[str, str]] = None) -> HybridSearchBuilder: ... def add_full_text_route(self, field_name: str, query: str, limit: int, weight: float = 1.0, options: Optional[Dict[str, str]] = None) -> HybridSearchBuilder: ... - def execute_local(self) -> GlobalIndexResult: ... def execute(self) -> GlobalIndexResult: ... class FullTextSearchBuilder: @@ -1018,7 +1017,7 @@ class FullTextSearchBuilder: def with_partition_filter(self, predicate: Dict[str, Any]) -> FullTextSearchBuilder: ... def new_full_text_scan(self) -> FullTextScan: ... def new_full_text_read(self) -> FullTextRead: ... - def execute_local(self) -> GlobalIndexResult: ... + def execute(self) -> GlobalIndexResult: ... class FullTextScan: def scan(self) -> FullTextScanPlan: ... diff --git a/bindings/python/src/full_text_search.rs b/bindings/python/src/full_text_search.rs index 8889ee7c8..7f3a441c7 100644 --- a/bindings/python/src/full_text_search.rs +++ b/bindings/python/src/full_text_search.rs @@ -125,7 +125,7 @@ impl PyFullTextSearchBuilder { .map_err(to_py_err)?, }) } - fn execute_local(&self, py: Python<'_>) -> PyResult { + fn execute(&self, py: Python<'_>) -> PyResult { let builder = self.config.builder(&self.table).map_err(to_py_err)?; py.detach(|| runtime().block_on(builder.execute_scored())) .map(PyGlobalIndexResult::new) diff --git a/bindings/python/src/hybrid_search.rs b/bindings/python/src/hybrid_search.rs index b946f5401..42ef92ee3 100644 --- a/bindings/python/src/hybrid_search.rs +++ b/bindings/python/src/hybrid_search.rs @@ -196,10 +196,6 @@ impl PyHybridSearchBuilder { } result.map(PyGlobalIndexResult::new).map_err(to_py_err) } - - fn execute_local(&self, py: Python<'_>) -> PyResult { - self.execute(py) - } } #[cfg(test)] diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index ffb864eb6..3bdfced80 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -167,7 +167,7 @@ impl PyVectorSearchBuilder { let inner = builder.new_read().map_err(to_py_err)?; Ok(PyVectorRead { inner }) } - fn execute_local(&self, py: Python<'_>) -> PyResult { + fn execute(&self, py: Python<'_>) -> PyResult { let mut builder = self.config.single(&self.table).map_err(to_py_err)?; if let Some(query) = &self.query { builder.with_query_vector(query.clone()); @@ -262,7 +262,7 @@ impl PyBatchVectorSearchBuilder { let inner = builder.new_read().map_err(to_py_err)?; Ok(PyBatchVectorRead { inner }) } - fn execute_batch_local(&self, py: Python<'_>) -> PyResult> { + fn execute(&self, py: Python<'_>) -> PyResult> { let mut builder = self.config.batch(&self.table).map_err(to_py_err)?; if let Some(query) = &self.query { builder.with_query_vectors(query.clone()); @@ -642,4 +642,77 @@ rejects(ValueError, sdk.Split.deserialize, b'bad') .unwrap(); }); } + + #[test] + fn search_builders_expose_only_execute_and_keep_empty_and_invalid_query_semantics() { + use paimon::catalog::Identifier; + use paimon::io::FileIOBuilder; + use paimon::spec::{ArrayType, DataType, FloatType, Schema, TableSchema, VarCharType}; + + let schema = Schema::builder() + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .column("body", DataType::VarChar(VarCharType::string_type())) + .option("bucket", "-1") + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .option("global-index.enabled", "true") + .build() + .unwrap(); + let table = Arc::new(Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "execute_contract"), + "memory:/execute_contract".into(), + TableSchema::new(0, &schema), + None, + )); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("sdk", module(py)).unwrap(); + locals + .set_item("table", crate::table::PyTable::new(table)) + .unwrap(); + py.run( + cr#" +for factory, expected in ( + (table.new_vector_search_builder, 'Vector column'), + (table.new_batch_vector_search_builder, 'Vector column'), + (table.new_full_text_search_builder, 'Limit must be positive'), + (table.new_hybrid_search_builder, 'Routes cannot be empty'), +): + builder = factory() + assert callable(builder.execute) + assert not hasattr(builder, 'execute_local') + assert not hasattr(builder, 'execute_batch_local') + try: + builder.execute() + except Exception as error: + assert expected in str(error), str(error) + else: + raise AssertionError('Unconfigured query was accepted') + +single = table.new_vector_search_builder().with_vector_column('embedding').with_query_vector([1., 0.]).with_limit(1) +assert isinstance(single.execute(), sdk.GlobalIndexResult) +assert len(single.execute()) == 0 +batch = table.new_batch_vector_search_builder().with_vector_column('embedding').with_query_vectors([[1., 0.], [0., 1.]]).with_limit(1) +results = batch.execute() +assert len(results) == 2 and all(isinstance(result, sdk.GlobalIndexResult) and len(result) == 0 for result in results) +text = table.new_full_text_search_builder().with_query('body', 'needle').with_limit(1) +assert isinstance(text.execute(), sdk.GlobalIndexResult) and len(text.execute()) == 0 +hybrid = table.new_hybrid_search_builder().add_vector_route('embedding', [1., 0.], 1).add_full_text_route('body', 'needle', 1).with_limit(1) +assert isinstance(hybrid.execute(), sdk.GlobalIndexResult) and len(hybrid.execute()) == 0 +try: + batch.with_query_vectors([]).execute() +except Exception as error: + assert 'Query vectors' in str(error), str(error) +else: + raise AssertionError('Empty batch query was accepted') +"#, + Some(&locals), + Some(&locals), + ).unwrap(); + }); + } } diff --git a/bindings/python/tests/test_vector_search.py b/bindings/python/tests/test_vector_search.py index 321ee2779..23762d68e 100644 --- a/bindings/python/tests/test_vector_search.py +++ b/bindings/python/tests/test_vector_search.py @@ -56,7 +56,7 @@ def test_filters_accumulate_and_partition_uses_table_schema(table): result = (_single(table).with_filter({'method': 'greaterThan', 'field': 'id', 'literals': [1]}) .with_filter({'method': 'lessThan', 'field': 'id', 'literals': [4]}) .with_partition_filter({'method': 'equal', 'field': 'pt', 'literals': [1]}) - .execute_local()) + .execute()) assert isinstance(result, GlobalIndexResult) assert len(result) == 1 assert list(result.row_ids().values()) == [1.0] @@ -91,12 +91,15 @@ def test_scan_read_keeps_snapshot_and_is_reusable(table): assert sorted(ids[:2]) == [1, 3] assert ids[2] == 2 assert rows.column('__paimon_search_score').to_pylist() == pytest.approx([1, 1, 1 / 3]) - assert result.snapshot_id() == plan.snapshot_id() == 1 + assert result.snapshot_id() == plan.snapshot().id() == 1 assert reader.read_plan(plan).row_ids() == result.row_ids() - assert len(builder.execute_local()) == 4 + assert len(builder.execute()) == 4 filtered = _single(table).with_filter({'method': 'equal', 'field': 'id', 'literals': [1]}) - with pytest.raises(Exception, match='plan.*reader|context|different'): - filtered.new_vector_search_read().read_plan(plan) + selected = filtered.new_vector_search_read().read_plan(plan) + assert len(selected) == 1 + assert selected.snapshot_id() == 1 + rows = pa.Table.from_batches(selected.new_read_builder().with_projection(['id']).read()) + assert rows.column('id').to_pylist() == [1] def test_batch_order_empty_results_and_scan_reuse(table): @@ -107,10 +110,10 @@ def test_batch_order_empty_results_and_scan_reuse(table): results = builder.new_batch_vector_search_read().read_batch_plan(plan) assert all(isinstance(r, GlobalIndexResult) for r in results) assert [r.row_ids() for r in results] == [ - _single(table).with_query_vector(q).with_limit(1).execute_local().row_ids() for q in queries] - assert [r.row_ids() for r in builder.execute_batch_local()] == [r.row_ids() for r in results] + _single(table).with_query_vector(q).with_limit(1).execute().row_ids() for q in queries] + assert [r.row_ids() for r in builder.execute()] == [r.row_ids() for r in results] builder.with_filter({'method': 'equal', 'field': 'id', 'literals': [99]}) - assert [len(r) for r in builder.execute_batch_local()] == [0, 0, 0] + assert [len(r) for r in builder.execute()] == [0, 0, 0] @pytest.mark.parametrize('query', [None, []]) @@ -121,16 +124,17 @@ def test_invalid_query_is_reported_before_planning(table, query): with pytest.raises(Exception, match='vector|Vector'): builder.new_vector_search_read() with pytest.raises(Exception, match='Query vectors'): - table.new_batch_vector_search_builder().with_vector_column('embedding').with_limit(1).with_query_vectors([]).execute_batch_local() + (table.new_batch_vector_search_builder().with_vector_column('embedding') + .with_limit(1).with_query_vectors([]).execute()) def test_single_vector_keyword_matches_stub(table): builder = _single(table).with_query_vector(vector=[1, 0]).with_limit(1) - assert builder.execute_local().row_ids() == _single(table).with_limit(1).execute_local().row_ids() + assert builder.execute().row_ids() == _single(table).with_limit(1).execute().row_ids() def test_external_result_uses_same_scan_type_without_search_context(table): - found = _single(table).with_limit(1).execute_local() + found = _single(table).with_limit(1).execute() external = GlobalIndexResult.from_scored_row_ids(found.row_ids()) builder = table.new_read_builder().with_projection(['id']) plan = builder.new_scan().with_global_index_result(external).plan() From 7ca55815d0426f5b43e8c011eadf529da300061c Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 23:22:12 +0800 Subject: [PATCH 5/6] Validate search source payloads and scope hybrid test bridge --- .../src/table/full_text_search_builder.rs | 142 +++++++++++++++++- .../table/full_text_search_builder/read.rs | 14 ++ .../paimon/src/table/hybrid_search_builder.rs | 2 +- crates/paimon/src/table/search_scan_split.rs | 99 ++++++++++++ 4 files changed, 253 insertions(+), 4 deletions(-) diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 5fc928a69..e5e81d6ac 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -1324,14 +1324,18 @@ mod tests { writer: &mut FullTextIndexWriter, query: &str, ) -> paimon_ftindex_core::FullTextSearchResult { - let mut bytes = Vec::new(); - writer.write(&mut PosWriter::new(&mut bytes)).unwrap(); - FullTextIndexReader::open(SliceReader::new(bytes)) + FullTextIndexReader::open(SliceReader::new(write_test_index(writer))) .unwrap() .search(query, 10) .unwrap() } + fn write_test_index(writer: &mut FullTextIndexWriter) -> Vec { + let mut bytes = Vec::new(); + writer.write(&mut PosWriter::new(&mut bytes)).unwrap(); + bytes + } + #[test] fn test_plan_full_text_index_searches_partitions_once_and_preserves_order() { let entries = [ @@ -1530,6 +1534,138 @@ mod tests { assert!(result.positions().unwrap().is_empty()); } + fn pk_full_text_indexed_plan(table: &Table) -> FullTextScanPlan { + let snapshot = crate::spec::Snapshot::builder() + .version(3) + .id(7) + .schema_id(0) + .base_manifest_list(String::new()) + .delta_manifest_list(String::new()) + .commit_user("capture".into()) + .commit_identifier(1) + .commit_kind(crate::spec::CommitKind::APPEND) + .time_millis(0) + .build(); + let data_file = crate::spec::DataFileMeta { + file_name: "d0".into(), + file_size: 1, + row_count: 2, + min_key: Vec::new(), + max_key: Vec::new(), + key_stats: crate::spec::stats::BinaryTableStats::empty(), + value_stats: crate::spec::stats::BinaryTableStats::empty(), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 0, + level: 1, + extra_files: Vec::new(), + creation_time: None, + delete_row_count: None, + embedded_index: None, + file_source: Some(1), + value_stats_cols: None, + external_path: None, + first_row_id: None, + write_cols: None, + column_max_sequence_numbers: None, + }; + let mut payload = full_text_index_entry("pk-ft", 0, 1).index_file; + // Java source metadata: version, level, file count, UTF name, row count. + payload.global_index_meta.as_mut().unwrap().source_meta = Some( + [ + 1i32.to_be_bytes().as_slice(), + 1i32.to_be_bytes().as_slice(), + 1i32.to_be_bytes().as_slice(), + 2u16.to_be_bytes().as_slice(), + b"d0", + 2i64.to_be_bytes().as_slice(), + ] + .concat(), + ); + let split = crate::table::SearchScanSplit::PrimaryKey { + data_split: crate::table::DataSplitBuilder::new() + .with_snapshot(7) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{}/bucket-0", table.location())) + .with_data_files(vec![data_file]) + .build() + .unwrap(), + payloads: vec![payload], + uncovered_data_files: vec![], + row_ranges_by_file: Default::default(), + }; + FullTextScanPlan::new(Some(snapshot), vec![split]).unwrap() + } + + #[tokio::test] + async fn pk_full_text_source_plan_rejects_incompatible_payloads_before_io() { + let table = pk_full_text_table("pk_ft_wrong_source", &[]); + let mut builder = table.new_full_text_search_builder(); + builder + .with_query("body", r#"{"match":{"query":"alpha"}}"#) + .with_limit(10); + let reader = builder.new_read().unwrap(); + let source = pk_full_text_indexed_plan(&table); + for invalid in ["field", "type", "missing-meta"] { + let mut splits = source.splits().to_vec(); + let crate::table::SearchScanSplit::PrimaryKey { payloads, .. } = &mut splits[0] else { + unreachable!() + }; + match invalid { + "field" => { + payloads[0] + .global_index_meta + .as_mut() + .unwrap() + .index_field_id = 0 + } + "type" => payloads[0].index_type = "lucene-vector".into(), + _ => payloads[0].global_index_meta = None, + } + let plan = FullTextScanPlan::new(source.snapshot().cloned(), splits).unwrap(); + let error = reader.read(plan).await.unwrap_err(); + assert!( + matches!(error, crate::Error::DataInvalid { ref message, .. } + if message.contains("does not match the reader's field and index type")), + "{invalid}: {error:?}" + ); + } + } + + #[tokio::test] + async fn pk_full_text_source_plan_reads_matching_payload_from_captured_snapshot() { + let table = pk_full_text_table("pk_ft_matching_source", &[]); + let mut writer = FullTextIndexWriter::new(FullTextIndexConfig::new()).unwrap(); + writer.add_document(0, "alpha").unwrap(); + writer.add_document(1, "beta").unwrap(); + table + .file_io() + .new_output(&format!("{}/index/pk-ft", table.location())) + .unwrap() + .write(Bytes::from(write_test_index(&mut writer))) + .await + .unwrap(); + let plan = pk_full_text_indexed_plan(&table); + let plan: FullTextScanPlan = + serde_json::from_slice(&serde_json::to_vec(&plan).unwrap()).unwrap(); + let result = table + .new_full_text_search_builder() + .with_query("body", r#"{"match":{"query":"alpha"}}"#) + .with_limit(10) + .new_read() + .unwrap() + .read(plan) + .await + .unwrap(); + assert_eq!(result.snapshot_id(), Some(7)); + let positions = result.positions().unwrap(); + assert_eq!(positions.len(), 1); + assert_eq!(positions[0].data_file_name, "d0"); + assert_eq!(positions[0].row_position, 0); + assert!(positions[0].score > 0.0); + } + // (d) The PK full-text path produces physical positions, not global row-ids: // Only the ranges-only `execute` API rejects the physical-position result. #[tokio::test] diff --git a/crates/paimon/src/table/full_text_search_builder/read.rs b/crates/paimon/src/table/full_text_search_builder/read.rs index 54f46ddb6..2cd518134 100644 --- a/crates/paimon/src/table/full_text_search_builder/read.rs +++ b/crates/paimon/src/table/full_text_search_builder/read.rs @@ -152,6 +152,20 @@ impl FullTextRead { source: None, }); } + if payloads.iter().any(|file| { + file.index_type != FULL_TEXT_INDEX_TYPE + || file + .global_index_meta + .as_ref() + .map(|meta| meta.index_field_id) + != self.context.primary_key_field_id + }) { + return Err(crate::Error::DataInvalid { + message: "Primary-key full-text source does not match the reader's field and index type" + .into(), + source: None, + }); + } let data_split = if data_split.snapshot_id() == 0 { let mut builder = crate::table::DataSplitBuilder::new() .with_snapshot(plan.snapshot_id().unwrap_or(0)) diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index 9a1c30610..6638e719e 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -321,7 +321,7 @@ impl<'a> HybridSearchBuilder<'a> { /// Schedule each route's shard work while preserving one shared snapshot /// and the configured core fusion ranker. - #[cfg(any(feature = "internal-api", test))] + #[cfg(any(feature = "internal-api", all(test, feature = "fulltext")))] pub(crate) async fn execute_with_executor( &self, executor: &dyn super::SearchTaskExecutor, diff --git a/crates/paimon/src/table/search_scan_split.rs b/crates/paimon/src/table/search_scan_split.rs index cf27c7cab..438a6e241 100644 --- a/crates/paimon/src/table/search_scan_split.rs +++ b/crates/paimon/src/table/search_scan_split.rs @@ -280,6 +280,12 @@ pub(super) fn global_entries_from_splits( } }; for file in scalar { + if normalize_queryable_global_index_type(&file.index_type).is_none() { + return Err(invalid_source(format!( + "Scalar search source '{}' has unsupported index type '{}'", + file.file_name, file.index_type + ))); + } let meta = file.global_index_meta.as_ref().ok_or_else(|| { invalid_source(format!( "Scalar search index '{}' has no global metadata", @@ -396,3 +402,96 @@ pub(super) async fn partition_row_ids_for_search_sources( } Ok(selected) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{CommitKind, GlobalIndexMeta, Snapshot}; + use crate::table::VectorScanPlan; + + fn snapshot() -> Snapshot { + Snapshot::builder() + .version(3) + .id(1) + .schema_id(0) + .base_manifest_list("base".into()) + .delta_manifest_list("delta".into()) + .commit_user("test".into()) + .commit_identifier(1) + .commit_kind(CommitKind::APPEND) + .time_millis(0) + .build() + } + + fn raw_source_with_scalar(index_type: &str) -> SearchScanSplit { + SearchScanSplit::Raw { + row_ranges: vec![RowRange::new(1, 1)], + scalar_index_files: vec![IndexFileMeta { + index_type: index_type.into(), + file_name: "scalar-index".into(), + file_size: 1, + row_count: 3, + deletion_vectors_ranges: None, + external_path: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: 0, + row_range_end: 2, + index_field_id: 1, + extra_field_ids: None, + index_meta: None, + source_meta: None, + }), + }], + index_type: None, + } + } + + #[test] + fn raw_source_scalar_metadata_cannot_add_search_payloads() { + // A search payload hidden in scalar metadata would otherwise be searched + // over [0, 2], extending the explicitly selected raw source [1, 1]. + for index_type in ["ivf-flat", "lumina", "lumina-vector-ann", "full-text"] { + let source = raw_source_with_scalar(index_type); + let error = VectorScanPlan::new(Some(snapshot()), vec![source.clone()]).unwrap_err(); + assert!(error.to_string().contains("unsupported index type")); + let wire = serde_json::json!({ + "snapshot": snapshot(), + "splits": [source.clone()], + }); + assert!(serde_json::from_value::(wire.clone()).is_err()); + #[cfg(feature = "fulltext")] + { + assert!( + crate::table::FullTextScanPlan::new(Some(snapshot()), vec![source]).is_err() + ); + assert!(serde_json::from_value::(wire).is_err()); + } + } + } + + #[test] + fn raw_source_accepts_queryable_scalar_metadata() { + for index_type in [ + "bitmap", + "Bitmap", + "btree", + "BTREE", + "multivalue", + "MultiValue", + "fm", + "FM", + ] { + let source = raw_source_with_scalar(index_type); + let plan = VectorScanPlan::new(Some(snapshot()), vec![source.clone()]).unwrap(); + let wire = serde_json::to_value(plan).unwrap(); + assert!(serde_json::from_value::(wire.clone()).is_ok()); + #[cfg(feature = "fulltext")] + { + assert!( + crate::table::FullTextScanPlan::new(Some(snapshot()), vec![source]).is_ok() + ); + assert!(serde_json::from_value::(wire).is_ok()); + } + } + } +} From cc9b330646edbb3882fdb0d08b732480b2df25f3 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 09:08:15 +0800 Subject: [PATCH 6/6] Fix CI dependency advisories and search plan test contracts --- Cargo.lock | 8 ++--- DEPENDENCIES.rust.tsv | 4 +-- benchmarks/tpcds/DEPENDENCIES.rust.tsv | 4 +-- bindings/c/DEPENDENCIES.rust.tsv | 4 +-- bindings/go/DEPENDENCIES.rust.tsv | 4 +-- bindings/python/DEPENDENCIES.rust.tsv | 4 +-- .../integration_tests/DEPENDENCIES.rust.tsv | 4 +-- .../datafusion/DEPENDENCIES.rust.tsv | 4 +-- .../datafusion/tests/read_tables.rs | 35 +++++++++++++++++-- .../paimon-rest-server/DEPENDENCIES.rust.tsv | 4 +-- crates/paimon/DEPENDENCIES.rust.tsv | 4 +-- .../table/global_index_result/search_tests.rs | 17 +++++---- 12 files changed, 63 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d73fbb31d..779c94e79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6115,9 +6115,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "aws-lc-rs", "ring", @@ -8170,9 +8170,9 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" dependencies = [ "rustls-pki-types", ] diff --git a/DEPENDENCIES.rust.tsv b/DEPENDENCIES.rust.tsv index 2313f8f4c..3a7378bf6 100644 --- a/DEPENDENCIES.rust.tsv +++ b/DEPENDENCIES.rust.tsv @@ -509,7 +509,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -681,7 +681,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X whoami@2.1.3 X X X wide@0.7.33 X X X winapi@0.3.9 X X diff --git a/benchmarks/tpcds/DEPENDENCIES.rust.tsv b/benchmarks/tpcds/DEPENDENCIES.rust.tsv index 24915be44..a7bb34556 100644 --- a/benchmarks/tpcds/DEPENDENCIES.rust.tsv +++ b/benchmarks/tpcds/DEPENDENCIES.rust.tsv @@ -333,7 +333,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -433,7 +433,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/bindings/c/DEPENDENCIES.rust.tsv b/bindings/c/DEPENDENCIES.rust.tsv index 9dbbc6e93..82566005e 100644 --- a/bindings/c/DEPENDENCIES.rust.tsv +++ b/bindings/c/DEPENDENCIES.rust.tsv @@ -294,7 +294,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -393,7 +393,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/bindings/go/DEPENDENCIES.rust.tsv b/bindings/go/DEPENDENCIES.rust.tsv index 9dbbc6e93..82566005e 100644 --- a/bindings/go/DEPENDENCIES.rust.tsv +++ b/bindings/go/DEPENDENCIES.rust.tsv @@ -294,7 +294,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -393,7 +393,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/bindings/python/DEPENDENCIES.rust.tsv b/bindings/python/DEPENDENCIES.rust.tsv index 563899bf6..fd4ef2cc8 100644 --- a/bindings/python/DEPENDENCIES.rust.tsv +++ b/bindings/python/DEPENDENCIES.rust.tsv @@ -433,7 +433,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -565,7 +565,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X whoami@2.1.3 X X X wide@0.7.33 X X X winapi@0.3.9 X X diff --git a/crates/integration_tests/DEPENDENCIES.rust.tsv b/crates/integration_tests/DEPENDENCIES.rust.tsv index 4d16dab49..6f7b6affb 100644 --- a/crates/integration_tests/DEPENDENCIES.rust.tsv +++ b/crates/integration_tests/DEPENDENCIES.rust.tsv @@ -264,7 +264,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -357,7 +357,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/crates/integrations/datafusion/DEPENDENCIES.rust.tsv b/crates/integrations/datafusion/DEPENDENCIES.rust.tsv index 091e6a231..787cf0b7a 100644 --- a/crates/integrations/datafusion/DEPENDENCIES.rust.tsv +++ b/crates/integrations/datafusion/DEPENDENCIES.rust.tsv @@ -428,7 +428,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -582,7 +582,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi@0.3.9 X X winapi-i686-pc-windows-gnu@0.4.0 X X diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index caa00fa9d..d18311ca3 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -2290,9 +2290,25 @@ mod vector_search_tests { #[tokio::test] async fn test_vector_search_without_matching_index_returns_empty() { - let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await; + let (ctx, catalog, _tmp) = create_lumina_vector_search_context().await; + let table = catalog + .get_table(&Identifier::new("default", "test_lumina_vector")) + .await + .expect("load table"); + let dropped = table + .new_global_index_drop_builder() + .with_index_column("embedding") + .with_index_type("lumina") + .execute() + .await + .expect("drop vector indexes"); + assert!( + dropped > 0, + "the fixture must have a vector index to remove" + ); + let batches = ctx - .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'missing_embedding', '[1.0]', 10)") + .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 10)") .await .expect("SQL should parse") .collect() @@ -2306,6 +2322,21 @@ mod vector_search_tests { ); } + #[tokio::test] + async fn test_vector_search_rejects_missing_column() { + let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await; + let error = ctx + .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'missing_embedding', '[1.0]', 10)") + .await + .expect("SQL should parse") + .collect() + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("Search column 'missing_embedding' does not exist")); + } + #[tokio::test] async fn test_vector_search_java_vindex_table() { let (ctx, _catalog, _tmp) = create_java_vindex_vector_search_context().await; diff --git a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv index 1ef833b0a..eadd3ccb0 100644 --- a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv +++ b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv @@ -268,7 +268,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -363,7 +363,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/crates/paimon/DEPENDENCIES.rust.tsv b/crates/paimon/DEPENDENCIES.rust.tsv index cceabf938..52f38e7cd 100644 --- a/crates/paimon/DEPENDENCIES.rust.tsv +++ b/crates/paimon/DEPENDENCIES.rust.tsv @@ -427,7 +427,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -585,7 +585,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X whoami@2.1.3 X X X wide@0.7.33 X X X winapi@0.3.9 X X diff --git a/crates/paimon/src/table/global_index_result/search_tests.rs b/crates/paimon/src/table/global_index_result/search_tests.rs index c38eb8008..05e99193d 100644 --- a/crates/paimon/src/table/global_index_result/search_tests.rs +++ b/crates/paimon/src/table/global_index_result/search_tests.rs @@ -176,7 +176,7 @@ async fn public_plan_can_be_reused_by_single_and_batch_readers_after_builder_dro } #[tokio::test] -async fn public_reader_rejects_a_plan_for_another_table_or_filter() { +async fn public_reader_applies_its_filter_and_requires_available_sources() { let table = de_vector_table().await; let mut builder = table.new_vector_search_builder(); builder @@ -185,15 +185,14 @@ async fn public_reader_rejects_a_plan_for_another_table_or_filter() { .with_limit(1); let plan = builder.new_scan().unwrap().plan().await.unwrap(); builder.with_filter(id_gt_filter(&table, 1)); - let error = builder + let result = builder .new_read() .unwrap() .read(plan.clone()) .await - .unwrap_err(); - assert!(error - .to_string() - .contains("same table, column and pre-filter")); + .unwrap(); + assert_eq!(result.row_ids().unwrap().row_ids, vec![2]); + assert_eq!(result.snapshot_id(), plan.snapshot_id()); let other = crate::table::vector_search_test_utils::vector_test_table_at( "memory:/different_plan_table", ); @@ -205,7 +204,7 @@ async fn public_reader_rejects_a_plan_for_another_table_or_filter() { .new_read() .unwrap(); let error = read.read(plan).await.unwrap_err(); - assert!(error - .to_string() - .contains("same table, column and pre-filter")); + let message = error.to_string(); + assert!(message.contains("NotFound"), "{message}"); + assert!(message.contains("different_plan_table/index/"), "{message}"); }