From cd2e2d1f14e83893513c4b487e432b1da04ad0f1 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 9 Oct 2026 21:45:08 +0800 Subject: [PATCH 1/3] feat(python): expose core vector search and scored results --- bindings/python/README.md | 24 + .../python/pypaimon_rust/datafusion.pyi | 63 +++ bindings/python/src/context.rs | 8 + bindings/python/src/lib.rs | 1 + bindings/python/src/table.rs | 8 + bindings/python/src/vector_search.rs | 430 ++++++++++++++++++ bindings/python/tests/test_vector_search.py | 122 +++++ .../src/table/batch_vector_search_builder.rs | 11 +- .../batch_vector_search_builder/tests.rs | 83 ++++ crates/paimon/src/table/de_vector_read.rs | 92 +++- .../paimon/src/table/de_vector_read/tests.rs | 183 +++++++- crates/paimon/src/table/de_vector_scan.rs | 20 + .../src/table/pk_vector_orchestrator.rs | 3 +- crates/paimon/src/table/source.rs | 59 ++- .../paimon/src/table/vector_search_builder.rs | 11 +- .../paimon/src/table/vector_search_common.rs | 26 +- .../paimon/src/table/vector_search_result.rs | 15 + 17 files changed, 1135 insertions(+), 24 deletions(-) create mode 100644 bindings/python/src/vector_search.rs create mode 100644 bindings/python/tests/test_vector_search.py diff --git a/bindings/python/README.md b/bindings/python/README.md index 6de276a07..b4710b858 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -56,6 +56,30 @@ ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp") For the full SQL reference, see the [SQL Integration docs](https://paimon.apache.org/docs/master/sql/). +### Native Vector Search + +A table with an existing vector column can search through the Rust core without +materializing user rows in Python: + +```python +search = (table.new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector([1.0, 0.0]) + .with_limit(10)) +plan = search.new_vector_search_scan().scan() +result = search.new_vector_search_read().read_plan(plan) +batches = result.new_read_builder().with_projection(["id"]).read() +``` + +The plan pins one snapshot and can be reused with different queries. Scalar and +partition filters are applied before ranking. Batch builders return one result +per input query, in input order, including empty results. + +Data Evolution results expose scored global row IDs through `row_ids()`. PK +results expose physical positions and Java-compatible scored split bytes; +`new_read_builder()` reads both result types from their source snapshot and +returns projected columns plus `__paimon_search_score` in relevance order. + ### Native Read / Write Beyond SQL, you can use the lower-level read and write APIs directly from Python. diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 5cce8692a..4ff2065da 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -221,6 +221,8 @@ class Table: def schema(self) -> TableSchema: ... def new_read_builder(self, options: Optional[Dict[str, str]] = None) -> ReadBuilder: ... def new_stream_scan(self) -> StreamTableScan: ... + def new_vector_search_builder(self) -> "VectorSearchBuilder": ... + def new_batch_vector_search_builder(self) -> "BatchVectorSearchBuilder": ... def new_batch_write_builder(self) -> "BatchWriteBuilder": ... def new_postpone_fixed_bucket_write_builder(self) -> "PostponeFixedBucketWriteBuilder": ... def new_stream_write_builder(self) -> "StreamWriteBuilder": ... @@ -534,3 +536,64 @@ class SQLContext: def sql_with_options( self, sql: str, options: SQLOptions ) -> List[pyarrow.RecordBatch]: ... + + +class VectorSearchBuilder: + def with_vector_column(self, name: str) -> "VectorSearchBuilder": ... + def with_limit(self, limit: int) -> "VectorSearchBuilder": ... + def with_query_vector(self, vector: Sequence[float]) -> "VectorSearchBuilder": ... + def with_option(self, key: str, value: str) -> "VectorSearchBuilder": ... + def with_options(self, options: Dict[str, str]) -> "VectorSearchBuilder": ... + def with_filter(self, predicate: Dict[str, Any]) -> "VectorSearchBuilder": ... + def with_partition_filter(self, predicate: Dict[str, Any]) -> "VectorSearchBuilder": + """Add a partition-only predicate, resolved by table field names.""" + ... + def new_vector_search_scan(self) -> "VectorScan": ... + def new_vector_search_read(self) -> "VectorRead": ... + def execute_local(self) -> "SearchResult": ... + +class BatchVectorSearchBuilder: + def with_vector_column(self, name: str) -> "BatchVectorSearchBuilder": ... + def with_limit(self, limit: int) -> "BatchVectorSearchBuilder": ... + def with_query_vectors(self, vectors: Sequence[Sequence[float]]) -> "BatchVectorSearchBuilder": ... + def with_option(self, key: str, value: str) -> "BatchVectorSearchBuilder": ... + def with_options(self, options: Dict[str, str]) -> "BatchVectorSearchBuilder": ... + def with_filter(self, predicate: Dict[str, Any]) -> "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"]: ... + +class VectorScan: + def scan(self) -> "VectorScanPlan": ... + +class VectorScanPlan: + def snapshot_id(self) -> Optional[int]: ... + +class VectorRead: + def read_plan(self, plan: VectorScanPlan) -> "SearchResult": ... + +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.""" + ... + + +class SearchResultReadBuilder: + def with_projection(self, columns: List[str]) -> "SearchResultReadBuilder": ... + def read(self) -> List[pyarrow.RecordBatch]: + """Read user columns and scores in relevance order from the search snapshot.""" + ... diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index b6a10ff89..1c50299c3 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -726,6 +726,14 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> let this = PyModule::new(py, "datafusion")?; 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/lib.rs b/bindings/python/src/lib.rs index b4e270f77..900e19ef9 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -30,6 +30,7 @@ mod table; mod udf; mod update_assignment; mod variant; +mod vector_search; mod write; // ---- #285: observability ---- mod partition; diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs index 2421a0abc..fd39b0c98 100644 --- a/bindings/python/src/table.rs +++ b/bindings/python/src/table.rs @@ -198,6 +198,14 @@ impl PyTable { } } + fn new_vector_search_builder(&self) -> crate::vector_search::PyVectorSearchBuilder { + crate::vector_search::PyVectorSearchBuilder::new(Arc::clone(&self.inner)) + } + + fn new_batch_vector_search_builder(&self) -> crate::vector_search::PyBatchVectorSearchBuilder { + crate::vector_search::PyBatchVectorSearchBuilder::new(Arc::clone(&self.inner)) + } + fn new_batch_write_builder(&self) -> PyBatchWriteBuilder { PyBatchWriteBuilder::new(Arc::clone(&self.inner)) } diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs new file mode 100644 index 000000000..e5e99eaa0 --- /dev/null +++ b/bindings/python/src/vector_search.rs @@ -0,0 +1,430 @@ +// 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 std::collections::HashMap; +use std::sync::Arc; + +use arrow::pyarrow::ToPyArrow; +use futures::TryStreamExt; +use paimon::spec::Predicate; +use paimon::table::{BatchVectorRead, Table, VectorRead, VectorScan, VectorScanPlan}; +use paimon::vector_search::SearchResult; +use paimon_datafusion::runtime::runtime; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict}; + +use crate::error::to_py_err; +use crate::predicate::dict_to_table_predicate; + +#[derive(Clone, Default)] +struct SearchConfig { + column: Option, + limit: Option, + options: Vec>, + filters: Vec, + partition_filters: Vec, +} + +impl SearchConfig { + fn single<'a>( + &self, + table: &'a Table, + ) -> paimon::Result> { + let mut builder = table.new_vector_search_builder(); + if let Some(column) = &self.column { + builder.with_vector_column(column); + } + if let Some(limit) = self.limit { + builder.with_limit(limit); + } + for options in &self.options { + builder.with_options(options.clone()); + } + for filter in &self.filters { + builder.with_filter(filter.clone()); + } + for filter in &self.partition_filters { + builder.with_partition_filter(filter.clone())?; + } + Ok(builder) + } + + fn batch<'a>( + &self, + table: &'a Table, + ) -> paimon::Result> { + let mut builder = table.new_batch_vector_search_builder(); + if let Some(column) = &self.column { + builder.with_vector_column(column); + } + if let Some(limit) = self.limit { + builder.with_limit(limit); + } + for options in &self.options { + builder.with_options(options.clone()); + } + for filter in &self.filters { + builder.with_filter(filter.clone()); + } + for filter in &self.partition_filters { + builder.with_partition_filter(filter.clone())?; + } + Ok(builder) + } +} + +#[pyclass(name = "VectorSearchBuilder", module = "pypaimon_rust.datafusion")] +pub struct PyVectorSearchBuilder { + table: Arc, + config: SearchConfig, + query: Option>, +} + +impl PyVectorSearchBuilder { + pub(crate) fn new(table: Arc
) -> Self { + Self { + table, + config: SearchConfig::default(), + query: None, + } + } +} + +#[pymethods] +impl PyVectorSearchBuilder { + fn with_vector_column(mut slf: PyRefMut<'_, Self>, name: String) -> PyRefMut<'_, Self> { + slf.config.column = Some(name); + slf + } + fn with_limit(mut slf: PyRefMut<'_, Self>, limit: usize) -> PyRefMut<'_, Self> { + slf.config.limit = Some(limit); + slf + } + fn with_query_vector(mut slf: PyRefMut<'_, Self>, vectors: Vec) -> PyRefMut<'_, Self> { + slf.query = Some(vectors); + slf + } + fn with_option(mut slf: PyRefMut<'_, Self>, key: String, value: String) -> PyRefMut<'_, Self> { + slf.config.options.push(HashMap::from([(key, value)])); + slf + } + fn with_options( + mut slf: PyRefMut<'_, Self>, + options: HashMap, + ) -> PyRefMut<'_, Self> { + slf.config.options.push(options); + slf + } + fn with_filter<'py>( + mut slf: PyRefMut<'py, Self>, + predicate: &Bound<'_, PyDict>, + ) -> PyResult> { + let predicate = dict_to_table_predicate(predicate, slf.table.schema(), true)?; + slf.config.filters.push(predicate); + Ok(slf) + } + fn with_partition_filter<'py>( + mut slf: PyRefMut<'py, Self>, + predicate: &Bound<'_, PyDict>, + ) -> PyResult> { + let predicate = dict_to_table_predicate(predicate, slf.table.schema(), true)?; + slf.config + .single(&slf.table) + .map_err(to_py_err)? + .with_partition_filter(predicate.clone()) + .map_err(to_py_err)?; + slf.config.partition_filters.push(predicate); + Ok(slf) + } + fn new_vector_search_scan(&self) -> PyResult { + let inner = self + .config + .single(&self.table) + .map_err(to_py_err)? + .new_scan() + .map_err(to_py_err)?; + Ok(PyVectorScan { inner }) + } + fn new_vector_search_read(&self) -> 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()); + } + let inner = builder.new_read().map_err(to_py_err)?; + Ok(PyVectorRead { inner }) + } + 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()); + } + let results = py + .detach(|| runtime().block_on(builder.execute())) + .map_err(to_py_err)?; + Ok(PySearchResult::new(results)) + } +} + +#[pyclass(name = "BatchVectorSearchBuilder", module = "pypaimon_rust.datafusion")] +pub struct PyBatchVectorSearchBuilder { + table: Arc
, + config: SearchConfig, + query: Option>>, +} + +impl PyBatchVectorSearchBuilder { + pub(crate) fn new(table: Arc
) -> Self { + Self { + table, + config: SearchConfig::default(), + query: None, + } + } +} + +#[pymethods] +impl PyBatchVectorSearchBuilder { + fn with_vector_column(mut slf: PyRefMut<'_, Self>, name: String) -> PyRefMut<'_, Self> { + slf.config.column = Some(name); + slf + } + fn with_limit(mut slf: PyRefMut<'_, Self>, limit: usize) -> PyRefMut<'_, Self> { + slf.config.limit = Some(limit); + slf + } + fn with_query_vectors( + mut slf: PyRefMut<'_, Self>, + vectors: Vec>, + ) -> PyRefMut<'_, Self> { + slf.query = Some(vectors); + slf + } + fn with_option(mut slf: PyRefMut<'_, Self>, key: String, value: String) -> PyRefMut<'_, Self> { + slf.config.options.push(HashMap::from([(key, value)])); + slf + } + fn with_options( + mut slf: PyRefMut<'_, Self>, + options: HashMap, + ) -> PyRefMut<'_, Self> { + slf.config.options.push(options); + slf + } + fn with_filter<'py>( + mut slf: PyRefMut<'py, Self>, + predicate: &Bound<'_, PyDict>, + ) -> PyResult> { + let predicate = dict_to_table_predicate(predicate, slf.table.schema(), true)?; + slf.config.filters.push(predicate); + Ok(slf) + } + fn with_partition_filter<'py>( + mut slf: PyRefMut<'py, Self>, + predicate: &Bound<'_, PyDict>, + ) -> PyResult> { + let predicate = dict_to_table_predicate(predicate, slf.table.schema(), true)?; + slf.config + .batch(&slf.table) + .map_err(to_py_err)? + .with_partition_filter(predicate.clone()) + .map_err(to_py_err)?; + slf.config.partition_filters.push(predicate); + Ok(slf) + } + fn new_vector_search_scan(&self) -> PyResult { + let inner = self + .config + .batch(&self.table) + .map_err(to_py_err)? + .new_scan() + .map_err(to_py_err)?; + Ok(PyVectorScan { inner }) + } + fn new_batch_vector_search_read(&self) -> 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()); + } + let inner = builder.new_read().map_err(to_py_err)?; + Ok(PyBatchVectorRead { inner }) + } + 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()); + } + let results = py + .detach(|| runtime().block_on(builder.execute())) + .map_err(to_py_err)?; + Ok(results.into_iter().map(PySearchResult::new).collect()) + } +} + +#[pyclass(name = "VectorScan", module = "pypaimon_rust.datafusion")] +pub struct PyVectorScan { + inner: VectorScan, +} + +#[pymethods] +impl PyVectorScan { + fn scan(&self, py: Python<'_>) -> PyResult { + let inner = py + .detach(|| runtime().block_on(self.inner.plan())) + .map_err(to_py_err)?; + Ok(PyVectorScanPlan { inner }) + } +} + +#[pyclass(name = "VectorScanPlan", module = "pypaimon_rust.datafusion")] +pub struct PyVectorScanPlan { + inner: VectorScanPlan, +} + +#[pymethods] +impl PyVectorScanPlan { + fn snapshot_id(&self) -> Option { + self.inner.snapshot_id() + } +} + +#[pyclass(name = "VectorRead", module = "pypaimon_rust.datafusion")] +pub struct PyVectorRead { + inner: VectorRead, +} + +#[pymethods] +impl PyVectorRead { + fn read_plan(&self, py: Python<'_>, plan: &PyVectorScanPlan) -> PyResult { + py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + .map(PySearchResult::new) + .map_err(to_py_err) + } +} + +#[pyclass(name = "BatchVectorRead", module = "pypaimon_rust.datafusion")] +pub struct PyBatchVectorRead { + inner: BatchVectorRead, +} + +#[pymethods] +impl PyBatchVectorRead { + fn read_batch_plan( + &self, + py: Python<'_>, + plan: &PyVectorScanPlan, + ) -> PyResult> { + py.detach(|| runtime().block_on(self.inner.read(plan.inner.clone()))) + .map(|results| results.into_iter().map(PySearchResult::new).collect()) + .map_err(to_py_err) + } +} + +#[pyclass(name = "SearchResult", module = "pypaimon_rust.datafusion")] +pub struct PySearchResult { + inner: SearchResult, +} + +impl PySearchResult { + fn new(inner: SearchResult) -> Self { + Self { inner } + } +} + +type PySearchPosition = (Py, i32, String, i64, f32); + +#[pymethods] +impl PySearchResult { + fn snapshot_id(&self) -> Option { + self.inner.snapshot_id() + } + fn new_read_builder(&self) -> PySearchResultReadBuilder { + PySearchResultReadBuilder { + result: self.inner.clone(), + projection: None, + } + } + fn __len__(&self) -> usize { + self.inner.len() + } + fn row_ids(&self) -> PyResult> { + let hits = self.inner.row_ids().map_err(to_py_err)?; + Ok(hits + .row_ids + .iter() + .copied() + .zip(hits.scores.iter().copied()) + .collect()) + } + fn positions(&self, py: Python<'_>) -> PyResult> { + Ok(self + .inner + .positions() + .map_err(to_py_err)? + .iter() + .map(|position| { + ( + PyBytes::new(py, &position.partition.to_serialized_bytes()).unbind(), + position.bucket, + position.data_file_name.clone(), + position.row_position, + position.score, + ) + }) + .collect()) + } + fn splits(&self, py: Python<'_>) -> PyResult>> { + Ok(self + .inner + .serialize_primary_key_splits() + .map_err(to_py_err)? + .iter() + .map(|bytes| PyBytes::new(py, bytes).unbind()) + .collect()) + } +} + +#[pyclass(name = "SearchResultReadBuilder", module = "pypaimon_rust.datafusion")] +pub struct PySearchResultReadBuilder { + result: SearchResult, + projection: Option>, +} + +#[pymethods] +impl PySearchResultReadBuilder { + fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec) -> PyRefMut<'_, Self> { + slf.projection = Some(columns); + slf + } + fn read(&self, py: Python<'_>) -> PyResult>> { + let batches = py + .detach(|| { + runtime().block_on(async { + let mut builder = self.result.new_read_builder(); + if let Some(projection) = &self.projection { + let columns: Vec<&str> = projection.iter().map(String::as_str).collect(); + builder.with_projection(&columns); + } + builder.read().await?.try_collect::>().await + }) + }) + .map_err(to_py_err)?; + batches + .iter() + .map(|batch| batch.to_pyarrow(py).map(Bound::unbind)) + .collect() + } +} diff --git a/bindings/python/tests/test_vector_search.py b/bindings/python/tests/test_vector_search.py new file mode 100644 index 000000000..21e3947e6 --- /dev/null +++ b/bindings/python/tests/test_vector_search.py @@ -0,0 +1,122 @@ +# 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. + +import pyarrow as pa +import pytest + +from pypaimon_rust.datafusion import PaimonCatalog, SQLContext + + +@pytest.fixture +def table(tmp_path): + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": str(tmp_path)}) + ctx.sql("CREATE SCHEMA paimon.db") + ctx.sql("CREATE TABLE paimon.db.t (id INT, embedding FLOAT[], pt INT) " + "PARTITIONED BY (pt) WITH ('bucket'='-1', 'data-evolution.enabled'='true', " + "'row-tracking.enabled'='true', 'global-index.enabled'='true', " + "'vector-index.search-mode'='full')") + table = PaimonCatalog({"warehouse": str(tmp_path)}).get_table("db.t") + _write(table, [1, 2, 3, 4], [[1, 0], [0, 1], [1, 0], None], [0, 0, 1, 1]) + return table + + +def _write(table, ids, vectors, partitions): + builder = table.new_batch_write_builder() + writer = builder.new_write() + commit = builder.new_commit() + try: + writer.write_arrow(pa.record_batch([ids, vectors, partitions], schema=pa.schema([ + ('id', pa.int32()), ('embedding', pa.list_(pa.float32())), ('pt', pa.int32())]))) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + + +def _single(table): + return table.new_vector_search_builder().with_vector_column('embedding').with_query_vector([1, 0]).with_limit(10) + + +def test_filters_accumulate_and_partition_uses_table_schema(table): + result = (_single(table).with_filter({'method': 'greaterThan', 'field': 'id', 'literals': [1]}) + .with_filter({'method': 'lessThan', 'field': 'id', 'literals': [4]}) + .with_partition_filter({'method': 'equal', 'field': 'pt', 'literals': [1]}) + .execute_local()) + assert len(result) == 1 + assert list(result.row_ids().values()) == [1.0] + assert result.snapshot_id() == 1 + with pytest.raises(Exception, match='global row IDs'): + result.positions() + with pytest.raises(Exception, match='primary-key indexed splits'): + result.splits() + + +@pytest.mark.parametrize('batch', [False, True]) +def test_partition_filter_rejects_data_fields(table, batch): + builder = table.new_batch_vector_search_builder() if batch else table.new_vector_search_builder() + with pytest.raises(Exception, match='Partition filter'): + builder.with_partition_filter({'method': 'equal', 'field': 'id', 'literals': [1]}) + with pytest.raises(ValueError, match='missing'): + builder.with_filter({'method': 'equal', 'field': 'missing', 'literals': [1]}) + + +def test_scan_read_keeps_snapshot_and_is_reusable(table): + builder = _single(table) + plan = builder.new_vector_search_scan().scan() + _write(table, [5], [[1, 0]], [0]) + reader = builder.new_vector_search_read() + result = reader.read_plan(plan) + assert len(result) == 3 + rows = pa.Table.from_batches(result.new_read_builder().with_projection(['id']).read()) + assert rows.column_names == ['id', '__paimon_search_score'] + ids = rows.column('id').to_pylist() + # Partition file publication determines global row IDs, so tied scores + # may order IDs 1 and 3 either way. Both must precede the weaker hit. + assert sorted(ids[:2]) == [1, 3] + assert ids[2] == 2 + assert rows.column('__paimon_search_score').to_pylist() == pytest.approx([1, 1, 1 / 3]) + assert result.snapshot_id() == plan.snapshot_id() == 1 + assert reader.read_plan(plan).row_ids() == result.row_ids() + assert len(builder.execute_local()) == 4 + filtered = _single(table).with_filter({'method': 'equal', 'field': 'id', 'literals': [1]}) + with pytest.raises(Exception, match='plan.*reader|context|different'): + filtered.new_vector_search_read().read_plan(plan) + + +def test_batch_order_empty_results_and_scan_reuse(table): + queries = [[1, 0], [0, 1], [1, 0]] + builder = (table.new_batch_vector_search_builder().with_vector_column('embedding') + .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 [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] + builder.with_filter({'method': 'equal', 'field': 'id', 'literals': [99]}) + assert [len(r) for r in builder.execute_batch_local()] == [0, 0, 0] + + +@pytest.mark.parametrize('query', [None, []]) +def test_invalid_query_is_reported_before_planning(table, query): + builder = table.new_vector_search_builder().with_vector_column('embedding').with_limit(1) + if query is not None: + builder.with_query_vector(query) + with pytest.raises(Exception, match='vector|Vector'): + builder.new_vector_search_read() + with pytest.raises(Exception, match='Query vectors'): + table.new_batch_vector_search_builder().with_vector_column('embedding').with_limit(1).with_query_vectors([]).execute_batch_local() diff --git a/crates/paimon/src/table/batch_vector_search_builder.rs b/crates/paimon/src/table/batch_vector_search_builder.rs index f80a7b248..68e83c5d1 100644 --- a/crates/paimon/src/table/batch_vector_search_builder.rs +++ b/crates/paimon/src/table/batch_vector_search_builder.rs @@ -21,6 +21,7 @@ use crate::spec::{CoreOptions, Predicate}; use crate::table::de_vector_scan::PreparedVectorSearchFilter; 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 roaring::RoaringTreemap; @@ -68,7 +69,7 @@ impl<'a> BatchVectorSearchBuilder<'a> { } pub fn with_options(&mut self, options: HashMap) -> &mut Self { - self.options = options; + self.options.extend(options); self } @@ -76,12 +77,18 @@ impl<'a> BatchVectorSearchBuilder<'a> { /// before vector Top-K. See [`crate::table::VectorSearchBuilder::with_filter`] for the /// primary-key and data-evolution execution semantics. pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { - self.filter = Some(filter); + add_filter(&mut self.filter, filter); self.include_row_ids = None; self.prepared_filter = None; self } + /// Add a partition-only predicate expressed against the table schema. + pub fn with_partition_filter(&mut self, filter: Predicate) -> crate::Result<&mut Self> { + validate_partition_filter(self.table, &filter)?; + Ok(self.with_filter(filter)) + } + /// Attach a prepared scalar pre-filter together with the exact table /// snapshot against which its row-ID allow-list was evaluated. pub fn with_prepared_filter( diff --git a/crates/paimon/src/table/batch_vector_search_builder/tests.rs b/crates/paimon/src/table/batch_vector_search_builder/tests.rs index 125e776f5..22e91097e 100644 --- a/crates/paimon/src/table/batch_vector_search_builder/tests.rs +++ b/crates/paimon/src/table/batch_vector_search_builder/tests.rs @@ -155,3 +155,86 @@ fn batch_resolve_materialize_read_type_default_rejects_reserved_user_column() { "batch default projection must reject reserved user column, got: {err:?}" ); } + +#[tokio::test] +async fn repeated_filters_are_conjoined_for_single_and_batch_searches() { + use crate::spec::{Datum, PredicateBuilder}; + let table = de_vector_table().await; + let predicates = PredicateBuilder::new(table.schema().fields()); + let lower = predicates.greater_than("id", Datum::Int(1)).unwrap(); + let upper = predicates.less_than("id", Datum::Int(3)).unwrap(); + let single = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(3) + .with_filter(lower.clone()) + .with_filter(upper.clone()) + .execute() + .await + .unwrap(); + assert_eq!(single.row_ids().unwrap().row_ids, vec![1]); + let batch = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(3) + .with_filter(lower) + .with_filter(upper) + .execute() + .await + .unwrap(); + assert_eq!(batch.len(), 2); + assert!(batch + .iter() + .all(|result| result.row_ids().unwrap().row_ids == vec![1])); +} + +#[tokio::test] +async fn repeated_options_keep_earlier_overrides() { + let table = de_vector_table().await.copy_with_options(HashMap::from([( + "fields.embedding.ivf.refine-factor".to_string(), + "invalid".to_string(), + )])); + let override_options = HashMap::from([("refine_factor".to_string(), "1".to_string())]); + let later = HashMap::from([("nprobe".to_string(), "1".to_string())]); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![0.0, 1.0]) + .with_limit(1) + .with_options(override_options.clone()) + .with_options(later.clone()) + .execute() + .await + .unwrap(); + assert_eq!(result.row_ids().unwrap().row_ids, vec![1]); + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![0.0, 1.0]]) + .with_limit(1) + .with_options(override_options) + .with_options(later) + .execute() + .await + .unwrap(); + assert_eq!(results[0].row_ids().unwrap().row_ids, vec![1]); +} + +#[test] +fn explicit_partition_filters_reject_data_fields() { + use crate::spec::{Datum, PredicateBuilder}; + let table = vector_test_table(); + let filter = PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(1)) + .unwrap(); + assert!(table + .new_vector_search_builder() + .with_partition_filter(filter.clone()) + .is_err()); + assert!(table + .new_batch_vector_search_builder() + .with_partition_filter(filter) + .is_err()); +} diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index e033f79fb..081f2a2a8 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -278,7 +278,7 @@ async fn execute_vindex_searches( vector_searches: Vec, source: S, file_name: String, - index_parallelism: usize, + context: VectorIndexSearchContext, guard: G, ) -> crate::Result>>> { let panic_context = if vector_searches.len() > 1 { @@ -288,9 +288,17 @@ async fn execute_vindex_searches( }; execute_global_index_with_guard(panic_context, guard, move || { let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options) - .with_batch_index_parallelism(index_parallelism); + .with_batch_index_parallelism(context.parallelism); reader - .visit_batch_vector_search(&vector_searches, |_| Ok(source)) + .visit_batch_vector_search_validated( + &vector_searches, + |_| Ok(source), + |metadata| { + context + .metrics + .check(RawVectorMetric::from_vindex(metadata.metric)) + }, + ) .map_err(|e| crate::Error::DataInvalid { message: format!("Failed to read vindex index file '{}': {}", file_name, e), source: Some(Box::new(e)), @@ -309,6 +317,45 @@ struct VectorSearchEvaluation<'a> { next_row_id: Option, } +#[derive(Clone)] +struct VectorMetricCheck { + persisted: Arc>>, + requested: Option, +} + +impl VectorMetricCheck { + fn check(&self, metric: RawVectorMetric) -> crate::Result<()> { + if self.requested.is_some_and(|requested| requested != metric) { + return Err(crate::Error::ConfigInvalid { + message: format!( + "Query vector metric {:?} does not match index metric {metric:?}", + self.requested + ), + }); + } + let mut persisted = + self.persisted + .lock() + .map_err(|error| crate::Error::UnexpectedError { + message: format!("Vector metric validation lock poisoned: {error}"), + source: None, + })?; + if persisted.is_some_and(|previous| previous != metric) { + return Err(crate::Error::DataInvalid { + message: "Cannot merge vector indexes with different metrics".to_string(), + source: None, + }); + } + *persisted = Some(metric); + Ok(()) + } +} + +struct VectorIndexSearchContext { + parallelism: usize, + metrics: VectorMetricCheck, +} + #[derive(Default)] struct IndexSearchTiming { permit_wait: Duration, @@ -454,6 +501,10 @@ async fn evaluate_batch_vector_search( .map(|entry| (entry, None)) .collect::>() }; + let metric_check = VectorMetricCheck { + persisted: Arc::new(std::sync::Mutex::new(None)), + requested: configured_raw_vector_metric_override(&search_options, field_name)?, + }; let mut permit_wait = Duration::ZERO; let mut file_reader_open = Duration::ZERO; let mut index_search = Duration::ZERO; @@ -516,6 +567,7 @@ async fn evaluate_batch_vector_search( let mut options = evaluation.table_options.clone(); options.extend(search_options.clone()); let input = evaluation.file_io.new_input(&path); + let metric_check = metric_check.clone(); async move { if let Some(local_filter) = shared_local_filter { let local_filter = Arc::new(local_filter); @@ -575,6 +627,8 @@ async fn evaluate_batch_vector_search( if let Some(start) = read_start { full_file_read = Some((start.elapsed(), data.len())); } + let metric = LuminaIndexMeta::deserialize(&io_meta.metadata)?.metric()?; + metric_check.check(RawVectorMetric::from_lumina(metric))?; execute_global_index_with_guard( "Lumina global-index batch search task failed", permit, @@ -618,7 +672,7 @@ async fn evaluate_batch_vector_search( vector_searches, source, file_name.clone(), - batch_index_parallelism, + VectorIndexSearchContext { parallelism: batch_index_parallelism, metrics: metric_check }, permit, ) .await?; @@ -651,7 +705,7 @@ async fn evaluate_batch_vector_search( vector_searches, Cursor::new(data), file_name.clone(), - batch_index_parallelism, + VectorIndexSearchContext { parallelism: batch_index_parallelism, metrics: metric_check }, permit, ) .await? @@ -727,7 +781,7 @@ async fn evaluate_batch_vector_search( evaluation, index_entries, field_id, - field_name, + &metric_check, vector_searches, merged, index_search_limit, @@ -766,12 +820,14 @@ async fn evaluate_batch_vector_search( evaluation.file_io, table_path, evaluation.table_options, + &vector_searches[0].options, index_entries, field_id, field_name, ) .await?; let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); + metric_check.check(metric)?; let (raw_results, raw_timing) = read_raw_batch_vector_search(table, vector_searches, &raw_ranges, metric).await?; if let Some(raw_timing) = raw_timing { @@ -965,7 +1021,7 @@ async fn maybe_rerank_indexed_batch_results( evaluation: VectorSearchEvaluation<'_>, index_entries: &[IndexManifestEntry], field_id: i32, - field_name: &str, + metric_check: &VectorMetricCheck, vector_searches: &[VectorSearch], results: Vec, index_search_limit: usize, @@ -1003,15 +1059,18 @@ async fn maybe_rerank_indexed_batch_results( let unique_candidates = union_candidates.len(); let raw_ranges = sorted_row_ids_to_row_ranges(union_candidates.iter())?; let metric_start = timing_enabled.then(Instant::now); + let field_name = &vector_searches[0].field_name; let metric = resolve_raw_vector_metric( evaluation.file_io, evaluation.table_path.trim_end_matches('/'), evaluation.table_options, + &vector_searches[0].options, index_entries, field_id, field_name, ) .await?; + metric_check.check(metric)?; let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); let (results, raw_timing) = @@ -1241,6 +1300,7 @@ async fn resolve_raw_vector_metric( file_io: &FileIO, table_path: &str, table_options: &HashMap, + search_options: &HashMap, index_entries: &[IndexManifestEntry], field_id: i32, field_name: &str, @@ -1320,13 +1380,19 @@ async fn resolve_raw_vector_metric( } } - configured_raw_vector_metric(table_options, field_name) + if let Some(metric) = configured_raw_vector_metric_override(search_options, field_name)? { + return Ok(metric); + } + Ok( + configured_raw_vector_metric_override(table_options, field_name)? + .unwrap_or(RawVectorMetric::L2), + ) } -fn configured_raw_vector_metric( +fn configured_raw_vector_metric_override( options: &HashMap, field_name: &str, -) -> crate::Result { +) -> crate::Result> { let direct_keys = [ format!("fields.{field_name}.distance.metric"), format!("fields.{field_name}.metric"), @@ -1337,7 +1403,7 @@ fn configured_raw_vector_metric( ]; for key in direct_keys { if let Some(value) = options.get(&key) { - return RawVectorMetric::parse(value); + return RawVectorMetric::parse(value).map(Some); } } @@ -1352,13 +1418,13 @@ fn configured_raw_vector_metric( }; if let Some(existing) = inferred { if existing != metric { - return Ok(RawVectorMetric::L2); + return Ok(None); } } else { inferred = Some(metric); } } - Ok(inferred.unwrap_or(RawVectorMetric::L2)) + Ok(inferred) } #[derive(Default)] diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs index 63e1ff801..87f8606c3 100644 --- a/crates/paimon/src/table/de_vector_read/tests.rs +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -432,16 +432,16 @@ fn test_configured_raw_vector_metric_precedence_and_conflict_default() { ); options.insert("metric".to_string(), "cosine".to_string()); assert_eq!( - configured_raw_vector_metric(&options, "embedding").unwrap(), - RawVectorMetric::InnerProduct + configured_raw_vector_metric_override(&options, "embedding").unwrap(), + Some(RawVectorMetric::InnerProduct) ); options.clear(); options.insert("foo.metric".to_string(), "cosine".to_string()); options.insert("bar.distance.metric".to_string(), "l2".to_string()); assert_eq!( - configured_raw_vector_metric(&options, "embedding").unwrap(), - RawVectorMetric::L2 + configured_raw_vector_metric_override(&options, "embedding").unwrap(), + None ); } @@ -465,6 +465,7 @@ async fn test_resolve_raw_vector_metric_uses_vindex_manifest_metadata() { &file_io, "memory:///test_table", &HashMap::new(), + &HashMap::new(), &[entry], 2, "embedding", @@ -503,6 +504,7 @@ async fn test_resolve_raw_vector_metric_falls_back_to_vindex_header() { &file_io, "memory:///test_table", &HashMap::new(), + &HashMap::new(), &[entry], 2, "embedding", @@ -1390,3 +1392,176 @@ async fn de_result_read_applies_scalar_filter_before_top_k() { assert_eq!(ids, vec![3, 2]); } + +#[tokio::test] +async fn raw_query_metric_options_precede_table_options() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_options = HashMap::from([( + "fields.embedding.distance.metric".to_string(), + "invalid".to_string(), + )]); + let options = HashMap::from([("ivf-flat.metric".to_string(), "cosine".to_string())]); + let metric = resolve_raw_vector_metric( + &file_io, + "memory:/table", + &table_options, + &options, + &[], + 2, + "embedding", + ) + .await + .unwrap(); + assert_eq!(metric, RawVectorMetric::Cosine); + let table_options = HashMap::from([("metric".to_string(), "inner_product".to_string())]); + let conflicting = HashMap::from([ + ("a.metric".to_string(), "cosine".to_string()), + ("b.metric".to_string(), "l2".to_string()), + ]); + let metric = resolve_raw_vector_metric( + &file_io, + "memory:/table", + &table_options, + &conflicting, + &[], + 2, + "embedding", + ) + .await + .unwrap(); + assert_eq!(metric, RawVectorMetric::InnerProduct); +} + +#[tokio::test] +async fn indexed_search_rejects_mixed_metrics_and_conflicting_query_metric() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let mut entries = Vec::new(); + for (ordinal, metric) in ["l2", "cosine"].into_iter().enumerate() { + let bytes = build_vindex_segment_bytes(metric); + let name = format!("metric-{ordinal}.idx"); + let mut entry = make_lumina_entry(&name, IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + entry.index_file.file_size = bytes.len() as i64; + let meta = entry.index_file.global_index_meta.as_mut().unwrap(); + meta.row_range_start = ordinal as i64 * 3; + meta.row_range_end = meta.row_range_start + 2; + file_io + .new_output(&format!("memory:///test_table/index/{name}")) + .unwrap() + .write(bytes::Bytes::from(bytes)) + .await + .unwrap(); + entries.push(entry); + } + let fields = vec![make_field(2, "embedding")]; + let options = HashMap::new(); + let query = VectorSearch::new(vec![1.0, 0.0], 10, "embedding".to_string()).unwrap(); + let error = evaluate_batch_vector_search( + eval_context(&file_io, &options, &fields, None), + &entries, + &[query.clone(), query.clone()], + ) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("Cannot merge vector indexes with different metrics"), + "{error}" + ); + let query = query.with_options(HashMap::from([( + "metric".to_string(), + "cosine".to_string(), + )])); + let error = evaluate_vector_search( + eval_context(&file_io, &options, &fields, None), + &entries[..1], + &query, + ) + .await + .unwrap_err(); + assert!( + error.to_string().contains("does not match index metric"), + "{error}" + ); +} + +#[tokio::test] +async fn filtered_rerank_rejects_metric_from_a_skipped_shard() { + use crate::spec::IndexManifest; + use crate::table::{TableCommit, TableWrite}; + use arrow_array::builder::ListBuilder; + let table = de_vector_table().await; + let element = Arc::new(ArrowField::new("item", ArrowDataType::Float32, false)); + let mut vectors = ListBuilder::new(Float32Builder::new()).with_field(element.clone()); + for values in [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]] { + vectors.values().append_slice(&values); + vectors.append(true); + } + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element), true), + ])), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(vectors.finish()), + ], + ) + .unwrap(); + let mut writer = TableWrite::new(&table, "metric-user".to_string()).unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + TableCommit::new(table.clone(), "metric-user".to_string()) + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let count = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .with_options(HashMap::from([( + "ivf-flat.distance.metric".to_string(), + "cosine".to_string(), + )])) + .execute() + .await + .unwrap(); + assert!(count > 0); + let manager = table.snapshot_manager(); + let snapshot = manager.get_latest_snapshot().await.unwrap().unwrap(); + let mut entries = IndexManifest::read( + table.file_io(), + &manager.manifest_path(snapshot.index_manifest().unwrap()), + ) + .await + .unwrap(); + entries.sort_by_key(|entry| { + entry + .index_file + .global_index_meta + .as_ref() + .map(|meta| meta.row_range_start) + }); + let query = VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()) + .unwrap() + .with_options(HashMap::from([( + "refine_factor".to_string(), + "2".to_string(), + )])) + .with_include_row_ids((3..6).collect()); + let evaluation = VectorSearchEvaluation { + table: Some(&table), + file_io: table.file_io(), + table_path: table.location(), + table_options: table.schema().options(), + schema_fields: table.schema().fields(), + next_row_id: snapshot.next_row_id(), + }; + let error = evaluate_vector_search(evaluation, &entries, &query) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("Cannot merge vector indexes with different metrics"), + "{error}" + ); +} diff --git a/crates/paimon/src/table/de_vector_scan.rs b/crates/paimon/src/table/de_vector_scan.rs index c5830a65b..e99c57586 100644 --- a/crates/paimon/src/table/de_vector_scan.rs +++ b/crates/paimon/src/table/de_vector_scan.rs @@ -18,6 +18,8 @@ //! Plans global-index vector searches against one snapshot and resolves scalar pre-filters. use crate::spec::{CoreOptions, IndexManifest, IndexManifestEntry, Predicate, ROW_ID_FIELD_NAME}; +use crate::table::bucket_filter::split_partition_and_data_predicates; +use crate::table::partition_filter::PartitionFilter; use crate::table::vector_scan::Scan; use crate::table::Table; use crate::vindex::vector_search_timing_enabled; @@ -281,6 +283,24 @@ impl Scan for DeVectorScan { } None => Vec::new(), }; + if let Some(filter) = &self.filter { + let (partition, _) = split_partition_and_data_predicates( + filter.clone(), + plan.table.schema().fields(), + plan.table.schema().partition_keys(), + ); + if let Some(partition) = partition { + let fields = plan.table.schema().partition_fields(); + let partition = PartitionFilter::from_predicate(partition, &fields); + let mut entries = Vec::with_capacity(plan.index_entries.len()); + for entry in plan.index_entries { + if partition.matches_entry(&entry.partition)? { + entries.push(entry); + } + } + plan.index_entries = entries; + } + } if let Some(timing) = &mut plan.timing { timing.manifest = manifest_start.map_or(Duration::ZERO, |start| start.elapsed()); } diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 5e37679b8..2f750b736 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -289,7 +289,8 @@ pub(crate) fn build_indexed_splits( .with_bucket(source.bucket()) .with_bucket_path(source.bucket_path().to_string()) .with_total_buckets(source.total_buckets()) - .with_data_files(vec![file_meta]); + .with_data_files(vec![file_meta]) + .with_raw_convertible(false); if let Some(df) = deletion_file { builder = builder.with_data_deletion_files(vec![Some(df)]); } diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index d50980abe..d68ff2a30 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -943,7 +943,27 @@ impl DataSplit { /// `IndexedSplit` (type 3) wrapping the DataSplit body plus the ranges. 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) + } + + /// Serialize aligned vector scores using Java's IndexedSplit wire format. + /// Scores belong to the selection, in ascending expanded-range order. + pub fn serialize_split_v1_with_scores(&self, scores: Option<&[f32]>) -> crate::Result> { 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, + }); + } + } let mut out = Vec::new(); out.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); @@ -963,7 +983,16 @@ impl DataSplit { out.extend_from_slice(&r.from().to_be_bytes()); out.extend_from_slice(&r.to().to_be_bytes()); } - out.push(0); // scores = null (vector scores are not modeled on the Rust split) + match scores { + None => out.push(0), + Some(scores) => { + out.push(1); + out.extend_from_slice(&(scores.len() as i32).to_be_bytes()); + for score in scores { + out.extend_from_slice(&score.to_be_bytes()); + } + } + } } } Ok(out) @@ -2399,6 +2428,34 @@ mod tests { assert_eq!(split.serialize_split_v1().unwrap(), expected); } + #[test] + fn serialize_indexed_scores_match_java_wire_layout() { + let split = v1_data_split_builder() + .with_row_ranges(vec![RowRange::new(1, 4), RowRange::new(11, 13)]) + .build() + .unwrap(); + let scores = [1.0f32, 0.75, 0.5, 0.25, 0.0, -0.5, -1.0]; + let mut expected = split.serialize_split_v1().unwrap(); + assert_eq!(expected.pop(), Some(0)); + expected.push(1); + expected.extend_from_slice(&7i32.to_be_bytes()); + for score in scores { + expected.extend_from_slice(&score.to_be_bytes()); + } + assert_eq!( + split.serialize_split_v1_with_scores(Some(&scores)).unwrap(), + expected + ); + assert!(split + .serialize_split_v1_with_scores(Some(&scores[..6])) + .is_err()); + assert!(v1_data_split_builder() + .build() + .unwrap() + .serialize_split_v1_with_scores(Some(&scores)) + .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. diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 176d75f9c..8e23077dd 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -20,6 +20,7 @@ use crate::spec::{CoreOptions, Predicate}; 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 std::collections::HashMap; @@ -61,7 +62,7 @@ impl<'a> VectorSearchBuilder<'a> { } pub fn with_options(&mut self, options: HashMap) -> &mut Self { - self.options = options; + self.options.extend(options); self } @@ -79,10 +80,16 @@ impl<'a> VectorSearchBuilder<'a> { /// the surviving files, so results stay exact. Sub-file row-range narrowing is /// not performed; a surviving file is re-read in full for the residual. pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { - self.filter = Some(filter); + add_filter(&mut self.filter, filter); self } + /// Add a partition-only predicate expressed against the table schema. + pub fn with_partition_filter(&mut self, filter: Predicate) -> crate::Result<&mut Self> { + validate_partition_filter(self.table, &filter)?; + Ok(self.with_filter(filter)) + } + /// Create a query-independent scan. Only the vector column must be configured. pub fn new_scan(&self) -> crate::Result { CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; diff --git a/crates/paimon/src/table/vector_search_common.rs b/crates/paimon/src/table/vector_search_common.rs index 61535a8ad..321bd9149 100644 --- a/crates/paimon/src/table/vector_search_common.rs +++ b/crates/paimon/src/table/vector_search_common.rs @@ -18,7 +18,9 @@ //! Shared vector-search option validation, index I/O helpers, and result ordering. use crate::lumina::is_lumina_index_type; -use crate::spec::{CoreOptions, DataField, ROW_ID_FIELD_NAME}; +use crate::spec::{ + field_idx_to_partition_idx, CoreOptions, DataField, Predicate, ROW_ID_FIELD_NAME, +}; use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; use crate::table::read_builder::resolve_projected_fields; use crate::table::Table; @@ -28,6 +30,28 @@ use arrow_array::{Int64Array, RecordBatch}; use arrow_select::interleave::interleave_record_batch; use std::collections::HashMap; +/// Explicit partition filters use table-field indices, like scalar filters. +/// Reject mixed predicates rather than silently discarding their data conjuncts. +pub(super) fn validate_partition_filter(table: &Table, filter: &Predicate) -> crate::Result<()> { + let mapping = + field_idx_to_partition_idx(table.schema().fields(), table.schema().partition_keys()); + if table.schema().partition_keys().is_empty() || !filter.references_only_mapped_fields(&mapping) + { + return Err(crate::Error::ConfigInvalid { + message: "Partition filter must reference only partition keys of a partitioned table" + .to_string(), + }); + } + Ok(()) +} + +pub(super) fn add_filter(current: &mut Option, filter: Predicate) { + *current = Some(match current.take() { + Some(previous) => Predicate::and(vec![previous, filter]), + None => filter, + }); +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(super) enum VectorIndexBackend { Lumina, diff --git a/crates/paimon/src/table/vector_search_result.rs b/crates/paimon/src/table/vector_search_result.rs index 0a69ba452..b2fb1f6b3 100644 --- a/crates/paimon/src/table/vector_search_result.rs +++ b/crates/paimon/src/table/vector_search_result.rs @@ -156,6 +156,21 @@ impl SearchResult { } } + /// 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]> { From 2a2692f5b85ab9317f4257982b37a7af3945ab6a Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 9 Oct 2026 22:57:09 +0800 Subject: [PATCH 2/3] fix: align scalar vector filtering with Java --- bindings/python/src/vector_search.rs | 4 +- bindings/python/tests/test_vector_search.py | 5 + crates/paimon/src/spec/core_options.rs | 10 + crates/paimon/src/table/de_vector_read.rs | 317 +++++++++++++----- .../paimon/src/table/de_vector_read/tests.rs | 2 + crates/paimon/src/table/de_vector_scan.rs | 17 +- .../paimon/src/table/de_vector_scan/tests.rs | 118 ++++++- .../paimon/src/table/global_index_scanner.rs | 42 +++ .../table/global_index_scanner/composite.rs | 3 + .../table/global_index_scanner/evaluator.rs | 90 +++-- .../src/table/global_index_scanner/tests.rs | 120 +++++++ .../paimon/src/table/vector_search_builder.rs | 16 +- 12 files changed, 613 insertions(+), 131 deletions(-) diff --git a/bindings/python/src/vector_search.rs b/bindings/python/src/vector_search.rs index e5e99eaa0..49fe05b61 100644 --- a/bindings/python/src/vector_search.rs +++ b/bindings/python/src/vector_search.rs @@ -114,8 +114,8 @@ impl PyVectorSearchBuilder { slf.config.limit = Some(limit); slf } - fn with_query_vector(mut slf: PyRefMut<'_, Self>, vectors: Vec) -> PyRefMut<'_, Self> { - slf.query = Some(vectors); + fn with_query_vector(mut slf: PyRefMut<'_, Self>, vector: Vec) -> PyRefMut<'_, Self> { + slf.query = Some(vector); slf } fn with_option(mut slf: PyRefMut<'_, Self>, key: String, value: String) -> PyRefMut<'_, Self> { diff --git a/bindings/python/tests/test_vector_search.py b/bindings/python/tests/test_vector_search.py index 21e3947e6..f167b6df5 100644 --- a/bindings/python/tests/test_vector_search.py +++ b/bindings/python/tests/test_vector_search.py @@ -120,3 +120,8 @@ def test_invalid_query_is_reported_before_planning(table, query): builder.new_vector_search_read() with pytest.raises(Exception, match='Query vectors'): table.new_batch_vector_search_builder().with_vector_column('embedding').with_limit(1).with_query_vectors([]).execute_batch_local() + + +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() diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index d740b08ec..b48208f23 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -25,6 +25,7 @@ pub(crate) const QUERY_AUTH_ENABLED_OPTION: &str = "query-auth.enabled"; const DATA_EVOLUTION_ENABLED_OPTION: &str = "data-evolution.enabled"; const DATA_EVOLUTION_NESTED_FIELD_ENABLED_OPTION: &str = "data-evolution.nested-field.enabled"; const FILE_INDEX_READ_ENABLED_OPTION: &str = "file-index.read.enabled"; +const GLOBAL_INDEX_FILTER_REFINE_FROM_DATA_OPTION: &str = "global-index.filter.refine-from-data"; const GLOBAL_INDEX_ENABLED_OPTION: &str = "global-index.enabled"; const GLOBAL_INDEX_SEARCH_MODE_OPTION: &str = "global-index.search-mode"; const SCALAR_INDEX_SEARCH_MODE_OPTION: &str = "scalar-index.search-mode"; @@ -984,6 +985,15 @@ impl<'a> CoreOptions<'a> { .unwrap_or(true) } + /// Verify candidate-only scalar index answers against data before vector + /// ranking (Java default: true). When false, those candidates are excluded. + pub fn global_index_filter_refine_from_data(&self) -> bool { + self.options + .get(GLOBAL_INDEX_FILTER_REFINE_FROM_DATA_OPTION) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(true) + } + /// Whether index files are stored in the bucket data-file directory rather /// than the table `index/` directory (option `index-file-in-data-file-dir`, /// default false). diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index 081f2a2a8..6e65d1331 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -22,12 +22,12 @@ use crate::lumina::reader::LuminaVectorGlobalIndexReader; use crate::lumina::{LuminaIndexMeta, LuminaVectorMetric}; use crate::spec::{ row_id_data_field, CoreOptions, DataField, DataType, FileKind, GlobalIndexSearchMode, - IndexFileMeta, IndexManifestEntry, ROW_ID_FIELD_NAME, + IndexFileMeta, IndexManifestEntry, Predicate, ROW_ID_FIELD_NAME, }; -use crate::table::de_vector_scan::DeVectorScanPlan; +use crate::table::de_vector_scan::{matching_row_ids_for_filter, DeVectorScanPlan}; use crate::table::global_index_scanner::{ deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, - unindexed_ranges_for_global_index_entries, RowRangeIndex, + unindexed_ranges_for_global_index_entries, GlobalIndexScanner, RowRangeIndex, }; use crate::table::index_file_path::IndexFileLocation; use crate::table::pk_vector_position_read::SEARCH_SCORE_COLUMN; @@ -188,6 +188,7 @@ impl Read for DeVectorRead { table, index_entries, include_row_ids, + filter, next_row_id, timing, skip_search, @@ -231,6 +232,7 @@ impl Read for DeVectorRead { table_options: pinned_table.schema().options(), schema_fields: pinned_table.schema().fields(), next_row_id, + filter: filter.as_ref(), }, &index_entries, &vector_searches, @@ -315,6 +317,7 @@ struct VectorSearchEvaluation<'a> { table_options: &'a HashMap, schema_fields: &'a [DataField], next_row_id: Option, + filter: Option<&'a Predicate>, } #[derive(Clone)] @@ -377,6 +380,84 @@ async fn evaluate_vector_search( crate::table::vector_search_common::take_only_result(results, "vector search")?.to_row_ranges() } +/// None means the scalar index cannot evaluate the filter: Java routes all +/// vector-index-covered ranges to raw scoring in that case, even in FAST mode. +async fn scalar_matched_row_ids( + evaluation: VectorSearchEvaluation<'_>, + entries: &[IndexManifestEntry], + vector_ranges: &[RowRange], + filter: &Predicate, +) -> crate::Result> { + let core = CoreOptions::new(evaluation.table_options); + // Java search pre-filters use per-column coverage and exclude multi-field + // definitions, even though ordinary table scans can use composite indexes. + let scalar_entries: Vec<_> = entries + .iter() + .filter(|entry| { + entry + .index_file + .global_index_meta + .as_ref() + .is_some_and(|meta| meta.extra_field_ids.as_ref().is_none_or(Vec::is_empty)) + }) + .cloned() + .collect(); + let scanner = GlobalIndexScanner::create_with_fm_options( + evaluation.file_io, + evaluation.table_path, + core.global_index_thread_num()?, + core.btree_index_fallback_scan_max_size()?, + core.btree_index_data_block_cache_size()?, + core.bitmap_index_fallback_scan_max_size()?, + &scalar_entries, + evaluation.schema_fields, + if scalar_entries + .iter() + .any(|entry| entry.index_file.index_type.eq_ignore_ascii_case("fm")) + { + crate::fm_index::FMOptions::from_options(evaluation.table_options)?.read + } else { + crate::fm_index::FMReadOptions::default() + }, + )?; + let Some(scanner) = scanner else { + return Ok(None); + }; + let Some((ranges, exact)) = scanner.matching_ranges_with_exactness(filter).await? else { + return Ok(None); + }; + let ranges = intersect_sorted_ranges(&ranges, vector_ranges); + if exact { + let mut row_ids = RoaringTreemap::new(); + for range in ranges { + let start = u64::try_from(range.from()).map_err(|_| crate::Error::DataInvalid { + message: "Scalar vector index returned a negative row ID".to_string(), + source: None, + })?; + let end = u64::try_from(range.to()).map_err(|_| crate::Error::DataInvalid { + message: "Scalar vector index returned a negative row ID".to_string(), + source: None, + })?; + row_ids.insert_range(start..=end); + } + return Ok(Some(row_ids)); + } + if !core.global_index_filter_refine_from_data() { + log::warn!("Scalar global index only provides candidates for {filter:?}; global-index.filter.refine-from-data is false, so indexed candidates are excluded"); + return Ok(Some(RoaringTreemap::new())); + } + if ranges.is_empty() { + return Ok(Some(RoaringTreemap::new())); + } + let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { + message: "Scalar vector filter refinement requires table context".to_string(), + source: None, + })?; + Ok(Some( + matching_row_ids_for_filter(table, filter, Some(ranges)).await?, + )) +} + async fn evaluate_batch_vector_search( evaluation: VectorSearchEvaluation<'_>, index_entries: &[IndexManifestEntry], @@ -435,6 +516,49 @@ async fn evaluate_batch_vector_search( return Ok(vec![ScoredRowIds::empty(); vector_searches.len()]); } + let vector_ranges = merge_row_ranges( + vector_entries + .iter() + .filter_map(|entry| { + entry + .index_file + .global_index_meta + .as_ref() + .map(|meta| RowRange::new(meta.row_range_start, meta.row_range_end)) + }) + .collect(), + ); + let mut indexed_searches = Cow::Borrowed(vector_searches); + let mut route_indexed_to_raw = false; + // Partition predicates have already pruned manifest entries. They do not + // make an otherwise exact scalar lookup candidate-only. + let scalar_filter = evaluation.filter.and_then(|filter| match evaluation.table { + Some(table) => { + let data = crate::table::bucket_filter::split_partition_and_data_predicates( + filter.clone(), + table.schema().fields(), + table.schema().partition_keys(), + ) + .1; + (!data.is_empty()).then(|| Predicate::and(data)) + } + None => Some(filter.clone()), + }); + if let Some(filter) = scalar_filter + .as_ref() + .filter(|_| !vector_entries.is_empty()) + { + match scalar_matched_row_ids(evaluation, index_entries, &vector_ranges, filter).await? { + Some(row_ids) => { + let row_ids = Arc::new(row_ids); + for search in indexed_searches.to_mut() { + search.set_shared_include_row_ids(Arc::clone(&row_ids)); + } + } + None => route_indexed_to_raw = true, + } + } + let deletion_vector_start = timing_enabled.then(Instant::now); let deleted_row_index = if core_options.data_evolution_enabled() { match evaluation.table { @@ -467,9 +591,9 @@ async fn evaluate_batch_vector_search( let index_search_limit = indexed_search_limit(max_limit, refine_factor)?; let vector_entry_count = vector_entries.len(); - let vector_search_plans = if let Some(include_row_ids) = - shared_batch_include_row_ids(vector_searches) - { + let vector_search_plans = if route_indexed_to_raw { + Vec::new() + } else if let Some(include_row_ids) = shared_batch_include_row_ids(&indexed_searches) { let ranges = vector_entries .iter() .map(|entry| { @@ -512,7 +636,7 @@ async fn evaluate_batch_vector_search( let mut refine = Duration::ZERO; let mut raw_fallback = Duration::ZERO; let mut merged = vec![ScoredRowIds::empty(); vector_searches.len()]; - if !vector_entries.is_empty() { + if !vector_search_plans.is_empty() { let index_search_start = timing_enabled.then(Instant::now); let concurrency = core_options.global_index_thread_num()?; if concurrency > tokio::sync::Semaphore::MAX_PERMITS { @@ -560,7 +684,7 @@ async fn evaluate_batch_vector_search( deleted_row_index.as_ref(), ) .min(i32::MAX as usize); - let mut vector_searches = vector_searches.to_vec(); + let mut vector_searches = indexed_searches.to_vec(); for vector_search in &mut vector_searches { vector_search.limit = index_limit; } @@ -790,7 +914,7 @@ async fn evaluate_batch_vector_search( refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); } - if search_mode != GlobalIndexSearchMode::Fast { + if route_indexed_to_raw || search_mode != GlobalIndexSearchMode::Fast { let raw_fallback_start = timing_enabled.then(Instant::now); let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { @@ -802,7 +926,7 @@ async fn evaluate_batch_vector_search( Vec::new() }; let field_ids = HashSet::from([field_id]); - let raw_ranges = unindexed_ranges_for_global_index_entries( + let mut raw_ranges = unindexed_ranges_for_global_index_entries( index_entries, &field_ids, search_mode, @@ -810,6 +934,10 @@ async fn evaluate_batch_vector_search( &detail_ranges, is_vector_global_index_file, ); + if route_indexed_to_raw { + raw_ranges.extend(vector_ranges); + raw_ranges = merge_row_ranges(raw_ranges); + } if !raw_ranges.is_empty() { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { message: "Vector raw search requires table context".to_string(), @@ -828,8 +956,14 @@ async fn evaluate_batch_vector_search( .await?; let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); metric_check.check(metric)?; - let (raw_results, raw_timing) = - read_raw_batch_vector_search(table, vector_searches, &raw_ranges, metric).await?; + let (raw_results, raw_timing) = read_raw_batch_vector_search( + table, + vector_searches, + &raw_ranges, + metric, + evaluation.filter, + ) + .await?; if let Some(raw_timing) = raw_timing { log::debug!( target: "paimon::vector_search", @@ -1074,7 +1208,7 @@ async fn maybe_rerank_indexed_batch_results( let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); let (results, raw_timing) = - read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric).await?; + read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric, None).await?; if let (Some(total_start), Some(raw_timing)) = (total_start, raw_timing) { log::debug!( target: "paimon::vector_search", @@ -1305,88 +1439,107 @@ async fn resolve_raw_vector_metric( field_id: i32, field_name: &str, ) -> crate::Result { + let check = VectorMetricCheck { + persisted: Arc::new(std::sync::Mutex::new(None)), + requested: configured_raw_vector_metric_override(search_options, field_name)?, + }; + let mut persisted = None; for entry in index_entries { - if entry.kind != FileKind::Add { + if entry.kind != FileKind::Add + || !entry + .index_file + .global_index_meta + .as_ref() + .is_some_and(|meta| meta.index_field_id == field_id) + { continue; } - let Some(global_meta) = entry.index_file.global_index_meta.as_ref() else { - continue; - }; - if global_meta.index_field_id != field_id { - continue; + if let Some(metric) = read_index_vector_metric(file_io, table_path, entry).await? { + check.check(metric)?; + persisted = Some(metric); } - let Some(backend) = VectorIndexBackend::from_index_type(&entry.index_file.index_type) - else { - continue; - }; - match backend { - VectorIndexBackend::Lumina => { - if let Some(index_meta) = global_meta.index_meta.as_ref() { - if !index_meta.is_empty() { - let metric = LuminaIndexMeta::deserialize(index_meta)?.metric()?; - return Ok(RawVectorMetric::from_lumina(metric)); - } + } + if let Some(metric) = persisted { + return Ok(metric); + } + if let Some(metric) = check.requested { + return Ok(metric); + } + Ok( + configured_raw_vector_metric_override(table_options, field_name)? + .unwrap_or(RawVectorMetric::L2), + ) +} + +/// Rerouted raw searches still validate each shard's persisted metric. +async fn read_index_vector_metric( + file_io: &FileIO, + table_path: &str, + entry: &IndexManifestEntry, +) -> crate::Result> { + let Some(global_meta) = entry.index_file.global_index_meta.as_ref() else { + return Ok(None); + }; + let Some(backend) = VectorIndexBackend::from_index_type(&entry.index_file.index_type) else { + return Ok(None); + }; + match backend { + VectorIndexBackend::Lumina => { + if let Some(index_meta) = global_meta.index_meta.as_ref() { + if !index_meta.is_empty() { + let metric = LuminaIndexMeta::deserialize(index_meta)?.metric()?; + return Ok(Some(RawVectorMetric::from_lumina(metric))); } } - VectorIndexBackend::Vindex => { - if let Some(index_meta) = global_meta.index_meta.as_ref() { - if let Ok(options) = - serde_json::from_slice::>(index_meta) - { - if let Some(metric) = options.get("metric") { - if let Some(metric) = - RawVectorMetric::parse_normalized(&normalize_metric(metric)) - { - return Ok(metric); - } + } + VectorIndexBackend::Vindex => { + if let Some(index_meta) = global_meta.index_meta.as_ref() { + if let Ok(options) = serde_json::from_slice::>(index_meta) { + if let Some(metric) = options.get("metric") { + if let Some(metric) = + RawVectorMetric::parse_normalized(&normalize_metric(metric)) + { + return Ok(Some(metric)); } } } - let path = IndexFileLocation::Global { table_path }.resolve( - &entry.index_file.file_name, - entry.index_file.external_path.as_deref(), - ); - let input = file_io.new_input(&path)?; - let read_error = |e| crate::Error::DataInvalid { + } + let path = IndexFileLocation::Global { table_path }.resolve( + &entry.index_file.file_name, + entry.index_file.external_path.as_deref(), + ); + let input = file_io.new_input(&path)?; + let read_error = |e| crate::Error::DataInvalid { + message: format!( + "Failed to read vindex index file '{}' for raw search metric: {}", + entry.index_file.file_name, e + ), + source: Some(Box::new(e)), + }; + let header_size = if entry.index_file.file_size > 0 { + (entry.index_file.file_size as u64).min(DISKANN_HEADER_SIZE as u64) + } else { + input + .metadata() + .await + .map_err(&read_error)? + .size + .min(DISKANN_HEADER_SIZE as u64) + }; + let file_reader = input.reader().await.map_err(&read_error)?; + let bytes = file_reader.read(0..header_size).await.map_err(read_error)?; + let reader = + VIndexReader::open(Cursor::new(bytes)).map_err(|e| crate::Error::DataInvalid { message: format!( - "Failed to read vindex index file '{}' for raw search metric: {}", - entry.index_file.file_name, e + "Failed to open paimon-vindex-core reader for raw search metric: {}", + e ), source: Some(Box::new(e)), - }; - let header_size = if entry.index_file.file_size > 0 { - (entry.index_file.file_size as u64).min(DISKANN_HEADER_SIZE as u64) - } else { - input - .metadata() - .await - .map_err(&read_error)? - .size - .min(DISKANN_HEADER_SIZE as u64) - }; - let file_reader = input.reader().await.map_err(&read_error)?; - let bytes = file_reader.read(0..header_size).await.map_err(read_error)?; - let reader = VIndexReader::open(Cursor::new(bytes)).map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to open paimon-vindex-core reader for raw search metric: {}", - e - ), - source: Some(Box::new(e)), - } })?; - return Ok(RawVectorMetric::from_vindex(reader.metadata().metric)); - } + return Ok(Some(RawVectorMetric::from_vindex(reader.metadata().metric))); } } - - if let Some(metric) = configured_raw_vector_metric_override(search_options, field_name)? { - return Ok(metric); - } - Ok( - configured_raw_vector_metric_override(table_options, field_name)? - .unwrap_or(RawVectorMetric::L2), - ) + Ok(None) } fn configured_raw_vector_metric_override( @@ -1444,6 +1597,7 @@ async fn read_raw_batch_vector_search( vector_searches: &[VectorSearch], raw_ranges: &[RowRange], metric: RawVectorMetric, + filter: Option<&Predicate>, ) -> crate::Result<(Vec, Option)> { let timing_enabled = vector_search_timing_enabled(); let total_start = timing_enabled.then(Instant::now); @@ -1475,6 +1629,9 @@ async fn read_raw_batch_vector_search( read_builder .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])? .with_row_ranges(raw_ranges); + if let Some(filter) = filter { + read_builder.with_filter(filter.clone()); + } let plan = read_builder.new_scan().plan().await?; let plan_elapsed = plan_start.map_or(Duration::ZERO, |start| start.elapsed()); let split_count = plan.splits().len(); diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs index 87f8606c3..d1fd1ba19 100644 --- a/crates/paimon/src/table/de_vector_read/tests.rs +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -61,6 +61,7 @@ fn eval_context<'a>( table_options: options, schema_fields: fields, next_row_id, + filter: None, } } @@ -1554,6 +1555,7 @@ async fn filtered_rerank_rejects_metric_from_a_skipped_shard() { table_options: table.schema().options(), schema_fields: table.schema().fields(), next_row_id: snapshot.next_row_id(), + filter: None, }; let error = evaluate_vector_search(evaluation, &entries, &query) .await diff --git a/crates/paimon/src/table/de_vector_scan.rs b/crates/paimon/src/table/de_vector_scan.rs index e99c57586..ac154ef2e 100644 --- a/crates/paimon/src/table/de_vector_scan.rs +++ b/crates/paimon/src/table/de_vector_scan.rs @@ -21,7 +21,7 @@ use crate::spec::{CoreOptions, IndexManifest, IndexManifestEntry, Predicate, ROW use crate::table::bucket_filter::split_partition_and_data_predicates; use crate::table::partition_filter::PartitionFilter; use crate::table::vector_scan::Scan; -use crate::table::Table; +use crate::table::{RowRange, Table}; use crate::vindex::vector_search_timing_enabled; use arrow_array::{Array, Int64Array}; use futures::TryStreamExt; @@ -54,14 +54,18 @@ fn same_vector_search_table(left: &Table, right: &Table) -> bool { && left.branch() == right.branch() } -async fn matching_row_ids_for_filter( +pub(super) async fn matching_row_ids_for_filter( table: &Table, filter: &Predicate, + ranges: Option>, ) -> crate::Result { let mut read_builder = table.new_read_builder(); read_builder .with_projection(&[ROW_ID_FIELD_NAME])? .with_filter(filter.clone()); + if let Some(ranges) = ranges { + read_builder.with_row_ranges(ranges); + } let plan = read_builder.new_scan().plan().await?; let read = read_builder.new_read()?; let mut stream = read.to_arrow(plan.splits())?; @@ -124,7 +128,7 @@ impl Table { }); }; let table = self.copy_with_resolved_snapshot(&snapshot).await?; - let include_row_ids = matching_row_ids_for_filter(&table, &filter).await?; + let include_row_ids = matching_row_ids_for_filter(&table, &filter, None).await?; Ok(PreparedVectorSearchFilter { table, include_row_ids: Arc::new(include_row_ids), @@ -162,6 +166,7 @@ pub(super) struct DeVectorScanPlan { pub(super) table: Table, pub(super) index_entries: Vec, pub(super) include_row_ids: Option>, + pub(super) filter: Option, pub(super) timing: Option, pub(super) next_row_id: Option, pub(super) skip_search: bool, @@ -174,6 +179,7 @@ impl DeVectorScanPlan { skip_search: true, index_entries: Vec::new(), include_row_ids: None, + filter: None, next_row_id: None, timing, } @@ -256,14 +262,11 @@ impl Scan for DeVectorScan { }; plan.next_row_id = snapshot.next_row_id(); + plan.filter = self.filter.clone(); plan.include_row_ids = if let Some(prepared) = self.prepared_filter.as_ref() { Some(Arc::clone(prepared.include_row_ids())) } else if let Some(include_row_ids) = self.include_row_ids.as_ref() { Some(Arc::clone(include_row_ids)) - } else if let Some(filter) = self.filter.as_ref() { - Some(Arc::new( - matching_row_ids_for_filter(&plan.table, filter).await?, - )) } else { None }; diff --git a/crates/paimon/src/table/de_vector_scan/tests.rs b/crates/paimon/src/table/de_vector_scan/tests.rs index 2016e6974..60427a881 100644 --- a/crates/paimon/src/table/de_vector_scan/tests.rs +++ b/crates/paimon/src/table/de_vector_scan/tests.rs @@ -86,7 +86,7 @@ async fn reader_reuses_queries_without_leaking_plan_filters() { let filter = id_gt_filter(&table, 1); let empty_filter = id_gt_filter(&table, 99); // One reader executes filtered, empty, and unfiltered plans. The plan owns - // the allow-list; neither it nor the empty result can alter later queries. + // the predicate; neither it nor the empty result can alter later queries. for (filter, expected) in [ (Some(&filter), [vec![2, 1], vec![1, 2]]), (Some(&empty_filter), [vec![], vec![]]), @@ -362,3 +362,119 @@ async fn de_scalar_filter_applies_to_unindexed_raw_fallback() { assert_eq!(result.row_ids().unwrap().row_ids, vec![3]); } + +#[tokio::test] +async fn unevaluable_scalar_filter_routes_covered_ranges_to_raw_in_every_mode() { + let table = de_vector_table().await; + for mode in ["fast", "full", "detail"] { + let table = table.copy_with_options(HashMap::from([( + "vector-index.search-mode".to_string(), + mode.to_string(), + )])); + let filter = id_gt_filter(&table, 0); + let mut plan = DeVectorScan::new(&table, Some(&filter), None, None) + .plan() + .await + .unwrap(); + // Keep the committed IVF shards but remove the scalar index from this + // plan. A data-derived ANN allow-list can lose hits with nprobe=1. + plan.index_entries + .retain(|entry| entry.index_file.index_type == "ivf-flat"); + let read = DeVectorRead::new( + "embedding", + &[&[1.0, 0.0], &[0.0, 1.0]], + 3, + &HashMap::from([("ivf.nprobe".to_string(), "1".to_string())]), + ) + .unwrap(); + let results = read.read(plan).await.unwrap(); + for result in results { + let mut ids = result.row_ids().unwrap().row_ids.clone(); + ids.sort_unstable(); + assert_eq!(ids, vec![0, 1, 2], "{mode}"); + } + } +} + +#[tokio::test] +async fn exact_and_partial_scalar_indexes_obey_refinement_switch() { + let table = de_vector_table().await; + let predicate = crate::spec::PredicateBuilder::new(table.schema().fields()); + let exact = id_gt_filter(&table, 1); + let partial = Predicate::and(vec![ + exact.clone(), + predicate.is_not_null("embedding").unwrap(), + ]); + let contradiction = Predicate::and(vec![exact.clone(), Predicate::negate(exact.clone())]); + for mode in ["fast", "full", "detail"] { + for refine in [false, true] { + let table = table.copy_with_options(HashMap::from([ + ("vector-index.search-mode".to_string(), mode.to_string()), + ( + "global-index.filter.refine-from-data".to_string(), + refine.to_string(), + ), + ])); + for (filter, expected) in [ + (&exact, 2), + (&partial, if refine { 2 } else { 0 }), + (&contradiction, 0), + ] { + let results = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 0.0], vec![0.0, 1.0]]) + .with_limit(3) + .with_options(HashMap::from([("ivf.nprobe".to_string(), "2".to_string())])) + .with_filter(filter.clone()) + .execute() + .await + .unwrap(); + assert_eq!(results.len(), 2); + for result in results { + assert_eq!( + result.row_ids().unwrap().len(), + expected, + "{mode}, refine={refine}" + ); + } + } + } + } +} + +#[tokio::test] +async fn scalar_filter_ignores_composite_index_definitions() { + let table = de_vector_table().await.copy_with_options(HashMap::from([ + ("vector-index.search-mode".to_string(), "fast".to_string()), + ( + "global-index.filter.refine-from-data".to_string(), + "false".to_string(), + ), + ])); + let filter = id_gt_filter(&table, 0); + let mut plan = DeVectorScan::new(&table, Some(&filter), None, None) + .plan() + .await + .unwrap(); + let embedding_id = + crate::table::find_field_id_by_name(table.schema().fields(), "embedding").unwrap(); + for entry in &mut plan.index_entries { + if entry.index_file.index_type == "btree" { + // A multi-field definition cannot supply Java's per-column scalar + // coverage, regardless of whether its leading field is indexed. + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .extra_field_ids = Some(vec![embedding_id]); + } + } + let results = DeVectorRead::new("embedding", &[&[1.0, 0.0]], 3, &HashMap::new()) + .unwrap() + .read(plan) + .await + .unwrap(); + assert_eq!(results[0].row_ids().unwrap().len(), 3); +} diff --git a/crates/paimon/src/table/global_index_scanner.rs b/crates/paimon/src/table/global_index_scanner.rs index b9f7f0586..71dd09a0b 100644 --- a/crates/paimon/src/table/global_index_scanner.rs +++ b/crates/paimon/src/table/global_index_scanner.rs @@ -366,6 +366,48 @@ impl GlobalIndexScanner { ) } + /// Scalar candidates and Java's exactness contract for searches which rank + /// rows before applying the row predicate. + pub(super) async fn matching_ranges_with_exactness( + &self, + predicate: &Predicate, + ) -> Result, bool)>> { + let Some(result) = self.evaluate(predicate).await? else { + return Ok(None); + }; + let exact = result.fully_evaluated + && self.is_exact_predicate(predicate, &result.evaluated_field_ids); + Ok(Some((result.row_ranges, exact))) + } + + fn is_exact_predicate( + &self, + predicate: &Predicate, + evaluated_fields: &std::collections::HashSet, + ) -> bool { + use crate::spec::PredicateOperator; + match predicate { + Predicate::Leaf { column, op, .. } => { + predicates::is_sorted_global_index_supported_op(*op) + && !matches!( + op, + PredicateOperator::Contains + | PredicateOperator::EndsWith + | PredicateOperator::Like + ) + && crate::table::find_field_id_by_name(&self.schema_fields, column) + .is_some_and(|id| evaluated_fields.contains(&id)) + } + Predicate::And(children) | Predicate::Or(children) => children + .iter() + .all(|child| self.is_exact_predicate(child, evaluated_fields)), + // The index evaluator declines NOT; its child's contributing + // field IDs cannot prove that this conjunct was evaluated. + Predicate::Not(_) | Predicate::AlwaysFalse => false, + Predicate::AlwaysTrue => true, + } + } + /// Matching positions without data-evolution's unindexed row-id extension. /// Source-backed PK indexes cover a validated immutable source group. pub(super) async fn matching_ranges( diff --git a/crates/paimon/src/table/global_index_scanner/composite.rs b/crates/paimon/src/table/global_index_scanner/composite.rs index f61d0e4f2..b84f6e31b 100644 --- a/crates/paimon/src/table/global_index_scanner/composite.rs +++ b/crates/paimon/src/table/global_index_scanner/composite.rs @@ -195,6 +195,9 @@ impl GlobalIndexScanner { .map(|field| field.id()) .collect(), indexed_coverage, + // Tuple bounds can leave suffix constraints for data reads. + // Search pre-filters exclude these definitions entirely. + fully_evaluated: false, }, })); } diff --git a/crates/paimon/src/table/global_index_scanner/evaluator.rs b/crates/paimon/src/table/global_index_scanner/evaluator.rs index 3c50b5a42..cc3a6a37e 100644 --- a/crates/paimon/src/table/global_index_scanner/evaluator.rs +++ b/crates/paimon/src/table/global_index_scanner/evaluator.rs @@ -59,6 +59,9 @@ pub(super) struct GlobalIndexScanResult { pub(super) row_ranges: Vec, pub(super) evaluated_field_ids: HashSet, pub(super) indexed_coverage: Vec, + /// False when an AND conjunct was declined or only some constraints were + /// selected for a field. Contributing field IDs alone cannot prove exactness. + pub(super) fully_evaluated: bool, } pub(super) async fn try_fold_bounded( @@ -149,10 +152,13 @@ impl GlobalIndexScanner { self.evaluate_leaf(&selected_entries, &predicates) .await .map(|result| { - result.map(|(row_ranges, indexed_coverage)| GlobalIndexScanResult { - row_ranges, - evaluated_field_ids: HashSet::from([field_id]), - indexed_coverage, + result.map(|(row_ranges, indexed_coverage, fully_evaluated)| { + GlobalIndexScanResult { + row_ranges, + evaluated_field_ids: HashSet::from([field_id]), + indexed_coverage, + fully_evaluated, + } }) }) } @@ -161,6 +167,7 @@ impl GlobalIndexScanner { let mut row_ranges: Option> = None; let mut indexed_coverage: Option> = None; let mut evaluated_field_ids = HashSet::new(); + let mut fully_evaluated = true; loop { let remaining_predicate = Predicate::And(children.clone()); let Some(composite) = self.evaluate_composite(&remaining_predicate).await? @@ -168,6 +175,7 @@ impl GlobalIndexScanner { break; }; let result = composite.result; + fully_evaluated &= result.fully_evaluated; row_ranges = Some(match row_ranges { None => result.row_ranges, Some(existing) => { @@ -250,40 +258,48 @@ impl GlobalIndexScanner { } } } + let fully_evaluated = selected_predicates.len() == predicates.len(); leaf_futures.push(async move { let result = self .evaluate_leaf(&selected_entries, &selected_predicates) .await?; - Ok((field_id, result)) + Ok((field_id, result, fully_evaluated)) }); } } let leaf_group_count = leaf_futures.len(); - let (mut row_ranges, mut indexed_coverage, mut evaluated_field_ids) = - try_fold_bounded( - leaf_futures, - leaf_group_count.max(1), - (row_ranges, indexed_coverage, evaluated_field_ids), - |(row_ranges, indexed_coverage, evaluated_field_ids), - (field_id, result)| { - if let Some((ranges, coverage)) = result { - *row_ranges = Some(match row_ranges.take() { - None => ranges, - Some(existing) => { - intersect_sorted_ranges(&existing, &ranges) - } - }); - *indexed_coverage = Some(match indexed_coverage.take() { - None => coverage, - Some(existing) => { - intersect_sorted_ranges(&existing, &coverage) - } - }); - evaluated_field_ids.insert(field_id); - } - }, - ) - .await?; + let ( + mut row_ranges, + mut indexed_coverage, + mut evaluated_field_ids, + mut fully_evaluated, + ) = try_fold_bounded( + leaf_futures, + leaf_group_count.max(1), + ( + row_ranges, + indexed_coverage, + evaluated_field_ids, + fully_evaluated, + ), + |(row_ranges, indexed_coverage, evaluated_field_ids, fully_evaluated), + (field_id, result, group_fully_evaluated)| { + *fully_evaluated &= group_fully_evaluated && result.is_some(); + if let Some((ranges, coverage, leaf_fully_evaluated)) = result { + *fully_evaluated &= leaf_fully_evaluated; + *row_ranges = Some(match row_ranges.take() { + None => ranges, + Some(existing) => intersect_sorted_ranges(&existing, &ranges), + }); + *indexed_coverage = Some(match indexed_coverage.take() { + None => coverage, + Some(existing) => intersect_sorted_ranges(&existing, &coverage), + }); + evaluated_field_ids.insert(field_id); + } + }, + ) + .await?; // Evaluate non-leaf children concurrently in predicate order. let stream = futures::stream::iter(0..non_leaf_children.len()) @@ -292,6 +308,7 @@ impl GlobalIndexScanner { futures::pin_mut!(stream); while let Some(result) = stream.try_next().await? { if let Some(child_result) = result { + fully_evaluated &= child_result.fully_evaluated; row_ranges = Some(match row_ranges { None => child_result.row_ranges, Some(existing) => { @@ -306,6 +323,8 @@ impl GlobalIndexScanner { &child_result.indexed_coverage, ), }); + } else { + fully_evaluated = false; } } @@ -313,10 +332,12 @@ impl GlobalIndexScanner { row_ranges, evaluated_field_ids, indexed_coverage: indexed_coverage.unwrap_or_default(), + fully_evaluated, })) } Predicate::Or(children) => { let mut all_ranges: Vec = Vec::new(); + let mut fully_evaluated = true; let mut evaluated_field_ids = HashSet::new(); let mut indexed_coverage: Option> = None; let stream = futures::stream::iter(0..children.len()) @@ -326,6 +347,7 @@ impl GlobalIndexScanner { while let Some(result) = stream.try_next().await? { match result { Some(child_result) => { + fully_evaluated &= child_result.fully_evaluated; all_ranges.extend(child_result.row_ranges); evaluated_field_ids.extend(child_result.evaluated_field_ids); indexed_coverage = Some(match indexed_coverage { @@ -348,6 +370,7 @@ impl GlobalIndexScanner { row_ranges, evaluated_field_ids, indexed_coverage: indexed_coverage.unwrap_or_default(), + fully_evaluated, })) } _ => Ok(None), @@ -362,7 +385,7 @@ impl GlobalIndexScanner { &self, entries: &[&GlobalIndexEntry], predicates: &[(PredicateOperator, &[Datum], &DataType)], - ) -> Result, Vec)>> { + ) -> Result, Vec, bool)>> { let Some(normalized_predicates) = predicates .iter() .map(|(op, literals, data_type)| { @@ -415,7 +438,7 @@ impl GlobalIndexScanner { ) }; if between.as_ref().is_some_and(|range| range.is_empty()) { - return Ok(Some((Vec::new(), coverage()))); + return Ok(Some((Vec::new(), coverage(), true))); } let effective_predicates = if between.is_some() { @@ -569,6 +592,7 @@ impl GlobalIndexScanner { let mut query_plans = Vec::with_capacity(entries.len()); let mut all_matching_ranges = Vec::new(); + let mut fully_evaluated = true; for (entry_idx, entry) in entries.iter().enumerate() { if (between.is_none() || between_all_matches[entry_idx]) && predicate_all_matches @@ -632,6 +656,7 @@ impl GlobalIndexScanner { if file_cannot_match { continue; } + fully_evaluated &= !file_has_unsupported_match; if !file_evaluated { if file_has_unsupported_match { return Ok(None); @@ -722,6 +747,7 @@ impl GlobalIndexScanner { Ok(Some(( crate::table::merge_row_ranges(all_ranges), coverage(), + fully_evaluated, ))) } diff --git a/crates/paimon/src/table/global_index_scanner/tests.rs b/crates/paimon/src/table/global_index_scanner/tests.rs index 0a79be3ac..c00c99673 100644 --- a/crates/paimon/src/table/global_index_scanner/tests.rs +++ b/crates/paimon/src/table/global_index_scanner/tests.rs @@ -3903,3 +3903,123 @@ async fn test_between_unmatched_file_drops_remaining_match() { whose key range is [0, 198] — got {ranges:?}" ); } + +#[tokio::test] +async fn scalar_exactness_keeps_string_candidates_and_dropped_conjuncts_inexact() { + let io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let mut fields = string_schema_fields(); + fields.push(DataField::new( + 2, + "other".to_string(), + fields[0].data_type().clone(), + )); + let mut entry = make_global_index_entry( + "unused.index", + 1, + 100, + 109, + &BTreeIndexMeta::new(Some(b"alpha".to_vec()), Some(b"zeta".to_vec()), false), + ); + entry.index_file.row_count = 10; + entry.index_file.file_size = 100; + let scanner = GlobalIndexScanner::create(&io, "memory:/absent", 1, 0, 0, &[entry], &fields) + .unwrap() + .unwrap(); + let builder = crate::spec::PredicateBuilder::new(&fields); + let exact = builder + .less_than("name", Datum::String("zz".to_string())) + .unwrap(); + let (ranges, is_exact) = scanner + .matching_ranges_with_exactness(&exact) + .await + .unwrap() + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(100, 109)]); + assert!(is_exact); + for candidate in [ + builder + .contains("name", Datum::String("needle".to_string())) + .unwrap(), + builder + .ends_with("name", Datum::String("needle".to_string())) + .unwrap(), + builder + .like("name", Datum::String("%needle%".to_string()), None) + .unwrap(), + builder + .equal("other", Datum::String("needle".to_string())) + .unwrap(), + Predicate::negate(exact.clone()), + ] { + let filter = Predicate::and(vec![exact.clone(), candidate]); + let (ranges, is_exact) = scanner + .matching_ranges_with_exactness(&filter) + .await + .unwrap() + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(100, 109)]); + assert!(!is_exact, "{filter:?}"); + } +} + +#[tokio::test] +async fn scalar_exactness_detects_unsupported_conjunct_on_same_fm_field() { + let io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let path = "memory:/fm_exactness"; + let output = VecFileWrite::new(); + let mut writer = + FMGlobalIndexWriter::new(Box::new(output.clone()), FMWriteOptions::default()).unwrap(); + writer.write(None, 0).await.unwrap(); + writer.write(Some(b"alpha"), 1).await.unwrap(); + let result = writer.finish().await.unwrap(); + io.new_output(&format!("{path}/index/name.fm")) + .unwrap() + .write(bytes::Bytes::from(output.to_vec())) + .await + .unwrap(); + let mut entry = make_global_index_entry_with_type( + "fm", + "name.fm", + 1, + 0, + 1, + &BTreeIndexMeta::new(None, None, true), + ); + entry.index_file.row_count = result.row_count as i64; + entry.index_file.file_size = output.to_vec().len() as i64; + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .index_meta = Some(result.index_meta); + let fields = string_schema_fields(); + let scanner = GlobalIndexScanner::create(&io, path, 1, i64::MAX, i64::MAX, &[entry], &fields) + .unwrap() + .unwrap(); + let builder = crate::spec::PredicateBuilder::new(&fields); + let is_null = builder.is_null("name").unwrap(); + let (ranges, exact) = scanner + .matching_ranges_with_exactness(&is_null) + .await + .unwrap() + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(0, 0)]); + assert!(exact); + let filter = Predicate::and(vec![ + is_null, + builder + .starts_with("name", Datum::String("a".to_string())) + .unwrap(), + ]); + let (ranges, exact) = scanner + .matching_ranges_with_exactness(&filter) + .await + .unwrap() + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(0, 0)]); + assert!( + !exact, + "FM must not mark a dropped startsWith conjunct exact" + ); +} diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 8e23077dd..8614795b2 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -70,15 +70,13 @@ impl<'a> VectorSearchBuilder<'a> { /// /// On the primary-key vector path this remains a residual allow-list over /// physical positions, mirroring Java `PrimaryKeyVectorRead`. On the - /// data-evolution/global-index path the predicate is evaluated through a - /// snapshot-pinned table read (which can use scalar global indexes such as - /// BTree), producing global row IDs that are localized for each vector-index - /// shard and passed to the vector backend as an include filter. - /// - /// The whole predicate is both pushed into the scan — where it prunes whole - /// data files by their column stats — and applied per row as a residual over - /// the surviving files, so results stay exact. Sub-file row-range narrowing is - /// not performed; a surviving file is re-read in full for the residual. + /// data-evolution/global-index path an exact scalar-index lookup supplies + /// the vector backend's row-ID include filter. Candidate-only answers are + /// refined against snapshot-pinned data when + /// `global-index.filter.refine-from-data=true`; otherwise those candidates + /// are excluded. If no scalar index can evaluate the predicate, indexed + /// ranges are scored from raw data even in FAST mode. Raw searches always + /// apply the row predicate before Top-K, matching Java. pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { add_filter(&mut self.filter, filter); self From 68da723a4014508bb6584992bb3d4d0819d7e86d Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 10 Oct 2026 00:11:46 +0800 Subject: [PATCH 3/3] fix: recover scalar index coverage gaps in vector search --- crates/paimon/src/table/de_vector_read.rs | 72 ++++++++++++++- .../paimon/src/table/de_vector_read/tests.rs | 72 +++++++++++++++ .../paimon/src/table/de_vector_scan/tests.rs | 92 +++++++++++++++++++ 3 files changed, 232 insertions(+), 4 deletions(-) diff --git a/crates/paimon/src/table/de_vector_read.rs b/crates/paimon/src/table/de_vector_read.rs index 6e65d1331..665953bc2 100644 --- a/crates/paimon/src/table/de_vector_read.rs +++ b/crates/paimon/src/table/de_vector_read.rs @@ -29,6 +29,7 @@ use crate::table::global_index_scanner::{ deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries, GlobalIndexScanner, RowRangeIndex, }; +use crate::table::global_index_types::normalize_queryable_global_index_type; use crate::table::index_file_path::IndexFileLocation; use crate::table::pk_vector_position_read::SEARCH_SCORE_COLUMN; use crate::table::row_id_predicate::intersect_sorted_ranges; @@ -471,7 +472,6 @@ async fn evaluate_batch_vector_search( let table_path = evaluation.table_path.trim_end_matches('/'); let core_options = CoreOptions::new(evaluation.table_options); - let search_mode = core_options.vector_index_search_mode()?; let field_name = &vector_searches[0].field_name; if vector_searches .iter() @@ -495,6 +495,12 @@ async fn evaluate_batch_vector_search( }); } + // Java's scan overlays query options for the vector/scalar search modes. + let mut effective_options = evaluation.table_options.clone(); + effective_options.extend(search_options.clone()); + let effective_core = CoreOptions::new(&effective_options); + let search_mode = effective_core.vector_index_search_mode()?; + let field_id = match find_field_id_by_name(evaluation.schema_fields, field_name) { Some(id) => id, None => return Ok(vec![ScoredRowIds::empty(); vector_searches.len()]), @@ -544,6 +550,10 @@ async fn evaluate_batch_vector_search( } None => Some(filter.clone()), }); + let scalar_search_mode = scalar_filter + .as_ref() + .map(|_| effective_core.scalar_index_search_mode()) + .transpose()?; if let Some(filter) = scalar_filter .as_ref() .filter(|_| !vector_entries.is_empty()) @@ -914,9 +924,14 @@ async fn evaluate_batch_vector_search( refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); } - if route_indexed_to_raw || search_mode != GlobalIndexSearchMode::Fast { + if route_indexed_to_raw + || search_mode != GlobalIndexSearchMode::Fast + || scalar_search_mode.is_some_and(|mode| mode != GlobalIndexSearchMode::Fast) + { let raw_fallback_start = timing_enabled.then(Instant::now); - let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail { + let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail + || scalar_search_mode == Some(GlobalIndexSearchMode::Detail) + { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { message: "Vector raw search in detail mode requires table context".to_string(), source: None, @@ -934,10 +949,25 @@ async fn evaluate_batch_vector_search( &detail_ranges, is_vector_global_index_file, ); + if let (Some(filter), Some(scalar_mode)) = (&scalar_filter, scalar_search_mode) { + let mut scalar_ranges = scalar_unindexed_ranges( + evaluation, + index_entries, + filter, + scalar_mode, + &detail_ranges, + ); + // Scalar FULL/DETAIL also recover gaps inside vector coverage. In + // vector FAST they must not expand beyond the indexed corpus. + if search_mode == GlobalIndexSearchMode::Fast { + scalar_ranges = intersect_sorted_ranges(&scalar_ranges, &vector_ranges); + } + raw_ranges.extend(scalar_ranges); + } if route_indexed_to_raw { raw_ranges.extend(vector_ranges); - raw_ranges = merge_row_ranges(raw_ranges); } + raw_ranges = merge_row_ranges(raw_ranges); if !raw_ranges.is_empty() { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { message: "Vector raw search requires table context".to_string(), @@ -1034,6 +1064,40 @@ async fn evaluate_batch_vector_search( Ok(results) } +/// Predicate coverage requires every referenced field to be covered by a +/// single-field scalar index. Reuse the table scanner's coverage accounting. +fn scalar_unindexed_ranges( + evaluation: VectorSearchEvaluation<'_>, + entries: &[IndexManifestEntry], + filter: &Predicate, + search_mode: GlobalIndexSearchMode, + data_ranges: &[RowRange], +) -> Vec { + if !evaluation.next_row_id.is_some_and(|next| next > 0) { + return Vec::new(); + } + let mut leaf_refs = Vec::new(); + crate::arrow::residual::collect_predicate_leaf_refs(filter, &mut leaf_refs); + let field_ids = leaf_refs + .into_iter() + .filter_map(|(column, _)| find_field_id_by_name(evaluation.schema_fields, column)) + .collect(); + unindexed_ranges_for_global_index_entries( + entries, + &field_ids, + search_mode, + evaluation.next_row_id, + data_ranges, + |file| { + normalize_queryable_global_index_type(&file.index_type).is_some() + && file + .global_index_meta + .as_ref() + .is_some_and(|meta| meta.extra_field_ids.as_ref().is_none_or(Vec::is_empty)) + }, + ) +} + fn is_vector_global_index_file(index_file: &IndexFileMeta) -> bool { VectorIndexBackend::from_index_type(&index_file.index_type).is_some() } diff --git a/crates/paimon/src/table/de_vector_read/tests.rs b/crates/paimon/src/table/de_vector_read/tests.rs index d1fd1ba19..ed64c2cba 100644 --- a/crates/paimon/src/table/de_vector_read/tests.rs +++ b/crates/paimon/src/table/de_vector_read/tests.rs @@ -1567,3 +1567,75 @@ async fn filtered_rerank_rejects_metric_from_a_skipped_shard() { "{error}" ); } + +#[test] +fn scalar_unindexed_coverage_intersects_all_predicate_fields() { + let io = FileIOBuilder::new("memory").build().unwrap(); + let options = HashMap::new(); + let fields = vec![make_field(1, "id"), make_field(2, "other")]; + let evaluation = eval_context(&io, &options, &fields, Some(20)); + let mut id = make_lumina_entry("id.btree", "btree", FileKind::Add, 1); + let mut other = make_lumina_entry("other.btree", "btree", FileKind::Add, 2); + other + .index_file + .global_index_meta + .as_mut() + .unwrap() + .row_range_end = 4; + let mut composite = id.clone(); + let meta = composite.index_file.global_index_meta.as_mut().unwrap(); + meta.row_range_end = 19; + meta.extra_field_ids = Some(vec![2]); + let builder = crate::spec::PredicateBuilder::new(&fields); + let predicates = vec![ + builder.equal("id", crate::spec::Datum::Int(7)).unwrap(), + builder.equal("other", crate::spec::Datum::Int(7)).unwrap(), + ]; + for filter in [ + Predicate::and(predicates.clone()), + Predicate::or(predicates), + ] { + let entries = vec![id.clone(), other.clone(), composite.clone()]; + assert_eq!( + scalar_unindexed_ranges( + evaluation, + &entries, + &filter, + GlobalIndexSearchMode::Full, + &[] + ), + vec![RowRange::new(5, 19)] + ); + assert_eq!( + scalar_unindexed_ranges( + evaluation, + &entries, + &filter, + GlobalIndexSearchMode::Detail, + &[RowRange::new(0, 9), RowRange::new(15, 17)] + ), + vec![RowRange::new(5, 9), RowRange::new(15, 17)] + ); + assert!(scalar_unindexed_ranges( + evaluation, + &entries, + &filter, + GlobalIndexSearchMode::Fast, + &[] + ) + .is_empty()); + // An index from another family cannot fill scalar coverage. + id.index_file.index_type = "full-text".to_string(); + assert_eq!( + scalar_unindexed_ranges( + evaluation, + &[id.clone(), other.clone()], + &filter, + GlobalIndexSearchMode::Full, + &[] + ), + vec![RowRange::new(0, 19)] + ); + id.index_file.index_type = "btree".to_string(); + } +} diff --git a/crates/paimon/src/table/de_vector_scan/tests.rs b/crates/paimon/src/table/de_vector_scan/tests.rs index 60427a881..aa1b14f94 100644 --- a/crates/paimon/src/table/de_vector_scan/tests.rs +++ b/crates/paimon/src/table/de_vector_scan/tests.rs @@ -478,3 +478,95 @@ async fn scalar_filter_ignores_composite_index_definitions() { .unwrap(); assert_eq!(results[0].row_ids().unwrap().len(), 3); } + +#[tokio::test] +async fn scalar_coverage_gaps_follow_scalar_and_vector_modes() { + for vector_covers_tail in [false, true] { + let table = de_vector_table().await; + let element = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vectors = ListBuilder::new(Float32Builder::new()).with_field(element.clone()); + vectors.values().append_slice(&[3.0, 3.0]); + vectors.append(true); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element), true), + ])), + vec![ + Arc::new(Int32Array::from(vec![4])) as ArrayRef, + Arc::new(vectors.finish()) as ArrayRef, + ], + ) + .unwrap(); + let mut writer = TableWrite::new(&table, "coverage-test".to_string()).unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + TableCommit::new(table.clone(), "coverage-test".to_string()) + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + if vector_covers_tail { + table + .new_vindex_index_build_builder(crate::vindex::IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .with_options(HashMap::from([( + "ivf-flat.nlist".to_string(), + "1".to_string(), + )])) + .execute() + .await + .unwrap(); + } + let filter = id_gt_filter(&table, 3); + for scalar_mode in ["fast", "full", "detail"] { + for vector_mode in ["fast", "full", "detail"] { + for builder_options in [false, true] { + let modes = HashMap::from([ + ( + "vector-index.search-mode".to_string(), + vector_mode.to_string(), + ), + ( + "scalar-index.search-mode".to_string(), + scalar_mode.to_string(), + ), + ]); + let mut table_options = if builder_options { + HashMap::from([ + ("vector-index.search-mode".to_string(), "fast".to_string()), + ("scalar-index.search-mode".to_string(), "fast".to_string()), + ]) + } else { + modes.clone() + }; + table_options.insert( + "global-index.filter.refine-from-data".to_string(), + "false".to_string(), + ); + let table = table.copy_with_options(table_options); + let mut builder = table.new_batch_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vectors(vec![vec![0.0, 0.0], vec![1.0, 0.0]]) + .with_limit(10) + .with_filter(filter.clone()); + if builder_options { + builder.with_options(modes); + } + let expected = if (vector_covers_tail && scalar_mode != "fast") + || (!vector_covers_tail && vector_mode != "fast") + { + vec![3] + } else { + vec![] + }; + let results = builder.execute().await.unwrap(); + assert_eq!(results.len(), 2); + for result in results { + assert_eq!(result.row_ids().unwrap().row_ids, expected, + "scalar={scalar_mode}, vector={vector_mode}, vector_covers_tail={vector_covers_tail}, builder_options={builder_options}"); + } + } + } + } + } +}