diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index 6eb989792..8a8d535fa 100644 --- a/bindings/c/src/tests.rs +++ b/bindings/c/src/tests.rs @@ -3118,7 +3118,7 @@ fn test_two_commits_same_builder() { // // Two storage shapes are exercised end-to-end through the C `execute_read` // terminal, each compared against an independent core Rust -// `SearchResultReadBuilder::read()` reference: +// `GlobalIndexResultReadBuilder::read()` reference: // // * A primary-key vector table backed by a real vindex IVF-flat ANN segment // built in-process (bucket-local ANN search, residual filter supported). diff --git a/bindings/c/src/vector_search.rs b/bindings/c/src/vector_search.rs index 31c362f05..f79d01d99 100644 --- a/bindings/c/src/vector_search.rs +++ b/bindings/c/src/vector_search.rs @@ -31,8 +31,8 @@ use std::collections::HashMap; use std::ffi::{c_char, c_void}; use paimon::spec::Predicate; +use paimon::table::GlobalIndexResult; use paimon::table::{ArrowRecordBatchStream, Table, VectorSearchBuilder}; -use paimon::vector_search::SearchResult; use crate::error::{check_non_null, paimon_error, validate_cstr, PaimonErrorCode}; use crate::result::{ @@ -320,7 +320,7 @@ pub unsafe extern "C" fn paimon_vector_search_builder_execute_read( } pub(crate) async fn materialize_search_result( - result: SearchResult, + result: GlobalIndexResult, projection: Option<&[String]>, ) -> paimon::Result { let mut reader = result.new_read_builder(); diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 1edff10eb..5b132e812 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -56,8 +56,8 @@ class Split: ... def row_count(self) -> int: ... def is_streaming(self) -> bool: ... - # Java SplitSerializer v1 binary: DataSplit v8/v9, or score-free - # IndexedSplit when the native split has row ranges. + # Java SplitSerializer v1 binary: DataSplit v8/v9, or IndexedSplit with + # selected row ranges and optional aligned scores. def serialize(self) -> bytes: ... class Plan: @@ -67,7 +67,32 @@ class Plan: def splits(self) -> List[Split]: ... def __len__(self) -> int: ... +class GlobalIndexResult: + """Index selection shared by table searches and scans; PK selections retain their snapshot.""" + @staticmethod + def from_splits(snapshot_id: Optional[int], splits: Iterable[Split]) -> "GlobalIndexResult": ... + @staticmethod + def from_row_ranges(ranges: List[Tuple[int, int]]) -> "GlobalIndexResult": ... + @staticmethod + def from_scored_row_ids(scores: Dict[int, float]) -> "GlobalIndexResult": ... + + def snapshot_id(self) -> Optional[int]: ... + def __len__(self) -> int: ... + def new_read_builder(self) -> "GlobalIndexResultReadBuilder": ... + def row_ids(self) -> Dict[int, float]: + """DE scores keyed by global row ID; rejects PK results, including empty ones.""" + ... + def positions(self) -> List[Tuple[bytes, int, str, int, float]]: + """PK partition bytes, bucket, file, position and score in relevance order.""" + ... + def splits(self) -> List[bytes]: + """PK physical selections in Java SplitSerializer v1 format, including scores.""" + ... + class TableScan: + def with_global_index_result(self, result: GlobalIndexResult) -> "TableScan": + """Reuse an index selection. Authorized scans ignore supplied index selections.""" + ... def with_row_position_slice(self, start: int, end: int) -> "TableScan": ... def with_row_position_shard(self, index: int, count: int) -> "TableScan": ... def with_chunk_shuffle(self, seed: str, chunk_size: int) -> "TableScan": @@ -731,7 +756,7 @@ class VectorSearchBuilder: ... def new_vector_search_scan(self) -> "VectorScan": ... def new_vector_search_read(self) -> "VectorRead": ... - def execute_local(self) -> "SearchResult": ... + def execute_local(self) -> "GlobalIndexResult": ... class BatchVectorSearchBuilder: def with_vector_column(self, name: str) -> "BatchVectorSearchBuilder": ... @@ -743,7 +768,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["SearchResult"]: ... + def execute_batch_local(self) -> List["GlobalIndexResult"]: ... class VectorScan: def scan(self) -> "VectorScanPlan": ... @@ -752,29 +777,13 @@ class VectorScanPlan: def snapshot_id(self) -> Optional[int]: ... class VectorRead: - def read_plan(self, plan: VectorScanPlan) -> "SearchResult": ... + def read_plan(self, plan: VectorScanPlan) -> "GlobalIndexResult": ... class BatchVectorRead: - def read_batch_plan(self, plan: VectorScanPlan) -> List["SearchResult"]: ... - -class SearchResult: - """Snapshot-scoped global row IDs or primary-key physical positions.""" - def snapshot_id(self) -> Optional[int]: ... - def __len__(self) -> int: ... - def new_read_builder(self) -> "SearchResultReadBuilder": ... - def row_ids(self) -> Dict[int, float]: - """DE scores keyed by global row ID; rejects PK results, including empty ones.""" - ... - def positions(self) -> List[Tuple[bytes, int, str, int, float]]: - """PK partition bytes, bucket, file, position and score in relevance order.""" - ... - def splits(self) -> List[bytes]: - """PK physical selections in Java SplitSerializer v1 format, including scores.""" - ... - + def read_batch_plan(self, plan: VectorScanPlan) -> List["GlobalIndexResult"]: ... -class SearchResultReadBuilder: - def with_projection(self, columns: List[str]) -> "SearchResultReadBuilder": ... +class GlobalIndexResultReadBuilder: + def with_projection(self, columns: List[str]) -> "GlobalIndexResultReadBuilder": ... def read(self) -> List[pyarrow.RecordBatch]: """Read user columns and scores in relevance order from the search snapshot.""" ... @@ -793,11 +802,7 @@ 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) -> HybridSearchResult: ... - -class HybridSearchResult: - def __len__(self) -> int: ... - def row_ids(self) -> Dict[int, float]: ... + def execute_local(self) -> GlobalIndexResult: ... class FullTextSearchBuilder: def with_query(self, field_name: str, query: str) -> FullTextSearchBuilder: ... @@ -806,7 +811,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) -> SearchResult: ... + def execute_local(self) -> GlobalIndexResult: ... class FullTextScan: def scan(self) -> FullTextScanPlan: ... @@ -815,7 +820,7 @@ class FullTextScanPlan: def snapshot_id(self) -> Optional[int]: ... class FullTextRead: - def read_plan(self, plan: FullTextScanPlan) -> SearchResult: ... + def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult: ... class ShardTableUpdator: def arrow_reader(self) -> RecordBatchReader: diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 1c29b4564..d312192a4 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -802,15 +802,14 @@ 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::()?; + 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 ca364d1a9..a3fc6a9bd 100644 --- a/bindings/python/src/full_text_search.rs +++ b/bindings/python/src/full_text_search.rs @@ -18,8 +18,8 @@ //! Thin bindings over core full-text Scan -> Plan -> Read. use crate::error::to_py_err; +use crate::global_index_result::PyGlobalIndexResult; use crate::predicate::dict_to_table_predicate; -use crate::search_result::PySearchResult; use paimon::spec::Predicate; use paimon::table::{FullTextRead, FullTextScan, FullTextScanPlan, Table}; use paimon_datafusion::runtime::runtime; @@ -123,10 +123,10 @@ impl PyFullTextSearchBuilder { .map_err(to_py_err)?, }) } - fn execute_local(&self, py: Python<'_>) -> PyResult { + fn execute_local(&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(PySearchResult::new) + .map(PyGlobalIndexResult::new) .map_err(to_py_err) } } @@ -161,9 +161,13 @@ pub struct PyFullTextRead { } #[pymethods] impl PyFullTextRead { - fn read_plan(&self, py: Python<'_>, plan: &PyFullTextScanPlan) -> PyResult { + fn read_plan( + &self, + py: Python<'_>, + plan: &PyFullTextScanPlan, + ) -> PyResult { py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - .map(PySearchResult::new) + .map(PyGlobalIndexResult::new) .map_err(to_py_err) } } diff --git a/bindings/python/src/search_result.rs b/bindings/python/src/global_index_result.rs similarity index 66% rename from bindings/python/src/search_result.rs rename to bindings/python/src/global_index_result.rs index 2e889e30f..df8f32185 100644 --- a/bindings/python/src/search_result.rs +++ b/bindings/python/src/global_index_result.rs @@ -16,21 +16,24 @@ // under the License. use crate::error::to_py_err; +use crate::read::extract_splits; use arrow::pyarrow::ToPyArrow; use futures::TryStreamExt; -use paimon::table::SearchResult; +use paimon::table::GlobalIndexResult; +use paimon::table::RowRange; use paimon_datafusion::runtime::runtime; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::PyBytes; use std::collections::HashMap; -#[pyclass(name = "SearchResult", module = "pypaimon_rust.datafusion")] -pub struct PySearchResult { - inner: SearchResult, +#[pyclass(name = "GlobalIndexResult", module = "pypaimon_rust.datafusion")] +pub struct PyGlobalIndexResult { + pub(crate) inner: GlobalIndexResult, } -impl PySearchResult { - pub(crate) fn new(inner: SearchResult) -> Self { +impl PyGlobalIndexResult { + pub(crate) fn new(inner: GlobalIndexResult) -> Self { Self { inner } } } @@ -38,12 +41,44 @@ impl PySearchResult { type PySearchPosition = (Py, i32, String, i64, f32); #[pymethods] -impl PySearchResult { +impl PyGlobalIndexResult { + #[staticmethod] + fn from_splits(snapshot_id: Option, splits: &Bound<'_, PyAny>) -> PyResult { + Ok(Self { + inner: GlobalIndexResult::from_splits(snapshot_id, extract_splits(splits)?), + }) + } + + #[staticmethod] + fn from_row_ranges(ranges: Vec<(i64, i64)>) -> PyResult { + let ranges = ranges + .into_iter() + .map(|(from, to)| { + if from > to { + Err(PyValueError::new_err("row range start must be <= end")) + } else { + Ok(RowRange::new(from, to)) + } + }) + .collect::>>()?; + Ok(Self { + inner: GlobalIndexResult::from_row_ranges(ranges), + }) + } + + #[staticmethod] + fn from_scored_row_ids(scores: HashMap) -> Self { + Self { + inner: GlobalIndexResult::from_scored_row_ids( + paimon::vector_search::ScoredRowIds::from_scored_map(scores), + ), + } + } fn snapshot_id(&self) -> Option { self.inner.snapshot_id() } - fn new_read_builder(&self) -> PySearchResultReadBuilder { - PySearchResultReadBuilder { + fn new_read_builder(&self) -> PyGlobalIndexResultReadBuilder { + PyGlobalIndexResultReadBuilder { result: self.inner.clone(), projection: None, } @@ -88,14 +123,17 @@ impl PySearchResult { } } -#[pyclass(name = "SearchResultReadBuilder", module = "pypaimon_rust.datafusion")] -pub struct PySearchResultReadBuilder { - result: SearchResult, +#[pyclass( + name = "GlobalIndexResultReadBuilder", + module = "pypaimon_rust.datafusion" +)] +pub struct PyGlobalIndexResultReadBuilder { + result: GlobalIndexResult, projection: Option>, } #[pymethods] -impl PySearchResultReadBuilder { +impl PyGlobalIndexResultReadBuilder { fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec) -> PyRefMut<'_, Self> { slf.projection = Some(columns); slf diff --git a/bindings/python/src/hybrid_search.rs b/bindings/python/src/hybrid_search.rs index 749bdd93f..0370b585a 100644 --- a/bindings/python/src/hybrid_search.rs +++ b/bindings/python/src/hybrid_search.rs @@ -18,13 +18,14 @@ //! Hybrid configuration forwarding; snapshot, routes and fusion belong to core. use crate::error::to_py_err; +use crate::global_index_result::PyGlobalIndexResult; use crate::predicate::dict_to_table_predicate; use paimon::spec::{Predicate, Snapshot}; use paimon::table::{HybridSearchRanker, HybridSearchRoute, Table}; -use paimon::vector_search::ScoredRowIds; use paimon_datafusion::runtime::runtime; use pyo3::prelude::*; use pyo3::types::PyDict; + use std::collections::HashMap; use std::sync::Arc; @@ -176,31 +177,10 @@ impl PyHybridSearchBuilder { Ok(slf) } - fn execute_local(&self, py: Python<'_>) -> PyResult { + fn execute_local(&self, py: Python<'_>) -> PyResult { let builder = self.builder(); py.detach(|| runtime().block_on(builder.execute_scored())) - .map(|inner| PyHybridSearchResult { inner }) + .map(PyGlobalIndexResult::new) .map_err(to_py_err) } } - -#[pyclass(name = "HybridSearchResult", module = "pypaimon_rust.datafusion")] -pub struct PyHybridSearchResult { - inner: ScoredRowIds, -} - -#[pymethods] -impl PyHybridSearchResult { - fn __len__(&self) -> usize { - self.inner.len() - } - - fn row_ids(&self) -> HashMap { - self.inner - .row_ids - .iter() - .copied() - .zip(self.inner.scores.iter().copied()) - .collect() - } -} diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index c55aa5bdf..4ae51342b 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -24,6 +24,7 @@ mod context; mod error; mod evolution; mod full_text_search; +mod global_index_result; mod hybrid_search; mod index_build; mod merge_into; @@ -31,7 +32,6 @@ mod oss_cpp; mod predicate; mod read; mod schema; -mod search_result; mod shard_update; mod table; mod udf; diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index fc2ae2fc5..8b7c167ba 100644 --- a/bindings/python/src/read.rs +++ b/bindings/python/src/read.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::global_index_result::PyGlobalIndexResult; use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; @@ -26,7 +27,8 @@ use futures::stream::BoxStream; use futures::TryStreamExt; use paimon::spec::{is_variant_extraction_row, DataField, DataType, Predicate, RowType}; use paimon::table::{ - ArrowRecordBatchStream, DataSplit, IncrementalScanMode, RowRange, StreamTableScan, Table, + ArrowRecordBatchStream, DataSplit, GlobalIndexResult, IncrementalScanMode, RowRange, + StreamTableScan, Table, }; use paimon_datafusion::runtime::runtime; use pyo3::exceptions::{PyBaseException, PyRuntimeError, PyTypeError, PyValueError}; @@ -105,7 +107,7 @@ fn apply_read_config( /// any iterable (list/tuple/generator). Runs under the GIL since it touches /// Python objects. A non-iterable argument or a non-`Split` element raises /// `TypeError`. -fn extract_splits(splits: &Bound<'_, PyAny>) -> PyResult> { +pub(crate) fn extract_splits(splits: &Bound<'_, PyAny>) -> PyResult> { let iter = splits .try_iter() .map_err(|_| PyTypeError::new_err("read() expects a sequence of Split objects"))?; @@ -344,6 +346,7 @@ impl PyReadBuilder { row_ranges: self.row_ranges.clone(), case_sensitive: self.case_sensitive, incremental_scan: None, + global_index_result: None, row_position_slice: None, row_position_shard: None, chunk_shuffle: None, @@ -407,6 +410,7 @@ pub struct PyTableScan { row_ranges: Option>, case_sensitive: bool, incremental_scan: Option, + global_index_result: Option, row_position_slice: Option<(u64, u64)>, row_position_shard: Option<(u64, u64)>, chunk_shuffle: Option, @@ -443,6 +447,11 @@ fn parse_incremental_scan_mode(mode: &str) -> PyResult { impl PyTableScan { fn core_scan(&self) -> PyResult> { let mut scan = self.read_builder()?.new_scan(); + if let Some(result) = &self.global_index_result { + scan = scan + .with_global_index_result(result.clone()) + .map_err(to_py_err)?; + } if let Some((start, end)) = self.row_position_slice { scan = scan .with_row_position_slice(start, end) @@ -471,6 +480,11 @@ impl PyTableScan { mode: IncrementalScanMode, ) -> PyResult> { let mut scan = self.read_builder()?.new_incremental_scan(mode, start, end); + if let Some(result) = &self.global_index_result { + scan = scan + .with_global_index_result(result.clone()) + .map_err(to_py_err)?; + } if let Some((start, end)) = self.row_position_slice { scan = scan .with_row_position_slice(start, end) @@ -645,6 +659,29 @@ impl PyStreamTableScan { #[pymethods] impl PyTableScan { + /// Supply an evaluated selection; validation and authorization stay in core. + fn with_global_index_result<'py>( + mut slf: PyRefMut<'py, Self>, + result: PyRef<'py, PyGlobalIndexResult>, + ) -> PyResult> { + let result = result.inner.clone(); + if let Some(incremental) = slf.incremental_scan { + slf.core_incremental_scan( + incremental.start_snapshot_id, + incremental.end_snapshot_id, + incremental.mode, + )? + .with_global_index_result(result.clone()) + .map_err(to_py_err)?; + } else { + slf.core_scan()? + .with_global_index_result(result.clone()) + .map_err(to_py_err)?; + } + slf.global_index_result = Some(result); + Ok(slf) + } + /// Select a half-open range of append-table row positions. fn with_row_position_slice( mut slf: PyRefMut<'_, Self>, @@ -972,7 +1009,7 @@ impl PySplit { /// /// This is deliberately separate from the constructor used by pickle: /// pickle bytes are an opaque Rust JSON encoding, while this method accepts - /// Java/Python-compatible DataSplit and score-free IndexedSplit frames. + /// Java/Python-compatible DataSplit and IndexedSplit frames, including scores. #[staticmethod] fn deserialize(state: &Bound<'_, PyBytes>) -> PyResult { Ok(Self { diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index 540c579d1..4c1173f22 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -25,8 +25,8 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use crate::error::to_py_err; +use crate::global_index_result::PyGlobalIndexResult; use crate::predicate::dict_to_table_predicate; -use crate::search_result::PySearchResult; #[derive(Clone, Default)] struct SearchConfig { @@ -165,7 +165,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_local(&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()); @@ -173,7 +173,7 @@ impl PyVectorSearchBuilder { let results = py .detach(|| runtime().block_on(builder.execute())) .map_err(to_py_err)?; - Ok(PySearchResult::new(results)) + Ok(PyGlobalIndexResult::new(results)) } } @@ -260,7 +260,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_batch_local(&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()); @@ -268,7 +268,7 @@ impl PyBatchVectorSearchBuilder { let results = py .detach(|| runtime().block_on(builder.execute())) .map_err(to_py_err)?; - Ok(results.into_iter().map(PySearchResult::new).collect()) + Ok(results.into_iter().map(PyGlobalIndexResult::new).collect()) } } @@ -306,9 +306,9 @@ pub struct PyVectorRead { #[pymethods] impl PyVectorRead { - fn read_plan(&self, py: Python<'_>, plan: &PyVectorScanPlan) -> PyResult { + fn read_plan(&self, py: Python<'_>, plan: &PyVectorScanPlan) -> PyResult { py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - .map(PySearchResult::new) + .map(PyGlobalIndexResult::new) .map_err(to_py_err) } } @@ -324,9 +324,9 @@ impl PyBatchVectorRead { &self, py: Python<'_>, plan: &PyVectorScanPlan, - ) -> PyResult> { + ) -> PyResult> { py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) - .map(|results| results.into_iter().map(PySearchResult::new).collect()) + .map(|results| results.into_iter().map(PyGlobalIndexResult::new).collect()) .map_err(to_py_err) } } diff --git a/bindings/python/tests/test_vector_search.py b/bindings/python/tests/test_vector_search.py index f167b6df5..321ee2779 100644 --- a/bindings/python/tests/test_vector_search.py +++ b/bindings/python/tests/test_vector_search.py @@ -18,7 +18,7 @@ import pyarrow as pa import pytest -from pypaimon_rust.datafusion import PaimonCatalog, SQLContext +from pypaimon_rust.datafusion import GlobalIndexResult, PaimonCatalog, SQLContext @pytest.fixture @@ -57,6 +57,7 @@ def test_filters_accumulate_and_partition_uses_table_schema(table): .with_filter({'method': 'lessThan', 'field': 'id', 'literals': [4]}) .with_partition_filter({'method': 'equal', 'field': 'pt', 'literals': [1]}) .execute_local()) + assert isinstance(result, GlobalIndexResult) assert len(result) == 1 assert list(result.row_ids().values()) == [1.0] assert result.snapshot_id() == 1 @@ -104,6 +105,7 @@ def test_batch_order_empty_results_and_scan_reuse(table): .with_query_vectors(queries).with_limit(1)) plan = builder.new_vector_search_scan().scan() 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] @@ -125,3 +127,18 @@ def test_invalid_query_is_reported_before_planning(table, query): 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() + + +def test_external_result_uses_same_scan_type_without_search_context(table): + found = _single(table).with_limit(1).execute_local() + 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() + expected = pa.Table.from_batches(found.new_read_builder().with_projection(['id']).read()) + actual = pa.Table.from_batches(builder.new_read().read(plan.splits())) + assert actual.column('id').to_pylist() == expected.column('id').to_pylist() + assert external.row_ids() == found.row_ids() + with pytest.raises(Exception, match='require a result returned by a table search'): + external.new_read_builder().read() + with pytest.raises(TypeError): + builder.new_scan().with_global_index_result(found.row_ids()) diff --git a/crates/integrations/datafusion/src/hybrid_search.rs b/crates/integrations/datafusion/src/hybrid_search.rs index 76a2889a7..fd8cddddb 100644 --- a/crates/integrations/datafusion/src/hybrid_search.rs +++ b/crates/integrations/datafusion/src/hybrid_search.rs @@ -226,7 +226,11 @@ impl TableProvider for HybridSearchTableProvider { .with_limit(self.limit) .with_ranker(&self.ranker) .map_err(to_datafusion_error)?; - builder.execute_scored().await.map_err(to_datafusion_error) + builder + .execute_scored() + .await + .and_then(|result| result.into_row_ids()) + .map_err(to_datafusion_error) }) .await?; diff --git a/crates/paimon/src/table/batch_vector_search_builder.rs b/crates/paimon/src/table/batch_vector_search_builder.rs index 68e83c5d1..2d42bc40c 100644 --- a/crates/paimon/src/table/batch_vector_search_builder.rs +++ b/crates/paimon/src/table/batch_vector_search_builder.rs @@ -23,7 +23,7 @@ use crate::table::vector_read::BatchVectorRead; use crate::table::vector_scan::{PlanContext, VectorScan}; use crate::table::vector_search_common::{add_filter, validate_partition_filter}; use crate::table::Table; -use crate::vector_search::SearchResult; +use crate::vector_search::GlobalIndexResult; use roaring::RoaringTreemap; use std::collections::HashMap; use std::sync::Arc; @@ -157,7 +157,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { } /// Search every query against one plan, including empty per-query results. - pub async fn execute(&self) -> crate::Result> { + pub async fn execute(&self) -> crate::Result> { let read = self.new_read()?; read.read(self.new_scan()?.plan().await?).await } diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index ae72aa1c4..6de0f5f3d 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -42,7 +42,7 @@ use crate::table::vector_search_common::{ use crate::table::{ find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, Table, }; -use crate::vector_search::{GlobalIndexIOMeta, ScoredRowIds, SearchResult, VectorSearch}; +use crate::vector_search::{GlobalIndexIOMeta, GlobalIndexResult, ScoredRowIds, VectorSearch}; use crate::vindex::executor::{ acquire_process_global_search_permit, drain_indexed_jobs, ensure_global_index_executor_capacity, execute_global_index_with_guard, @@ -185,7 +185,7 @@ pub(super) async fn materialize_row_ids( impl Read for DeVectorRead { type Plan = DeVectorScanPlan; - async fn read(&self, plan: DeVectorScanPlan) -> crate::Result> { + async fn read(&self, plan: DeVectorScanPlan) -> crate::Result> { let DeVectorScanPlan { table, index_entries, @@ -201,7 +201,7 @@ impl Read for DeVectorRead { .into_iter() .zip(&self.vector_searches) .map(|(hits, query)| { - SearchResult::from_row_ids(table.clone(), query.field_name.clone(), hits) + GlobalIndexResult::from_row_ids(table.clone(), query.field_name.clone(), hits) }) .collect() }; diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs index 0d76e3093..293499f7d 100644 --- a/crates/paimon/src/table/de_vector_read/tests.rs +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -95,7 +95,7 @@ fn make_lumina_entry( } } -// ---- Task B: search-and-read (`SearchResultReadBuilder::read`) tests ---- +// ---- Task B: search-and-read (`GlobalIndexResultReadBuilder::read`) tests ---- /// Build a small materialization batch: user column `id: Int32`, the internal /// `_PKEY_VECTOR_POSITION: Int64`, and `__paimon_search_score: Float32` (mirroring diff --git a/crates/paimon/src/table/de_vector_scan/tests.rs b/crates/paimon/src/table/de_vector_scan/tests.rs index dc35c06f9..0e16c0452 100644 --- a/crates/paimon/src/table/de_vector_scan/tests.rs +++ b/crates/paimon/src/table/de_vector_scan/tests.rs @@ -23,7 +23,7 @@ use crate::table::vector_search_test_utils::{ de_vector_table, id_gt_filter, vector_test_table, vector_test_table_at, }; use crate::table::{TableCommit, TableWrite}; -use crate::vector_search::SearchResult; +use crate::vector_search::GlobalIndexResult; use arrow_array::builder::{Float32Builder, ListBuilder}; use arrow_array::{ArrayRef, Float32Array, Int32Array, RecordBatch}; use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; @@ -341,7 +341,7 @@ async fn de_scalar_filter_with_no_matching_rows_returns_empty() { .await .unwrap(); assert_eq!(results.len(), 2); - assert!(results.iter().all(SearchResult::is_empty)); + assert!(results.iter().all(GlobalIndexResult::is_empty)); } #[tokio::test] diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 11d85bfbd..4b7e4c799 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -126,8 +126,8 @@ impl<'a> FullTextSearchBuilder<'a> { /// Search one snapshot without materializing user columns. Results carry /// global row IDs or primary-key file positions and can be read later with - /// [`crate::table::SearchResult::new_read_builder`]. - pub async fn execute_scored(&self) -> crate::Result { + /// [`crate::table::GlobalIndexResult::new_read_builder`]. + pub async fn execute_scored(&self) -> crate::Result { let read = self.new_read()?; read.read(self.new_scan()?.scan().await?).await } 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 5992e8dcd..d37882857 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 @@ -135,7 +135,7 @@ fn query(table: &Table) -> FullTextSearchBuilder<'_> { builder } -fn ids(result: crate::table::SearchResult) -> Vec { +fn ids(result: crate::table::GlobalIndexResult) -> Vec { let mut ids = result.into_row_ids().unwrap().row_ids; ids.sort_unstable(); ids @@ -382,7 +382,7 @@ async fn empty_string_filters_preserve_java_search_refinement() { } } -async fn user_ids(table: &Table, result: crate::table::SearchResult) -> Vec { +async fn user_ids(table: &Table, result: crate::table::GlobalIndexResult) -> Vec { let mut builder = table.new_read_builder(); builder.with_projection(&["id", ROW_ID_FIELD_NAME]).unwrap(); let plan = builder.new_scan().plan().await.unwrap(); 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 934affe9f..6cee33041 100644 --- a/crates/paimon/src/table/full_text_search_builder/read.rs +++ b/crates/paimon/src/table/full_text_search_builder/read.rs @@ -81,7 +81,10 @@ impl FullTextRead { }) } - pub async fn read(&self, plan: FullTextScanPlan) -> crate::Result { + pub async fn read( + &self, + plan: FullTextScanPlan, + ) -> 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 { @@ -118,7 +121,7 @@ impl FullTextRead { let result = read .search_route(source, &self.search.query_text, self.search.limit) .await?; - return crate::table::SearchResult::from_full_text_primary_key( + return crate::table::GlobalIndexResult::from_full_text_primary_key( table, result.snapshot_id, result.candidates, @@ -142,7 +145,7 @@ impl FullTextRead { &search, ) .await?; - Ok(crate::table::SearchResult::from_full_text_row_ids( + Ok(crate::table::GlobalIndexResult::from_global_search( table, crate::vector_search::ScoredRowIds::new(hits.row_ids, hits.scores), )) diff --git a/crates/paimon/src/table/global_index_result.rs b/crates/paimon/src/table/global_index_result.rs new file mode 100644 index 000000000..914b09a31 --- /dev/null +++ b/crates/paimon/src/table/global_index_result.rs @@ -0,0 +1,416 @@ +// 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. + +//! Global-index selections shared by searches and table scans. + +use super::{merge_row_ranges, DataSplit, Plan, RowRange, Table}; +use crate::spec::BinaryRow; +use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit; +use crate::table::pk_vector_orchestrator::{ + build_indexed_splits, PkVectorCandidate, PkVectorSearchSplit, +}; +use crate::vector_search::ScoredRowIds; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use std::sync::Arc; + +mod read; +pub use read::GlobalIndexResultReadBuilder; + +/// One scored physical row in a primary-key table, in relevance order. +#[derive(Debug, Clone)] +pub struct PrimaryKeySearchPosition { + pub partition: BinaryRow, + pub bucket: i32, + pub data_file_name: String, + pub row_position: i64, + pub score: f32, +} + +#[derive(Debug, Clone)] +enum GlobalIndexSelection { + RowRanges(Vec), + ScoredRowIds(ScoredRowIds), + Splits(Vec), +} + +/// Context retained only by table searches, for ranked reads without re-searching. +#[derive(Debug, Clone)] +struct IndexReadContext { + table: Arc, + vector_column: Option, + positions: Option>, +} + +/// An evaluated global-index selection, accepted directly by table scans. +/// +/// Data Evolution results address stable global row IDs. Primary-key results +/// retain physical splits and their source snapshot, as in Java's +/// `GlobalIndexSplitResult`. Table searches also retain the context needed by +/// [`new_read_builder`](Self::new_read_builder) to read rows in relevance order. +#[derive(Debug, Clone)] +pub struct GlobalIndexResult { + selection: GlobalIndexSelection, + snapshot_id: Option, + read_context: Option, +} + +impl GlobalIndexResult { + pub fn from_row_ranges(ranges: Vec) -> Self { + Self { + selection: GlobalIndexSelection::RowRanges(merge_row_ranges(ranges)), + snapshot_id: None, + read_context: None, + } + } + + pub fn from_scored_row_ids(hits: ScoredRowIds) -> Self { + Self { + selection: GlobalIndexSelection::ScoredRowIds(hits), + snapshot_id: None, + read_context: None, + } + } + + pub fn from_splits(snapshot_id: Option, splits: Vec) -> Self { + Self { + selection: GlobalIndexSelection::Splits(splits), + snapshot_id, + read_context: None, + } + } + + pub(super) fn from_row_ids( + table: Arc
, + vector_column: String, + hits: ScoredRowIds, + ) -> Self { + let snapshot_id = table.travel_snapshot().map(|snapshot| snapshot.id()); + Self { + snapshot_id, + selection: GlobalIndexSelection::ScoredRowIds(hits), + read_context: Some(IndexReadContext { + table, + vector_column: Some(vector_column), + positions: None, + }), + } + } + + pub(super) fn from_global_search(table: Arc
, hits: ScoredRowIds) -> Self { + let snapshot_id = table.travel_snapshot().map(|snapshot| snapshot.id()); + Self { + snapshot_id, + selection: GlobalIndexSelection::ScoredRowIds(hits), + read_context: Some(IndexReadContext { + table, + vector_column: None, + positions: None, + }), + } + } + + #[cfg(feature = "fulltext")] + pub(super) fn from_full_text_primary_key( + table: Arc
, + snapshot_id: i64, + candidates: Vec, + source_splits: &[super::pk_full_text_scan::PrimaryKeyFullTextSearchSplit], + ) -> crate::Result { + let positions = candidates + .iter() + .map(|c| PrimaryKeySearchPosition { + partition: c.partition.clone(), + bucket: c.bucket, + data_file_name: c.data_file_name.clone(), + row_position: c.row_position, + score: c.score, + }) + .collect(); + let splits = + super::pk_full_text_read::build_full_text_indexed_splits(candidates, source_splits)?; + Self::from_indexed_splits(table, snapshot_id, positions, splits) + } + + pub(super) fn from_primary_key( + table: Arc
, + snapshot_id: i64, + candidates: Vec, + source_splits: &[PkVectorSearchSplit], + metric: VectorSearchMetric, + ) -> crate::Result { + let positions = candidates + .iter() + .map(|c| PrimaryKeySearchPosition { + partition: c.partition.clone(), + bucket: c.bucket, + data_file_name: c.data_file_name.clone(), + row_position: c.row_position, + score: metric.distance_to_score(c.distance), + }) + .collect(); + let splits = build_indexed_splits(candidates, source_splits, metric)?; + Self::from_indexed_splits(table, snapshot_id, positions, splits) + } + + fn from_indexed_splits( + table: Arc
, + snapshot_id: i64, + positions: Vec, + splits: Vec, + ) -> crate::Result { + let splits = splits + .into_iter() + .map(|selected| { + let split = selected.split.with_selected_row_ranges(selected.row_ranges); + match selected.scores { + Some(scores) => split.with_scores(scores), + None => Ok(split), + } + }) + .collect::>>()?; + Ok(Self { + selection: GlobalIndexSelection::Splits(splits), + snapshot_id: (snapshot_id != 0).then_some(snapshot_id), + read_context: Some(IndexReadContext { + table, + vector_column: None, + positions: Some(positions), + }), + }) + } + + /// Source snapshot, or `None` when no snapshot was supplied or resolved. + pub fn snapshot_id(&self) -> Option { + self.snapshot_id + } + + /// Source table retained by a search. External selections have no read context. + pub fn table(&self) -> Option<&Table> { + self.read_context + .as_ref() + .map(|context| context.table.as_ref()) + } + + /// Number of selected rows (a lower bound for splits with unknown row counts). + pub fn len(&self) -> usize { + match &self.selection { + GlobalIndexSelection::RowRanges(ranges) => { + ranges.iter().fold(0usize, |count, range| { + count.saturating_add( + usize::try_from(i128::from(range.to()) - i128::from(range.from()) + 1) + .unwrap_or(usize::MAX), + ) + }) + } + GlobalIndexSelection::ScoredRowIds(hits) => hits.len(), + GlobalIndexSelection::Splits(splits) => splits.iter().fold(0usize, |count, split| { + count.saturating_add(usize::try_from(split.row_count()).unwrap_or(usize::MAX)) + }), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Scored global row IDs, in relevance order; rejects physical PK selections. + pub fn row_ids(&self) -> crate::Result<&ScoredRowIds> { + match &self.selection { + GlobalIndexSelection::ScoredRowIds(hits) => Ok(hits), + _ => Err(no_scored_row_ids()), + } + } + + pub fn into_row_ids(self) -> crate::Result { + match self.selection { + GlobalIndexSelection::ScoredRowIds(hits) => Ok(hits), + _ => Err(no_scored_row_ids()), + } + } + + /// Scored physical positions in relevance order, available on PK search results. + pub fn positions(&self) -> crate::Result<&[PrimaryKeySearchPosition]> { + if !self.is_split_result() { + return Err(crate::Error::Unsupported { + message: "Global row-ID results use global row IDs, not primary-key file positions" + .into(), + }); + } + self.read_context + .as_ref() + .and_then(|context| context.positions.as_deref()) + .ok_or_else(|| crate::Error::Unsupported { + message: + "Ranked physical positions are only available on primary-key search results" + .into(), + }) + } + + /// Snapshot-scoped PK selections, including ranges and aligned scores. + pub fn splits(&self) -> crate::Result<&[DataSplit]> { + match &self.selection { + GlobalIndexSelection::Splits(splits) => Ok(splits), + _ => Err(crate::Error::Unsupported { + message: "Global row-ID results do not carry primary-key indexed splits".into(), + }), + } + } + + /// Snapshot-scoped PK selections in Java SplitSerializer v1 format. + pub fn serialize_primary_key_splits(&self) -> crate::Result>> { + self.splits()? + .iter() + .map(DataSplit::serialize_split_v1) + .collect() + } + + pub(super) fn is_split_result(&self) -> bool { + matches!(self.selection, GlobalIndexSelection::Splits(_)) + } + + fn indexed_splits(&self) -> crate::Result> { + self.splits()? + .iter() + .map(|split| { + Ok(PkVectorIndexedSplit { + split: split.clone(), + row_ranges: split + .row_ranges() + .ok_or_else(|| crate::Error::DataInvalid { + message: "Ranked primary-key splits must carry physical row ranges" + .into(), + source: None, + })? + .to_vec(), + scores: split.scores().map(<[f32]>::to_vec), + }) + }) + .collect() + } + + /// Read selected user columns and scores in relevance order. External selections + /// must be read through a table scan because they do not retain a source table. + pub fn new_read_builder(&self) -> GlobalIndexResultReadBuilder<'_> { + GlobalIndexResultReadBuilder::new(self) + } + + pub(crate) fn row_ranges(&self) -> crate::Result>> { + match &self.selection { + GlobalIndexSelection::RowRanges(ranges) => Ok(Some(merge_row_ranges(ranges.clone()))), + GlobalIndexSelection::Splits(_) => Ok(None), + GlobalIndexSelection::ScoredRowIds(hits) => { + if hits.row_ids.len() != hits.scores.len() { + return Err(crate::Error::DataInvalid { + message: "Scored row IDs and scores must have equal lengths".into(), + source: None, + }); + } + let mut ranges = hits + .row_ids + .iter() + .map(|id| { + let id = i64::try_from(*id).map_err(|_| crate::Error::DataInvalid { + message: "Global row ID exceeds the signed row-range address space" + .into(), + source: None, + })?; + Ok(RowRange::new(id, id)) + }) + .collect::>>()?; + ranges.sort_unstable_by_key(RowRange::from); + if ranges.windows(2).any(|r| r[0].from() == r[1].from()) { + return Err(crate::Error::DataInvalid { + message: "Scored global row IDs must be unique".into(), + source: None, + }); + } + Ok(Some(merge_row_ranges(ranges))) + } + } + } + + fn score_map(&self) -> crate::Result>> { + let GlobalIndexSelection::ScoredRowIds(hits) = &self.selection else { + return Ok(None); + }; + if hits.row_ids.len() != hits.scores.len() { + return Err(crate::Error::DataInvalid { + message: "Scored row IDs and scores must have equal lengths".into(), + source: None, + }); + } + let mut scores = std::collections::HashMap::with_capacity(hits.len()); + for (id, score) in hits.row_ids.iter().zip(&hits.scores) { + let id = i64::try_from(*id).map_err(|_| crate::Error::DataInvalid { + message: "Global row ID exceeds the signed row-range address space".into(), + source: None, + })?; + if scores.insert(id, *score).is_some() { + return Err(crate::Error::DataInvalid { + message: "Scored global row IDs must be unique".into(), + source: None, + }); + } + } + Ok(Some(scores)) + } + + pub(crate) fn attach_scores(&self, plan: Plan) -> crate::Result { + let Some(scores) = self.score_map()? else { + return Ok(plan); + }; + let snapshot_id = plan.snapshot_id(); + let splits = plan + .into_splits() + .into_iter() + .map(|split| { + let selected = split + .row_ranges() + .ok_or_else(|| crate::Error::DataInvalid { + message: "Scored global-index splits must carry row ranges".into(), + source: None, + })?; + let mut values = Vec::new(); + for range in selected { + for id in range.from()..=range.to() { + values.push(*scores.get(&id).ok_or_else(|| crate::Error::DataInvalid { + message: format!("Missing score for selected global row ID {id}"), + source: None, + })?); + } + } + split.with_scores(values) + }) + .collect::>>()?; + let plan = Plan::new(splits); + Ok(match snapshot_id { + Some(id) => plan.with_snapshot_id(id), + None => plan, + }) + } +} + +fn no_scored_row_ids() -> crate::Error { + crate::Error::Unsupported { + message: "Scored global row IDs are unavailable for primary-key physical file positions or unscored row ranges".into(), + } +} + +#[cfg(test)] +mod search_tests; +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/global_index_result/read.rs b/crates/paimon/src/table/global_index_result/read.rs new file mode 100644 index 000000000..5cf02b531 --- /dev/null +++ b/crates/paimon/src/table/global_index_result/read.rs @@ -0,0 +1,100 @@ +// 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. + +//! Ranked materialization of a snapshot-scoped global-index result. + +use super::{GlobalIndexResult, GlobalIndexSelection}; +use crate::spec::CoreOptions; +use crate::table::data_file_reader::DataFileReader; +use crate::table::de_vector_read::{materialize_row_ids, DeVectorRead}; +use crate::table::pk_vector_read::materialize_positions; +use crate::table::vector_search_common::resolve_materialize_read_type; +use crate::table::ArrowRecordBatchStream; + +/// Reads selected user columns and `__paimon_search_score`, in relevance order. +/// Internal global row IDs and PK positions are not output. Reading a result +/// never reruns search or replans its primary-key source files. +#[derive(Debug, Clone)] +pub struct GlobalIndexResultReadBuilder<'a> { + result: &'a GlobalIndexResult, + projection: Option>, +} + +impl<'a> GlobalIndexResultReadBuilder<'a> { + pub(super) fn new(result: &'a GlobalIndexResult) -> Self { + Self { + result, + projection: None, + } + } + + /// Select user columns in this order. Defaults to every user table column. + pub fn with_projection(&mut self, columns: &[&str]) -> &mut Self { + self.projection = Some(columns.iter().map(|name| name.to_string()).collect()); + self + } + + pub async fn read(&self) -> crate::Result { + let context = + self.result + .read_context + .as_ref() + .ok_or_else(|| crate::Error::Unsupported { + message: "Direct result reads require a result returned by a table search" + .into(), + })?; + let table = &context.table; + let core_options = CoreOptions::new(table.schema().options()); + core_options.ensure_read_authorized()?; + match &self.result.selection { + GlobalIndexSelection::ScoredRowIds(hits) => { + let read_type = match &context.vector_column { + Some(column) => { + DeVectorRead::read_type(table, column, self.projection.as_deref())? + } + None => resolve_materialize_read_type(table, self.projection.as_deref())?, + }; + materialize_row_ids(table, hits, read_type).await + } + GlobalIndexSelection::Splits(_) => { + let positions = + context + .positions + .as_deref() + .ok_or_else(|| crate::Error::Unsupported { + message: "Direct primary-key reads require ranked search positions" + .into(), + })?; + let splits = self.result.indexed_splits()?; + let read_type = resolve_materialize_read_type(table, self.projection.as_deref())?; + let reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + read_type, + Vec::new(), + ) + .with_table_options(table.schema().options().clone()); + materialize_positions(positions, &splits, &reader).await + } + GlobalIndexSelection::RowRanges(_) => Err(crate::Error::Unsupported { + message: "Unscored row ranges must be read through a table scan".into(), + }), + } + } +} diff --git a/crates/paimon/src/table/search_result/tests.rs b/crates/paimon/src/table/global_index_result/search_tests.rs similarity index 100% rename from crates/paimon/src/table/search_result/tests.rs rename to crates/paimon/src/table/global_index_result/search_tests.rs diff --git a/crates/paimon/src/table/global_index_result/tests.rs b/crates/paimon/src/table/global_index_result/tests.rs new file mode 100644 index 000000000..100ff6abc --- /dev/null +++ b/crates/paimon/src/table/global_index_result/tests.rs @@ -0,0 +1,375 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::*; +use crate::table::vector_search_test_utils::{ + de_vector_table, id_gt_filter, pk_vector_table, vector_test_table_at, +}; +use crate::vector_search::ScoredRowIds; +use std::collections::HashMap; + +fn selected_split() -> DataSplit { + DataSplit::deserialize_split_v1(include_bytes!("../goldens/split_v1_data.bin")) + .unwrap() + .with_selected_row_ranges(vec![RowRange::new(1, 1)]) + .with_scores(vec![0.75]) + .unwrap() +} + +#[tokio::test] +async fn pk_supplied_result_preserves_snapshot_splits_and_ignores_planning_hints() { + // There are no snapshots or manifests to resolve. These belong to the + // captured result, exactly as Java PrimaryKeyBatchScan.preProcessPlan. + let table = pk_vector_table(&[("pk-vector.index.columns", "embedding")]); + let split = selected_split(); + let mut builder = table.new_read_builder(); + builder.with_limit(0).with_filter(id_gt_filter(&table, 100)); + let scan = builder + .new_scan() + .with_global_index_result(GlobalIndexResult::from_splits( + Some(42), + vec![split.clone()], + )) + .unwrap() + .with_shard(1, 2) + .unwrap(); + let plan = scan.plan().await.unwrap(); + let (traced, trace) = scan.plan_with_trace().await.unwrap(); + for plan in [plan, traced] { + assert_eq!(plan.snapshot_id(), Some(42)); + assert_eq!(plan.splits(), std::slice::from_ref(&split)); + assert_eq!( + plan.splits()[0].serialize_split_v1().unwrap(), + split.serialize_split_v1().unwrap() + ); + } + assert_eq!(trace.snapshot_id, Some(42)); + assert_eq!(trace.manifest_entries_read, 0); + assert_eq!(trace.final_splits, 1); + assert_eq!(trace.final_files, split.data_files().len()); + assert_eq!( + trace.planned_data_file_bytes, + split + .data_files() + .iter() + .map(|f| f.file_size as u64) + .sum::() + ); +} + +#[tokio::test] +async fn empty_pk_result_retains_its_snapshot_and_does_not_scan_latest() { + let table = pk_vector_table(&[("pk-vector.index.columns", "embedding")]); + for snapshot_id in [None, Some(7)] { + let scan = table + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_splits(snapshot_id, vec![])) + .unwrap(); + let plan = scan.plan().await.unwrap(); + assert!(plan.splits().is_empty()); + assert_eq!(plan.snapshot_id(), snapshot_id); + } +} + +#[tokio::test] +async fn result_address_spaces_and_conflicting_selections_are_validated() { + let pk = pk_vector_table(&[("pk-vector.index.columns", "embedding")]); + let de = de_vector_table().await; + let append = vector_test_table_at("memory:/append_index_result"); + assert!(pk + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_row_ranges(vec![])) + .is_err()); + for table in [&de, &append] { + assert!(table + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_splits(None, vec![])) + .is_err()); + } + assert!(append + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_row_ranges(vec![])) + .is_err()); + let mut builder = de.new_read_builder(); + builder.with_row_ranges(vec![]); + assert!(builder + .new_scan() + .with_global_index_result(GlobalIndexResult::from_row_ranges(vec![])) + .is_err()); + let scan = de + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_row_ranges(vec![])) + .unwrap(); + assert!(scan.clone().with_row_ranges(vec![]).plan().await.is_err()); + assert!(scan.clone().with_scan_all_files().plan().await.is_err()); + assert!(scan.with_chunk_shuffle(1, 1).is_err()); + assert!(de + .new_read_builder() + .new_scan() + .with_chunk_shuffle(1, 1) + .unwrap() + .with_global_index_result(GlobalIndexResult::from_row_ranges(vec![])) + .is_err()); +} + +#[tokio::test] +async fn de_scores_follow_final_ranges_after_distribution_and_limit() { + let table = de_vector_table().await; + let hits = + GlobalIndexResult::from_scored_row_ids(ScoredRowIds::new(vec![2, 0], vec![0.8, 0.1])); + let mut builder = table.new_read_builder(); + builder.with_projection(&["id"]).unwrap(); + let plan = builder + .new_scan() + .with_global_index_result(hits.clone()) + .unwrap() + .plan() + .await + .unwrap(); + let ranges = plan + .splits() + .iter() + .flat_map(|s| s.row_ranges().unwrap()) + .cloned() + .collect::>(); + assert_eq!(ranges, vec![RowRange::new(0, 0), RowRange::new(2, 2)]); + let scores = plan + .splits() + .iter() + .flat_map(|s| s.scores().unwrap()) + .copied() + .collect::>(); + assert_eq!(scores, vec![0.1, 0.8]); + let selected = builder + .new_scan() + .with_global_index_result(hits) + .unwrap() + .with_row_position_slice(1, 3) + .unwrap() + .plan() + .await + .unwrap(); + assert_eq!(selected.snapshot_id(), plan.snapshot_id()); + assert_eq!(selected.splits().len(), 1); + assert_eq!( + selected.splits()[0].row_ranges(), + Some([RowRange::new(2, 2)].as_slice()) + ); + assert_eq!(selected.splits()[0].scores(), Some([0.8].as_slice())); + let wire = selected.splits()[0].serialize_split_v1().unwrap(); + assert_eq!( + DataSplit::deserialize_split_v1(&wire).unwrap(), + selected.splits()[0] + ); + builder.with_limit(1); + let plan = builder + .new_scan() + .with_global_index_result(GlobalIndexResult::from_scored_row_ids(ScoredRowIds::new( + vec![2, 0], + vec![0.8, 0.1], + ))) + .unwrap() + .plan() + .await + .unwrap(); + // LIMIT is a conservative split hint; every retained position keeps its score. + assert!(plan + .splits() + .iter() + .all(|s| s.scores().unwrap().len() == s.row_count() as usize)); +} + +#[tokio::test] +async fn de_search_results_can_be_reused_by_an_ordinary_scan() { + let table = de_vector_table().await; + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(2) + .execute() + .await + .unwrap(); + let expected = result + .row_ids() + .unwrap() + .row_ids + .iter() + .copied() + .zip(result.row_ids().unwrap().scores.iter().copied()) + .collect::>(); + let plan = result + .table() + .unwrap() + .new_read_builder() + .new_scan() + .with_global_index_result(result.clone()) + .unwrap() + .plan() + .await + .unwrap(); + let mut actual = HashMap::new(); + for split in plan.splits() { + let ids = split + .row_ranges() + .unwrap() + .iter() + .flat_map(|r| r.from()..=r.to()); + actual.extend( + ids.zip(split.scores().unwrap().iter().copied()) + .map(|(id, score)| (id as u64, score)), + ); + } + assert_eq!(plan.snapshot_id(), result.snapshot_id()); + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn invalid_scored_ids_do_not_become_a_full_scan() { + let table = de_vector_table().await; + for hits in [ + ScoredRowIds { + row_ids: vec![0], + scores: vec![], + }, + ScoredRowIds::new(vec![0, 0], vec![1.0, 2.0]), + ScoredRowIds::new(vec![u64::MAX], vec![1.0]), + ] { + assert!(table + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_scored_row_ids(hits)) + .unwrap() + .plan() + .await + .is_err()); + } +} + +#[tokio::test] +async fn supplied_pk_result_cannot_bypass_required_authorization() { + let table = pk_vector_table(&[("pk-vector.index.columns", "embedding")]).copy_with_options( + HashMap::from([("query-auth.enabled".into(), "true".into())]), + ); + let scan = table + .new_read_builder() + .new_scan() + .with_global_index_result(GlobalIndexResult::from_splits( + Some(42), + vec![selected_split()], + )) + .unwrap(); + assert!(scan.plan().await.is_err()); + assert!(scan.plan_with_trace().await.is_err()); +} + +#[tokio::test] +async fn de_combined_incremental_results_retain_the_window_and_scores() { + use crate::table::{IncrementalScanMode, TableCommit}; + use arrow_array::{types::Float32Type, Int32Array, ListArray, RecordBatch}; + use std::sync::Arc; + + let table = de_vector_table().await; + let start = table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap() + .unwrap(); + let builder = table.new_write_builder().with_commit_user("next").unwrap(); + let mut writer = builder.new_write().unwrap(); + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from(vec![4])) as arrow_array::ArrayRef, + ), + ( + "embedding", + Arc::new(ListArray::from_iter_primitive::([Some( + vec![Some(1.0), Some(0.0)], + )])) as arrow_array::ArrayRef, + ), + ]) + .unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + TableCommit::new(table.clone(), "next".into()) + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let end = table + .snapshot_manager() + .get_latest_snapshot_id() + .await + .unwrap() + .unwrap(); + let scan = table + .new_read_builder() + .new_incremental_scan(IncrementalScanMode::Delta, start, end) + .with_global_index_result(GlobalIndexResult::from_scored_row_ids(ScoredRowIds::new( + vec![0, 3], + vec![0.1, 0.9], + ))) + .unwrap(); + assert!(scan.plan().await.is_err()); + let plan = scan.plan_combined().await.unwrap(); + assert_eq!(plan.snapshot_id(), Some(end)); + assert_eq!(plan.splits().len(), 1); + assert_eq!( + plan.splits()[0].row_ranges(), + Some([RowRange::new(3, 3)].as_slice()) + ); + assert_eq!(plan.splits()[0].scores(), Some([0.9].as_slice())); + assert!(table + .new_read_builder() + .new_incremental_scan(IncrementalScanMode::Delta, start, end) + .with_global_index_result(GlobalIndexResult::from_splits(Some(end), vec![])) + .is_err()); +} + +#[tokio::test] +async fn external_selections_need_a_table_scan_for_reads() { + let ranges = GlobalIndexResult::from_row_ranges(vec![RowRange::new(0, i64::MAX)]); + assert_eq!( + ranges.len(), + usize::try_from(i128::from(i64::MAX) + 1).unwrap_or(usize::MAX) + ); + let split = selected_split(); + let physical = GlobalIndexResult::from_splits(Some(42), vec![split.clone()]); + assert_eq!(physical.len(), 1); + assert_eq!( + physical.serialize_primary_key_splits().unwrap(), + vec![split.serialize_split_v1().unwrap()] + ); + for result in [ + ranges, + physical, + GlobalIndexResult::from_scored_row_ids(ScoredRowIds::new(vec![1], vec![0.5])), + ] { + let error = match result.new_read_builder().read().await { + Ok(_) => panic!("external selection unexpectedly retained a source table"), + Err(error) => error, + }; + assert!(error + .to_string() + .contains("require a result returned by a table search")); + } +} diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index 602a3337b..bd6f8b9e5 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -20,6 +20,7 @@ //! Reference: `org.apache.paimon.table.source.HybridSearchBuilder`. use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; use arrow_array::RecordBatch; use futures::{stream, StreamExt, TryStreamExt}; @@ -311,10 +312,10 @@ impl<'a> HybridSearchBuilder<'a> { } pub async fn execute(&self) -> crate::Result> { - self.execute_scored().await?.to_row_ranges() + self.execute_scored().await?.row_ids()?.to_row_ranges() } - pub async fn execute_scored(&self) -> crate::Result { + pub async fn execute_scored(&self) -> crate::Result { let core = CoreOptions::new(self.table.schema().options()); core.ensure_read_authorized()?; let limit = self.validate_search()?; @@ -339,10 +340,20 @@ impl<'a> HybridSearchBuilder<'a> { self.validate_global_routes()?; // Resolve once even for an empty table or a retained tag. Re-resolving // selectors independently in the routes can observe different versions. - let Some(table) = self.resolve_pinned_route_table().await? else { - return Ok(ScoredRowIds::empty()); + let (table, hits) = match self.resolve_pinned_route_table().await? { + Some(table) => { + let hits = self.execute_global_routes(&table, limit).await?; + (table, hits) + } + None => ( + self.table.copy_with_pinned_snapshot(None), + ScoredRowIds::empty(), + ), }; - self.execute_global_routes(&table, limit).await + Ok(super::GlobalIndexResult::from_global_search( + Arc::new(table), + hits, + )) } fn validate_search(&self) -> crate::Result { @@ -650,11 +661,7 @@ impl<'a> HybridSearchBuilder<'a> { .iter() .map(PrimaryKeySearchPosition::from_vector_position) .collect::>>()?; - let source_splits = result - .indexed_splits()? - .iter() - .map(|split| split.split.clone()) - .collect(); + let source_splits = result.splits()?.to_vec(); Ok(PkRoute { positions, source_splits, 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 df2f8e255..8c1031018 100644 --- a/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs +++ b/crates/paimon/src/table/hybrid_search_builder/parity_tests.rs @@ -133,7 +133,16 @@ async fn full_text_route_uses_java_structured_query() { ) .unwrap(); builder.with_limit(4); - assert_eq!(builder.execute_scored().await.unwrap().row_ids, vec![2]); + assert_eq!( + builder + .execute_scored() + .await + .unwrap() + .into_row_ids() + .unwrap() + .row_ids, + vec![2] + ); } #[tokio::test] @@ -187,8 +196,8 @@ async fn both_routes_filter_before_top_k_and_accumulate_predicates() { .with_partition_filter(pb.equal("pt", Datum::Int(1)).unwrap()) .unwrap(); let result = builder.execute_scored().await.unwrap(); - assert_eq!(result.row_ids, vec![6]); - assert!((result.scores[0] - 2.0 / 61.0).abs() < 1e-6); + assert_eq!(result.row_ids().unwrap().row_ids, vec![6]); + assert!((result.row_ids().unwrap().scores[0] - 2.0 / 61.0).abs() < 1e-6); builder .with_partition_filter(pb.equal("pt", Datum::Int(0)).unwrap()) .unwrap(); @@ -213,6 +222,8 @@ async fn scalar_coverage_and_refinement_are_shared_by_both_routes() { .execute_scored() .await .unwrap() + .into_row_ids() + .unwrap() .row_ids; ids.sort_unstable(); // Vector FAST reads indexed ranges from data when no scalar index can @@ -242,6 +253,8 @@ async fn scalar_coverage_and_refinement_are_shared_by_both_routes() { .execute_scored() .await .unwrap() + .into_row_ids() + .unwrap() .row_ids; ids.sort_unstable(); assert_eq!(ids, if refine { vec![1, 2] } else { vec![] }); @@ -272,15 +285,24 @@ async fn concurrent_routes_preserve_weights_limits_and_java_fusion() { .with_ranker(ranker) .unwrap(); let result = builder.execute_scored().await.unwrap(); - assert_eq!(result.row_ids, ids, "ranker {ranker}"); - for (actual, expected) in result.scores.iter().zip(scores) { + assert_eq!(result.row_ids().unwrap().row_ids, ids, "ranker {ranker}"); + for (actual, expected) in result.row_ids().unwrap().scores.iter().zip(scores) { assert!( (actual - expected).abs() < 1e-6, "ranker {ranker}: {actual} != {expected}" ); } builder.with_limit(1); - assert_eq!(builder.execute_scored().await.unwrap().row_ids, vec![3]); + assert_eq!( + builder + .execute_scored() + .await + .unwrap() + .into_row_ids() + .unwrap() + .row_ids, + vec![3] + ); } } @@ -315,6 +337,8 @@ async fn retained_tag_keeps_both_routes_on_the_expired_snapshot() { .execute_scored() .await .unwrap() + .into_row_ids() + .unwrap() .row_ids; old.sort_unstable(); assert_eq!(old, vec![0, 1, 2, 3]); @@ -357,6 +381,8 @@ async fn caller_pinned_empty_and_snapshot_views_do_not_read_latest() { .execute_scored() .await .unwrap() + .into_row_ids() + .unwrap() .row_ids; old.sort_unstable(); assert_eq!(old, vec![0, 1, 2, 3]); diff --git a/crates/paimon/src/table/incremental_scan.rs b/crates/paimon/src/table/incremental_scan.rs index 5eff66b15..b56a65a67 100644 --- a/crates/paimon/src/table/incremental_scan.rs +++ b/crates/paimon/src/table/incremental_scan.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{DataSplit, Plan, ScanTrace, SnapshotManager, Table, TableScan}; +use super::{DataSplit, GlobalIndexResult, Plan, ScanTrace, SnapshotManager, Table, TableScan}; use crate::spec::{CommitKind, CoreOptions}; /// Java IncrementalDeltaStartingScanner.betweenTimestamps. Resolve both bounds @@ -307,6 +307,19 @@ impl<'a> IncrementalScan<'a> { Ok(self) } + /// Intersect a DE global-row-ID selection with the combined incremental + /// batch. Unlike PK split results, this must retain the snapshot window. + /// Use [`Self::plan_combined`] to retain aligned scores on the final splits. + pub fn with_global_index_result(mut self, result: GlobalIndexResult) -> crate::Result { + if result.is_split_result() { + return Err(crate::Error::Unsupported { + message: "Split-based results require an ordinary batch scan".into(), + }); + } + self.scan = self.scan.with_global_index_result(result)?; + Ok(self) + } + pub async fn plan(&self) -> crate::Result { let mode = self.resolve_mode(); self.validate_per_snapshot_plan(mode).await?; @@ -326,9 +339,10 @@ impl<'a> IncrementalScan<'a> { if self.scan.has_row_position_selection() || self.scan.has_chunk_shuffle() || self.scan.has_shard() + || self.scan.has_global_index_result() { return Err(crate::Error::Unsupported { - message: "Incremental row-position selection and chunk_shuffle require combined Delta planning; sharding requires combined Delta or Changelog planning" + message: "Incremental row-position selection and chunk_shuffle require combined Delta planning; sharding and global-index results require combined Delta or Changelog planning" .into(), }); } @@ -437,8 +451,10 @@ impl<'a> IncrementalScan<'a> { }) .map(str::to_owned) .collect(); - scan.plan_incremental_lists_with_trace(&end_snapshot, manifest_lists, trace) - .await + let plan = scan + .plan_incremental_lists_with_trace(&end_snapshot, manifest_lists, trace) + .await?; + self.scan.attach_global_index_scores(plan) } fn resolve_mode(&self) -> IncrementalScanMode { diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index f7caeb626..bf59ab569 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -78,6 +78,7 @@ mod full_text_search_builder; mod global_index_build_builder; pub(crate) mod global_index_build_common; mod global_index_drop_builder; +mod global_index_result; pub(crate) mod global_index_scanner; mod global_index_search_filter; mod global_index_types; @@ -138,7 +139,6 @@ mod row_position_selection; mod row_sidecar; mod scan_trace; pub(crate) mod schema_manager; -pub(crate) mod search_result; mod shard_table_updater; pub(crate) mod snapshot_commit; mod snapshot_manager; @@ -202,6 +202,9 @@ pub use full_text_search_builder::{ use futures::stream::BoxStream; pub use global_index_build_builder::GlobalIndexBuildBuilder; pub use global_index_drop_builder::GlobalIndexDropBuilder; +pub use global_index_result::{ + GlobalIndexResult, GlobalIndexResultReadBuilder, PrimaryKeySearchPosition, +}; pub use global_index_types::{normalize_global_index_type, SUPPORTED_GLOBAL_INDEX_TYPES_FOR_DROP}; pub use hybrid_search_builder::{ HybridSearchBuilder, HybridSearchRanker, HybridSearchRoute, HybridSearchRouteKind, @@ -222,7 +225,6 @@ pub use read_builder::ReadBuilder; pub use rest_env::RESTEnv; pub use scan_trace::ScanTrace; pub use schema_manager::SchemaManager; -pub use search_result::{PrimaryKeySearchPosition, SearchResult, SearchResultReadBuilder}; pub use shard_table_updater::ShardTableUpdater; pub use snapshot_commit::{RESTSnapshotCommit, RenamingSnapshotCommit, SnapshotCommit}; pub use snapshot_manager::SnapshotManager; diff --git a/crates/paimon/src/table/pk_full_text_read.rs b/crates/paimon/src/table/pk_full_text_read.rs index a5523a302..e83504db0 100644 --- a/crates/paimon/src/table/pk_full_text_read.rs +++ b/crates/paimon/src/table/pk_full_text_read.rs @@ -17,7 +17,7 @@ //! Primary-key full-text search read foundation: the score-based candidate type, //! its cross-bucket score-descending Top-K, and the FAST-only search over -//! physical positions. SearchResult materializes the selected rows later. +//! physical positions. GlobalIndexResult materializes the selected rows later. use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; @@ -255,12 +255,12 @@ pub(crate) struct PrimaryKeyFullTextRouteResult<'a> { /// FAST-only primary-key full-text search. Given a planned set of /// per-bucket search splits, it searches each bucket through the full-text archive -/// reader and fuses physical positions by score. SearchResult owns the selections +/// reader and fuses physical positions by score. GlobalIndexResult owns the selections /// and their later materialization, as Java PrimaryKeyScoredResult does. /// /// The planning + FAST-mode guard live in the caller (`FullTextSearchBuilder`); /// this type searches a prepared plan. Its data reader loads deletion vectors; -/// projected user columns are read later by SearchResult. +/// projected user columns are read later by GlobalIndexResult. pub(crate) struct PrimaryKeyFullTextRead { file_io: FileIO, materialize_reader: DataFileReader, @@ -287,7 +287,7 @@ impl PrimaryKeyFullTextRead { /// best-score-first Top-`limit`, WITHOUT materializing any rows. Returns the /// surviving candidates together with the route source context (the plan's /// resolved snapshot id and its per-bucket source splits) so a caller can fuse - /// them with another route before materializing through SearchResult. + /// them with another route before materializing through GlobalIndexResult. /// `query` is passed verbatim to the archive reader /// (spec D2). pub(crate) async fn search_route<'p>( @@ -519,7 +519,7 @@ mod read_tests { None, ); let route = self.search_route(plan, query, limit).await?; - crate::table::SearchResult::from_full_text_primary_key( + crate::table::GlobalIndexResult::from_full_text_primary_key( Arc::new(table), route.snapshot_id, route.candidates, diff --git a/crates/paimon/src/table/pk_search_position.rs b/crates/paimon/src/table/pk_search_position.rs index 2d8c36fa4..aa522228a 100644 --- a/crates/paimon/src/table/pk_search_position.rs +++ b/crates/paimon/src/table/pk_search_position.rs @@ -222,7 +222,7 @@ mod tests { #[test] fn from_vector_position_preserves_score_and_validates_position() { - // SearchResult already converted distance to score; fusion must preserve it. + // GlobalIndexResult already converted distance to score; fusion must preserve it. let mut hit = vector_position(0.5); let position = PrimaryKeySearchPosition::from_vector_position(&hit).unwrap(); assert_eq!(position.score(), 0.5); diff --git a/crates/paimon/src/table/pk_vector_indexed_split_read.rs b/crates/paimon/src/table/pk_vector_indexed_split_read.rs index 4feeacb2f..7fefb6f19 100644 --- a/crates/paimon/src/table/pk_vector_indexed_split_read.rs +++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs @@ -49,8 +49,9 @@ fn data_invalid(message: impl Into) -> crate::Error { /// `scores`, when present, is aligned to the expanded-range order (ascending /// position); `None` means no `__paimon_search_score` column in the output. /// -/// Deliberately NOT reusing `DataSplit.row_ranges`, whose ranges mean stable/global -/// row ids on the append/data-evolution path. Not serialized. +/// This is the position reader's explicit physical selection. `GlobalIndexResult` +/// retains the same ranges in its PK IndexedSplits and recreates this reader input +/// at materialization. DE split ranges use a separate, global row-ID address space. #[derive(Debug, Clone)] pub(crate) struct PkVectorIndexedSplit { pub split: DataSplit, diff --git a/crates/paimon/src/table/pk_vector_read.rs b/crates/paimon/src/table/pk_vector_read.rs index 59f74c443..7194d35a6 100644 --- a/crates/paimon/src/table/pk_vector_read.rs +++ b/crates/paimon/src/table/pk_vector_read.rs @@ -43,7 +43,7 @@ use crate::table::vector_search_common::{ reorder_and_strip_position, vindex_concurrency_limits, RankedRow, VectorIndexBackend, }; use crate::table::{ArrowRecordBatchStream, RowRange, Table}; -use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; +use crate::vector_search::{GlobalIndexIOMeta, GlobalIndexResult, VectorSearch}; use crate::vindex::pkvector::ann::{AnnSegmentSource, PkVectorAnnSearcher, VindexAnnSearcher}; use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture}; use crate::vindex::pkvector::metric::VectorSearchMetric; @@ -92,7 +92,7 @@ impl PkVectorRead { impl Read for PkVectorRead { type Plan = PkVectorScanPlan; - async fn read(&self, plan: PkVectorScanPlan) -> crate::Result> { + async fn read(&self, plan: PkVectorScanPlan) -> crate::Result> { let core = CoreOptions::new(self.table.schema().options()); let queries: Vec<&[f32]> = self.queries.iter().map(Vec::as_slice).collect(); let candidates = search_pk_candidates_batch_with_plan( @@ -111,7 +111,7 @@ impl Read for PkVectorRead { candidates .into_iter() .map(|candidates| { - SearchResult::from_primary_key( + GlobalIndexResult::from_primary_key( table.clone(), plan.snapshot_id, candidates, diff --git a/crates/paimon/src/table/pk_vector_read/tests.rs b/crates/paimon/src/table/pk_vector_read/tests.rs index 59df08f93..a31f00038 100644 --- a/crates/paimon/src/table/pk_vector_read/tests.rs +++ b/crates/paimon/src/table/pk_vector_read/tests.rs @@ -1066,7 +1066,7 @@ async fn reader_uses_planned_splits_after_index_manifest_is_removed() { .collect::>(), vec![1.0, 0.5] ); - assert_eq!(result.indexed_splits().unwrap().len(), 1); + assert_eq!(result.splits().unwrap().len(), 1); let batches: Vec = result .new_read_builder() @@ -1125,15 +1125,19 @@ async fn execute_returns_scored_positions_and_publishes_diagnostics() { // Every position retains the source file and snapshot needed for a later read. assert_eq!(result.snapshot_id(), Some(1), "first commit -> snapshot 1"); - let splits = result.indexed_splits().unwrap(); + let splits = result.splits().unwrap(); assert_eq!(splits.len(), 1); - assert_eq!(splits[0].split.snapshot_id(), 1); + assert_eq!(splits[0].snapshot_id(), 1); + assert_eq!(splits[0].scores(), Some([1.0, 0.5].as_slice())); + assert_eq!( + splits[0].row_ranges(), + Some([crate::table::RowRange::new(1, 2)].as_slice()) + ); assert!( positions.iter().all(|p| { - p.data_file_name == splits[0].split.data_files()[0].file_name - && p.bucket == splits[0].split.bucket() - && p.partition.to_serialized_bytes() - == splits[0].split.partition().to_serialized_bytes() + p.data_file_name == splits[0].data_files()[0].file_name + && p.bucket == splits[0].bucket() + && p.partition.to_serialized_bytes() == splits[0].partition().to_serialized_bytes() }), "positions must refer to their retained source split" ); @@ -1177,7 +1181,7 @@ async fn execute_empty_plan_yields_empty_pk_result() { "no data -> no positions" ); assert!( - result.indexed_splits().unwrap().is_empty(), + result.splits().unwrap().is_empty(), "no data -> no source splits" ); assert_eq!(result.snapshot_id(), None); diff --git a/crates/paimon/src/table/search_result.rs b/crates/paimon/src/table/search_result.rs deleted file mode 100644 index 4c5861e0d..000000000 --- a/crates/paimon/src/table/search_result.rs +++ /dev/null @@ -1,293 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Snapshot-scoped search results for global row IDs and primary-key file positions. - -use crate::spec::{BinaryRow, CoreOptions}; -use crate::table::data_file_reader::DataFileReader; -use crate::table::de_vector_read::{materialize_row_ids, DeVectorRead}; -use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit; -use crate::table::pk_vector_orchestrator::{ - build_indexed_splits, PkVectorCandidate, PkVectorSearchSplit, -}; -use crate::table::pk_vector_read::materialize_positions; -use crate::table::vector_search_common::resolve_materialize_read_type; -use crate::table::{ArrowRecordBatchStream, Table}; -use crate::vector_search::ScoredRowIds; -use crate::vindex::pkvector::metric::VectorSearchMetric; -use std::sync::Arc; - -/// One scored physical row in a primary-key table. Positions are local to the -/// named file; neither the primary key nor a global row ID is synthesized. -#[derive(Debug, Clone)] -pub struct PrimaryKeySearchPosition { - pub partition: BinaryRow, - pub bucket: i32, - pub data_file_name: String, - pub row_position: i64, - pub score: f32, -} - -#[derive(Debug, Clone)] -enum SearchHits { - DataEvolution { - // Vector results retain their target-column validation at read time; - // full-text columns have already been validated by FullTextRead. - vector_column: Option, - hits: ScoredRowIds, - }, - PrimaryKey { - positions: Vec, - splits: Vec, - }, -} - -/// Ranked search hits together with the source snapshot needed to read them. -/// -/// Single/batch vector and full-text searches return this type. DE hits use global row IDs; -/// PK hits use physical file positions, as in Java's `PrimaryKeyScoredResult`. -/// Searching does not materialize projected user columns. Use -/// [`new_read_builder`](Self::new_read_builder) when rows are needed. -#[derive(Debug, Clone)] -pub struct SearchResult { - table: Arc
, - snapshot_id: Option, - hits: SearchHits, -} - -impl SearchResult { - pub(super) fn from_row_ids( - table: Arc
, - vector_column: String, - hits: ScoredRowIds, - ) -> Self { - let snapshot_id = table.travel_snapshot().map(|snapshot| snapshot.id()); - Self { - table, - snapshot_id, - hits: SearchHits::DataEvolution { - vector_column: Some(vector_column), - hits, - }, - } - } - - #[cfg(feature = "fulltext")] - pub(super) fn from_full_text_row_ids(table: Arc
, hits: ScoredRowIds) -> Self { - let snapshot_id = table.travel_snapshot().map(|snapshot| snapshot.id()); - Self { - table, - snapshot_id, - hits: SearchHits::DataEvolution { - vector_column: None, - hits, - }, - } - } - - #[cfg(feature = "fulltext")] - pub(super) fn from_full_text_primary_key( - table: Arc
, - snapshot_id: i64, - candidates: Vec, - source_splits: &[super::pk_full_text_scan::PrimaryKeyFullTextSearchSplit], - ) -> crate::Result { - let positions = candidates - .iter() - .map(|c| PrimaryKeySearchPosition { - partition: c.partition.clone(), - bucket: c.bucket, - data_file_name: c.data_file_name.clone(), - row_position: c.row_position, - score: c.score, - }) - .collect(); - let splits = - super::pk_full_text_read::build_full_text_indexed_splits(candidates, source_splits)?; - Ok(Self { - table, - snapshot_id: (snapshot_id != 0).then_some(snapshot_id), - hits: SearchHits::PrimaryKey { positions, splits }, - }) - } - - pub(super) fn from_primary_key( - table: Arc
, - snapshot_id: i64, - candidates: Vec, - source_splits: &[PkVectorSearchSplit], - metric: VectorSearchMetric, - ) -> crate::Result { - let positions = candidates - .iter() - .map(|c| PrimaryKeySearchPosition { - partition: c.partition.clone(), - bucket: c.bucket, - data_file_name: c.data_file_name.clone(), - row_position: c.row_position, - score: metric.distance_to_score(c.distance), - }) - .collect(); - let splits = build_indexed_splits(candidates, source_splits, metric)?; - Ok(Self { - table, - snapshot_id: (snapshot_id != 0).then_some(snapshot_id), - hits: SearchHits::PrimaryKey { positions, splits }, - }) - } - - /// Source snapshot, or `None` when the table had no resolved snapshot. - pub fn snapshot_id(&self) -> Option { - self.snapshot_id - } - - /// The source table retained by the search; DE scans are pinned to this snapshot. - pub fn table(&self) -> &Table { - &self.table - } - - pub fn len(&self) -> usize { - match &self.hits { - SearchHits::DataEvolution { hits, .. } => hits.len(), - SearchHits::PrimaryKey { positions, .. } => positions.len(), - } - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Scored global row IDs, in relevance order. PK positions cannot be converted - /// into this address space, including when a PK search has no hits. - pub fn row_ids(&self) -> crate::Result<&ScoredRowIds> { - match &self.hits { - SearchHits::DataEvolution { hits, .. } => Ok(hits), - SearchHits::PrimaryKey { .. } => Err(no_global_row_ids()), - } - } - - pub fn into_row_ids(self) -> crate::Result { - match self.hits { - SearchHits::DataEvolution { hits, .. } => Ok(hits), - SearchHits::PrimaryKey { .. } => Err(no_global_row_ids()), - } - } - - /// Scored physical positions in relevance order, scoped to this result's snapshot. - pub fn positions(&self) -> crate::Result<&[PrimaryKeySearchPosition]> { - match &self.hits { - SearchHits::PrimaryKey { positions, .. } => Ok(positions), - SearchHits::DataEvolution { .. } => Err(crate::Error::Unsupported { - message: "data-evolution search results use global row IDs, not primary-key file positions".to_string(), - }), - } - } - - /// Snapshot-scoped PK selections in Java SplitSerializer v1 format. - /// The ranges are physical file positions and scores are aligned to them. - pub fn serialize_primary_key_splits(&self) -> crate::Result>> { - self.indexed_splits()? - .iter() - .map(|selected| { - selected - .split - .clone() - .with_selected_row_ranges(selected.row_ranges.clone()) - .serialize_split_v1_with_scores(selected.scores.as_deref()) - }) - .collect() - } - - /// Snapshot-scoped file metadata and selections, reused when hybrid fusion - /// builds its final selections without scanning the source table again. - pub(super) fn indexed_splits(&self) -> crate::Result<&[PkVectorIndexedSplit]> { - match &self.hits { - SearchHits::PrimaryKey { splits, .. } => Ok(splits), - SearchHits::DataEvolution { .. } => Err(crate::Error::Unsupported { - message: "data-evolution search results do not carry primary-key indexed splits" - .to_string(), - }), - } - } - - pub fn new_read_builder(&self) -> SearchResultReadBuilder<'_> { - SearchResultReadBuilder { - result: self, - projection: None, - } - } -} - -fn no_global_row_ids() -> crate::Error { - crate::Error::Unsupported { - message: "primary-key search results use physical file positions, not global row IDs" - .to_string(), - } -} - -/// Reads selected user columns and `__paimon_search_score`, in relevance order. -/// Internal global row IDs and PK positions are not output. Reading a result -/// never reruns search or replans its primary-key source files. -#[derive(Debug, Clone)] -pub struct SearchResultReadBuilder<'a> { - result: &'a SearchResult, - projection: Option>, -} - -impl SearchResultReadBuilder<'_> { - /// Select user columns in this order. Defaults to every user table column. - pub fn with_projection(&mut self, columns: &[&str]) -> &mut Self { - self.projection = Some(columns.iter().map(|name| name.to_string()).collect()); - self - } - - pub async fn read(&self) -> crate::Result { - let table = &self.result.table; - let core_options = CoreOptions::new(table.schema().options()); - core_options.ensure_read_authorized()?; - match &self.result.hits { - SearchHits::DataEvolution { - vector_column, - hits, - } => { - let read_type = match vector_column { - Some(column) => { - DeVectorRead::read_type(table, column, self.projection.as_deref())? - } - None => resolve_materialize_read_type(table, self.projection.as_deref())?, - }; - materialize_row_ids(table, hits, read_type).await - } - SearchHits::PrimaryKey { positions, splits } => { - let read_type = resolve_materialize_read_type(table, self.projection.as_deref())?; - let reader = DataFileReader::new( - table.file_io().clone(), - table.schema_manager().clone(), - table.schema().id(), - table.schema().fields().to_vec(), - read_type, - Vec::new(), - ) - .with_table_options(table.schema().options().clone()); - materialize_positions(positions, splits, &reader).await - } - } - } -} - -#[cfg(test)] -mod tests; diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 623e1419e..355b76949 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -494,7 +494,7 @@ pub struct DataSplit { partition: Arc, bucket: i32, bucket_path: Arc, - total_buckets: i32, + total_buckets: Option, data_files: Arc<[DataFileMeta]>, /// Deletion file for each data file, same order as `data_files`. /// `None` at index `i` means no deletion file for `data_files[i]` (matches Java getDeletionFiles() / List with null elements). @@ -504,6 +504,9 @@ pub struct DataSplit { /// Primary-key IndexedSplits use file-local physical positions, independent /// of row tracking. row_ranges: Option>, + /// Optional IndexedSplit scores, aligned to expanded row ranges. + #[serde(default)] + scores: Option>, /// Whether the split can be read raw, without the merge reader: its /// physical rows are exactly its logical rows (modulo deletion files). /// Mirrors Java `DataSplit#rawConvertible`. @@ -565,7 +568,7 @@ impl DataSplit { &self.bucket_path } pub fn total_buckets(&self) -> i32 { - self.total_buckets + self.total_buckets.unwrap_or(1) } pub fn data_files(&self) -> &[DataFileMeta] { @@ -581,10 +584,24 @@ impl DataSplit { self.row_ranges.as_deref() } + pub fn scores(&self) -> Option<&[f32]> { + self.scores.as_deref() + } + + /// Attach scores in ascending expanded-range order, as in Java IndexedSplit. + pub fn with_scores(mut self, scores: Vec) -> crate::Result { + validate_scores(self.row_ranges(), &scores)?; + self.scores = Some(scores.into()); + Ok(self) + } + /// Narrow an already planned split without rebuilding or copying its file /// and deletion-vector metadata. Ranges use the same coordinate system as /// the existing `row_ranges` field. pub(crate) fn with_selected_row_ranges(mut self, ranges: Vec) -> Self { + // Scores describe the old selection; callers must supply scores for + // the new one explicitly rather than silently misaligning them. + self.scores = None; self.row_ranges = Some(ranges.into()); self } @@ -600,6 +617,7 @@ impl DataSplit { .as_ref() .map(|files| vec![files[file_index].clone()].into()); split.row_ranges = None; + split.scores = None; split.raw_convertible = false; split } @@ -797,8 +815,13 @@ impl DataSplit { out.extend_from_slice(&p); out.extend_from_slice(&self.bucket.to_be_bytes()); write_java_utf(&mut out, &self.bucket_path)?; - out.push(1); // totalBuckets present (version >= 6) - out.extend_from_slice(&self.total_buckets.to_be_bytes()); + match self.total_buckets { + Some(total) => { + out.push(1); + out.extend_from_slice(&total.to_be_bytes()); + } + None => out.push(0), + } out.extend_from_slice(&0i32.to_be_bytes()); // deprecated beforeFiles count out.push(0); // beforeDeletionFiles = null list out.extend_from_slice(&(self.data_files.len() as i32).to_be_bytes()); @@ -935,7 +958,9 @@ impl DataSplit { if let Some(dels) = data_deletion_files { builder = builder.with_data_deletion_files(dels); } - builder.build() + let mut split = builder.build()?; + split.total_buckets = total_buckets; + Ok(split) } /// Java `SplitSerializer#serialize` frame: magic + version + type id + body. The cross-language @@ -944,7 +969,7 @@ impl DataSplit { /// uses `QueryAuthSplit` (type 5), retaining the complete captured REST response. Byte-compatible with /// `compatibility/split-v1-data` / `split-v1-indexed`. pub fn serialize_split_v1(&self) -> crate::Result> { - self.serialize_split_v1_with_scores(None) + self.serialize_split_v1_with_scores(self.scores()) } /// Serialize aligned vector scores using Java's IndexedSplit wire format. @@ -975,18 +1000,7 @@ impl DataSplit { } self.ensure_serializable_without_grant()?; if let Some(scores) = scores { - let count = self.row_ranges.as_ref().and_then(|ranges| { - ranges.iter().try_fold(0usize, |count, range| { - let len = range.to().checked_sub(range.from())?.checked_add(1)?; - count.checked_add(usize::try_from(len).ok()?) - }) - }); - if count != Some(scores.len()) || scores.len() > i32::MAX as usize { - return Err(crate::Error::DataInvalid { - message: "IndexedSplit scores must align with selected row ranges".to_string(), - source: None, - }); - } + validate_scores(self.row_ranges(), scores)?; } let mut out = Vec::new(); out.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); @@ -1025,7 +1039,7 @@ impl DataSplit { /// Reverse of [`DataSplit::serialize_split_v1`]: parse a Java `SplitSerializer` frame into a /// `DataSplit` (type 1), an `IndexedSplit` (type 3, which carries `row_ranges`), or a /// `QueryAuthSplit` (type 5, retaining the authorization-required marker without minting - /// a grant). Other type ids, and IndexedSplit vector scores, are unsupported. + /// a grant). IndexedSplit scores retain their range alignment. Other type ids are unsupported. /// Consumes the entire buffer. pub fn deserialize_split_v1(data: &[u8]) -> crate::Result { let mut cur = data; @@ -1105,14 +1119,21 @@ impl DataSplit { } ranges.push(RowRange::new(from, to)); } - let scores_flag = read_u8(cur)?; - if scores_flag != 0 { - return Err(crate::Error::Unsupported { - message: "IndexedSplit scores are not supported".into(), - }); - } let mut body = body; body.row_ranges = Some(ranges.into()); + if read_u8(cur)? != 0 { + let count = read_i32(cur)?; + if count < 0 || count as usize > cur.len() / 4 { + return Err(crate::Error::DataInvalid { + message: "Invalid IndexedSplit scores count".into(), + source: None, + }); + } + let scores = (0..count) + .map(|_| read_i32(cur).map(|bits| f32::from_bits(bits as u32))) + .collect::>>()?; + body = body.with_scores(scores)?; + } body } other => { @@ -1125,6 +1146,22 @@ impl DataSplit { } } +fn validate_scores(ranges: Option<&[RowRange]>, scores: &[f32]) -> crate::Result<()> { + let count = ranges.and_then(|ranges| { + ranges.iter().try_fold(0usize, |count, range| { + let len = range.to().checked_sub(range.from())?.checked_add(1)?; + count.checked_add(usize::try_from(len).ok()?) + }) + }); + if count != Some(scores.len()) || scores.len() > i32::MAX as usize { + return Err(crate::Error::DataInvalid { + message: "IndexedSplit scores must align with selected row ranges".into(), + source: None, + }); + } + Ok(()) +} + /// Java `DataSplit#MAGIC` / `VERSION` for the serialize format. const SPLIT_MAGIC: i64 = -2394839472490812314; const SPLIT_VERSION: i32 = 9; @@ -1538,10 +1575,11 @@ impl DataSplitBuilder { partition: Arc::new(partition), bucket: self.bucket, bucket_path: bucket_path.into(), - total_buckets: self.total_buckets, + total_buckets: Some(self.total_buckets), data_files: data_files.into(), data_deletion_files: self.data_deletion_files.map(Into::into), row_ranges: self.row_ranges.map(Into::into), + scores: None, raw_convertible: self.raw_convertible, is_streaming: self.is_streaming, }) @@ -2683,9 +2721,7 @@ mod tests { .is_err()); } - // The scores-stripped IndexedSplit frame the Rust serializer emits (see - // `serialize_indexed_split_v1_matches_golden`): the Java golden's trailing vector scores are - // replaced by a `false` scores flag, which is the exact byte form `deserialize_split_v1` supports. + // A score-free frame derived from the Java fixture. fn indexed_split_v1_no_scores() -> Vec { let golden = include_bytes!("goldens/split_v1_indexed.bin"); let scores_len = 1 + 4 + 3 * 4; // writeBoolean(true) + count + 3 floats @@ -2727,15 +2763,58 @@ mod tests { } #[test] - fn deserialize_split_v1_rejects_indexed_scores() { - // The raw Java golden carries vector scores, which the Rust split cannot model. + fn deserialize_split_v1_rejects_misaligned_indexed_scores() { + // This historical Java fixture has three scores for seven positions. let golden = include_bytes!("goldens/split_v1_indexed.bin"); assert!(matches!( DataSplit::deserialize_split_v1(golden), - Err(crate::Error::Unsupported { .. }) + Err(crate::Error::DataInvalid { .. }) )); } + #[test] + fn indexed_scores_round_trip_and_reject_malformed_payloads() { + let split = DataSplit::deserialize_split_v1(&indexed_split_v1_no_scores()).unwrap(); + let values = vec![0.5, 1.25, -2.0, 4.0, 5.0, 6.0, 7.0]; + let split = split.with_scores(values.clone()).unwrap(); + let bytes = split.serialize_split_v1().unwrap(); + let restored = DataSplit::deserialize_split_v1(&bytes).unwrap(); + assert_eq!(restored.scores(), Some(values.as_slice())); + assert_eq!(restored, split); + assert_eq!(restored.serialize_split_v1().unwrap(), bytes); + assert_eq!( + serde_json::from_slice::(&serde_json::to_vec(&split).unwrap()).unwrap(), + split + ); + assert!(restored.clone().with_scores(vec![1.0]).is_err()); + assert!(restored + .with_selected_row_ranges(vec![RowRange::new(1, 1)]) + .scores() + .is_none()); + let count_offset = bytes.len() - values.len() * 4 - 4; + for count in [-1, 0, 6, i32::MAX] { + let mut malformed = bytes.clone(); + malformed[count_offset..count_offset + 4].copy_from_slice(&count.to_be_bytes()); + assert!(DataSplit::deserialize_split_v1(&malformed).is_err()); + } + for length in count_offset..bytes.len() { + assert!(DataSplit::deserialize_split_v1(&bytes[..length]).is_err()); + } + } + + #[test] + fn absent_total_buckets_round_trips_without_materializing_a_default() { + let mut split = v1_data_split_builder().build().unwrap(); + let present_len = split.serialize_split_v1().unwrap().len(); + split.total_buckets = None; + let bytes = split.serialize_split_v1().unwrap(); + assert_eq!(bytes.len(), present_len - 4); + let restored = DataSplit::deserialize_split_v1(&bytes).unwrap(); + assert_eq!(restored.total_buckets(), 1); + assert_eq!(restored.total_buckets, None); + assert_eq!(restored.serialize_split_v1().unwrap(), bytes); + } + #[test] fn deserialize_split_v1_rejects_unsupported_type() { // hand-build a frame with type id 2 diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 345021b37..c949080a1 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -23,6 +23,7 @@ use super::bucket_filter::compute_target_buckets; use super::chunk_shuffle::{chunk_shuffle_splits, ChunkShuffle}; use super::format_table_scan::FormatTableScan; +use super::global_index_result::GlobalIndexResult; use super::global_index_scanner::RowRangeIndex; use super::global_index_types::normalize_queryable_global_index_type; use super::kv_file_reader::retain_primary_key_conjuncts; @@ -1269,6 +1270,25 @@ impl<'a> TableScan<'a> { } } + /// Supply an already evaluated global-index selection. PK split results + /// bypass snapshot resolution and further pruning, preserving their source + /// metadata exactly. Query-authorized scans ignore supplied selections, + /// since indexes describe raw values that masks may change. + pub fn with_global_index_result(mut self, result: GlobalIndexResult) -> crate::Result { + match &mut self.0 { + TableScanKind::Paimon(scan) => { + scan.split_selection + .get_or_insert_with(Default::default) + .global_index_result = Some(result); + scan.validate_global_index_result()?; + Ok(self) + } + TableScanKind::Format(_) => Err(crate::Error::Unsupported { + message: "Format tables do not support global-index results".into(), + }), + } + } + /// Select a half-open range of logical positions in an append snapshot. /// /// Data-evolution positions are assigned before group statistics, projection @@ -1303,10 +1323,13 @@ impl<'a> TableScan<'a> { message: "chunk_shuffle cannot combine with limit".to_string(), }); } - if scan.row_ranges.is_some() || scan.row_position_selection().is_some() { + if scan.row_ranges.is_some() + || scan.row_position_selection().is_some() + || scan.global_index_result().is_some() + { return Err(crate::Error::Unsupported { message: - "chunk_shuffle cannot combine with row ranges or positional selection" + "chunk_shuffle cannot combine with row ranges, global index results or positional selection" .to_string(), }); } @@ -1328,6 +1351,7 @@ impl<'a> TableScan<'a> { only_real_buckets: scan.only_real_buckets(), snapshot_mode: scan.snapshot_mode(), snapshot_levels: scan.snapshot_levels().cloned(), + global_index_result: None, })); Ok(Self(TableScanKind::Paimon(scan))) } @@ -1371,6 +1395,7 @@ impl<'a> TableScan<'a> { only_real_buckets: false, snapshot_mode: None, snapshot_levels: None, + global_index_result: None, })); } } @@ -1397,6 +1422,17 @@ impl<'a> TableScan<'a> { matches!(&self.0, TableScanKind::Paimon(scan) if scan.shard().is_some()) } + pub(super) fn has_global_index_result(&self) -> bool { + matches!(&self.0, TableScanKind::Paimon(scan) if scan.global_index_result().is_some()) + } + + pub(super) fn attach_global_index_scores(&self, plan: Plan) -> crate::Result { + match &self.0 { + TableScanKind::Paimon(scan) => scan.attach_index_scores(plan), + TableScanKind::Format(_) => Ok(plan), + } + } + fn with_row_position_selection(self, selection: RowPositionSelection) -> crate::Result { match self.0 { TableScanKind::Paimon(mut scan) if scan.table.schema().primary_keys().is_empty() => { @@ -1433,6 +1469,7 @@ impl<'a> TableScan<'a> { only_real_buckets: scan.only_real_buckets(), snapshot_mode: scan.snapshot_mode(), snapshot_levels: scan.snapshot_levels().cloned(), + global_index_result: scan.global_index_result().cloned(), })); Ok(Self(TableScanKind::Paimon(scan))) } @@ -1564,6 +1601,7 @@ struct ScanSplitSelection { only_real_buckets: bool, snapshot_mode: Option, snapshot_levels: Option>, + global_index_result: Option, } // A stream scan owns its table so clients can retain the scan across polls. @@ -1614,6 +1652,52 @@ struct PaimonTableScan<'a> { } impl<'a> PaimonTableScan<'a> { + fn global_index_result(&self) -> Option<&GlobalIndexResult> { + self.split_selection + .as_deref() + .and_then(|s| s.global_index_result.as_ref()) + } + + fn validate_global_index_result(&self) -> crate::Result<()> { + let Some(result) = self.global_index_result() else { + return Ok(()); + }; + let options = self.table.schema().core_options(); + let message = if self.row_ranges.is_some() { + Some("with_global_index_result and with_row_ranges are mutually exclusive") + } else if self.chunk_shuffle().is_some() { + Some("chunk_shuffle cannot combine with global index results") + } else if self.scan_all_files || self.is_streaming() { + Some("Global-index results require a batch data scan") + } else { + if result.is_split_result() && self.table.schema().primary_keys().is_empty() { + Some("Split-based global-index results require a primary-key table") + } else if !result.is_split_result() && !options.data_evolution_enabled() { + Some("Row-range global-index results require a Data Evolution table") + } else { + None + } + }; + match message { + Some(message) => Err(crate::Error::Unsupported { + message: message.into(), + }), + None => Ok(()), + } + } + + fn supplied_split_plan(&self) -> Option { + if self.table.schema().core_options().query_auth_enabled() { + return None; + } + let result = self.global_index_result()?; + let mut plan = Plan::new(result.splits().ok()?.to_vec()); + if let Some(id) = result.snapshot_id() { + plan = plan.with_snapshot_id(id); + } + Some(plan) + } + fn snapshot_levels(&self) -> Option<&std::ops::RangeInclusive> { self.split_selection .as_deref() @@ -1748,7 +1832,24 @@ impl<'a> PaimonTableScan<'a> { /// for `scan.version`; the strict selectors mirror Java's typed /// `scan.snapshot-id` / `scan.tag-name` handling. pub async fn plan(&self) -> crate::Result { + let plan = self.plan_selection().await?; + self.attach_index_scores(plan) + } + + fn attach_index_scores(&self, plan: Plan) -> crate::Result { + match self.global_index_result() { + Some(result) if !self.table.schema().core_options().query_auth_enabled() => { + result.attach_scores(plan) + } + _ => Ok(plan), + } + } + + async fn plan_selection(&self) -> crate::Result { let scan = self.authorized_scan().await?; + if let Some(plan) = scan.supplied_split_plan() { + return Ok(plan); + } if let Some((start, end)) = scan .table .schema() @@ -1771,11 +1872,26 @@ impl<'a> PaimonTableScan<'a> { /// Plan the full scan and return metadata-pruning trace counters. pub async fn plan_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { + let (plan, trace) = self.plan_selection_with_trace().await?; + Ok((self.attach_index_scores(plan)?, trace)) + } + + async fn plan_selection_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { let scan = self.authorized_scan().await?; let mut trace = ScanTrace { limit: scan.limit, ..Default::default() }; + if let Some(plan) = scan.supplied_split_plan() { + trace.snapshot_id = plan.snapshot_id(); + trace.planned_data_file_bytes = plan.planned_data_file_bytes(); + trace.record_final_plan( + plan.splits().len(), + plan.splits().len(), + plan.splits().iter().map(|s| s.data_files().len()).sum(), + ); + return Ok((plan, trace)); + } if let Some((start, end)) = scan .table .schema() @@ -1946,6 +2062,7 @@ impl<'a> PaimonTableScan<'a> { } async fn authorized_scan(&self) -> crate::Result> { + self.validate_global_index_result()?; // Timestamp-window planning delegates to combined incremental planning // within the same operation. Reuse its scope, not another REST request. if self.query_auth_grant.is_some() { @@ -1955,7 +2072,21 @@ impl<'a> PaimonTableScan<'a> { let grant = self.authorize_query().await?; Ok(match grant { Some(grant) => Cow::Owned(self.restricted_by(&grant)), - None => Cow::Borrowed(self), + None => match self + .global_index_result() + .map(GlobalIndexResult::row_ranges) + .transpose()? + .flatten() + { + Some(ranges) => { + let mut scan = self.clone(); + scan.row_ranges = Some(ranges); + // The selection is now an ordinary row-ID restriction. + scan.split_selection.as_mut().unwrap().global_index_result = None; + Cow::Owned(scan) + } + _ => Cow::Borrowed(self), + }, }) } diff --git a/crates/paimon/src/table/vector_read.rs b/crates/paimon/src/table/vector_read.rs index be2c55f0d..c8ae740d2 100644 --- a/crates/paimon/src/table/vector_read.rs +++ b/crates/paimon/src/table/vector_read.rs @@ -24,7 +24,7 @@ use crate::table::pk_vector_search_params::PkVectorSearchParams; use crate::table::vector_scan::{PlanContext, VectorScanPlan, VectorScanWork}; use crate::table::vector_search_common::{take_only_result, targets_primary_key_column}; use crate::table::Table; -use crate::vector_search::SearchResult; +use crate::vector_search::GlobalIndexResult; use std::collections::HashMap; use std::future::Future; @@ -38,7 +38,7 @@ pub(super) trait Read: Sync { fn read( &self, plan: Self::Plan, - ) -> impl Future>> + Send; + ) -> impl Future>> + Send; } /// Searches a common plan with one query, without replanning or materializing rows. @@ -47,7 +47,7 @@ pub struct VectorRead { } impl VectorRead { - pub async fn read(&self, plan: VectorScanPlan) -> crate::Result { + pub async fn read(&self, plan: VectorScanPlan) -> crate::Result { take_only_result(self.batch.read(plan).await?, "vector search") } } @@ -88,7 +88,7 @@ impl BatchVectorRead { Ok(Self { context, reader }) } - pub async fn read(&self, plan: VectorScanPlan) -> crate::Result> { + pub async fn read(&self, plan: VectorScanPlan) -> crate::Result> { if self.context != plan.context { return Err(crate::Error::DataInvalid { message: "vector plan and reader must use the same table, column and pre-filter" diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 8614795b2..d8fc37f8a 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -22,7 +22,7 @@ use crate::table::vector_read::{BatchVectorRead, VectorRead}; use crate::table::vector_scan::{PlanContext, VectorScan}; use crate::table::vector_search_common::{add_filter, validate_partition_filter}; use crate::table::Table; -use crate::vector_search::SearchResult; +use crate::vector_search::GlobalIndexResult; use std::collections::HashMap; pub struct VectorSearchBuilder<'a> { @@ -119,7 +119,7 @@ impl<'a> VectorSearchBuilder<'a> { /// Search locally using the same Scan -> Plan -> Read API exposed to engines. /// Use the result's `new_read_builder()` to materialize projected columns. - pub async fn execute(&self) -> crate::Result { + pub async fn execute(&self) -> crate::Result { let read = self.new_read()?; read.read(self.new_scan()?.plan().await?).await } diff --git a/crates/paimon/src/table/vector_search_builder/tests.rs b/crates/paimon/src/table/vector_search_builder/tests.rs index 3638254ee..40b27adef 100644 --- a/crates/paimon/src/table/vector_search_builder/tests.rs +++ b/crates/paimon/src/table/vector_search_builder/tests.rs @@ -614,7 +614,7 @@ async fn result_read_empty_plan_lumina_array_float_is_admitted() { #[tokio::test] async fn result_read_projection_reserved_name_fails_loud() { // Projecting a reserved metadata / row-id column must fail loud. The guard - // lives in `resolve_materialize_read_type`, which `SearchResultReadBuilder::read` invokes + // lives in `resolve_materialize_read_type`, which `GlobalIndexResultReadBuilder::read` invokes // before the empty-plan early return; assert on the resolver directly here. let table = pk_vector_table(&[ ("pk-vector.index.columns", "embedding"), diff --git a/crates/paimon/src/vector_search.rs b/crates/paimon/src/vector_search.rs index 1ebc3f716..dc5327857 100644 --- a/crates/paimon/src/vector_search.rs +++ b/crates/paimon/src/vector_search.rs @@ -15,9 +15,7 @@ // specific language governing permissions and limitations // under the License. -pub use crate::table::search_result::{ - PrimaryKeySearchPosition, SearchResult, SearchResultReadBuilder, -}; +pub use crate::table::{GlobalIndexResult, GlobalIndexResultReadBuilder, PrimaryKeySearchPosition}; use std::cmp::Ordering; use std::collections::{BinaryHeap, HashMap}; @@ -116,7 +114,7 @@ impl GlobalIndexIOMeta { } /// Global row IDs and aligned scores used by DE index evaluation and ranking. -/// Table searches return [`SearchResult`], which also represents PK file positions. +/// Table searches return [`GlobalIndexResult`], which also represents PK file positions. #[derive(Debug, Clone)] pub struct ScoredRowIds { pub row_ids: Vec, diff --git a/crates/paimon/tests/pk_vector_batch_test.rs b/crates/paimon/tests/pk_vector_batch_test.rs index c7fcb32d8..9181939d8 100644 --- a/crates/paimon/tests/pk_vector_batch_test.rs +++ b/crates/paimon/tests/pk_vector_batch_test.rs @@ -373,7 +373,7 @@ async fn drain_ids_and_scores(stream: ArrowRecordBatchStream) -> (Vec, Vec< (ids, scores) } -/// Single-query `SearchResultReadBuilder::read` into `(id, score)` tuples. +/// Single-query `GlobalIndexResultReadBuilder::read` into `(id, score)` tuples. async fn single_read(table: &Table, query: Vec, limit: usize) -> (Vec, Vec) { let mut builder = table.new_vector_search_builder(); builder