You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by lz...@apache.org on 2026/04/23 12:00:53 UTC
(paimon) branch master updated: [index] vector search support predict in python (#7691)
This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new fe359ed414 [index] vector search support predict in python (#7691)
fe359ed414 is described below
commit fe359ed41416bac826d45853b7ca7a768489e962
Author: jerry <ji...@gmail.com>
AuthorDate: Thu Apr 23 20:00:47 2026 +0800
[index] vector search support predict in python (#7691)
---
.../paimon/lumina/index/JavaPyLuminaE2ETest.java | 136 ++++++
paimon-python/dev/run_mixed_tests.sh | 40 +-
.../globalindex/btree/btree_index_reader.py | 5 +-
.../pypaimon/globalindex/global_index_evaluator.py | 4 +-
.../pypaimon/globalindex/global_index_meta.py | 10 +-
.../pypaimon/globalindex/global_index_scanner.py | 70 ++-
.../lumina/lumina_vector_global_index_reader.py | 4 +-
.../tantivy_full_text_global_index_reader.py | 6 +-
.../globalindex/union_global_index_reader.py | 120 +++++
.../pypaimon/table/source/vector_search_builder.py | 109 ++++-
.../pypaimon/table/source/vector_search_read.py | 49 +-
.../pypaimon/table/source/vector_search_scan.py | 60 ++-
.../pypaimon/table/source/vector_search_split.py | 3 +-
.../pypaimon/tests/e2e/java_py_read_write_test.py | 45 ++
.../pypaimon/tests/vector_search_filter_test.py | 525 +++++++++++++++++++++
15 files changed, 1147 insertions(+), 39 deletions(-)
diff --git a/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/JavaPyLuminaE2ETest.java b/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/JavaPyLuminaE2ETest.java
index dd0212bda9..8826288965 100644
--- a/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/JavaPyLuminaE2ETest.java
+++ b/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/JavaPyLuminaE2ETest.java
@@ -25,8 +25,10 @@ import org.apache.paimon.fs.FileIOFinder;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.local.LocalFileIO;
import org.apache.paimon.globalindex.GlobalIndexBuilderUtils;
+import org.apache.paimon.globalindex.GlobalIndexParallelWriter;
import org.apache.paimon.globalindex.GlobalIndexSingletonWriter;
import org.apache.paimon.globalindex.ResultEntry;
+import org.apache.paimon.globalindex.btree.BTreeGlobalIndexerFactory;
import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.io.CompactIncrement;
import org.apache.paimon.io.DataIncrement;
@@ -212,4 +214,138 @@ public class JavaPyLuminaE2ETest {
assertThat(indexEntries.get(0).indexFile().indexType())
.isEqualTo(LuminaVectorGlobalIndexerFactory.IDENTIFIER);
}
+
+ @Test
+ @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true")
+ public void testLuminaVectorWithBTreeIndexWrite() throws Exception {
+ String tableName = "test_lumina_vector_btree_filter";
+ Path tablePath = new Path(warehouse.toString() + "/default.db/" + tableName);
+
+ int dimension = 4;
+
+ RowType rowType =
+ RowType.of(
+ new DataType[] {DataTypes.INT(), new ArrayType(new FloatType())},
+ new String[] {"id", "embedding"});
+
+ Options options = new Options();
+ options.set(PATH, tablePath.toString());
+ options.set(ROW_TRACKING_ENABLED, true);
+ options.set(DATA_EVOLUTION_ENABLED, true);
+ options.set(GLOBAL_INDEX_ENABLED, true);
+ options.setString(LuminaVectorIndexOptions.DIMENSION.key(), String.valueOf(dimension));
+ options.setString(LuminaVectorIndexOptions.DISTANCE_METRIC.key(), "l2");
+ options.setString(LuminaVectorIndexOptions.ENCODING_TYPE.key(), "rawf32");
+
+ TableSchema tableSchema =
+ SchemaUtils.forceCommit(
+ new SchemaManager(LocalFileIO.create(), tablePath),
+ new Schema(
+ rowType.getFields(),
+ Collections.emptyList(),
+ Collections.emptyList(),
+ options.toMap(),
+ ""));
+
+ AppendOnlyFileStoreTable table =
+ new AppendOnlyFileStoreTable(
+ FileIOFinder.find(tablePath),
+ tablePath,
+ tableSchema,
+ CatalogEnvironment.empty());
+
+ // Same 6 vectors as testLuminaVectorIndexWrite, paired with ids 0..5.
+ float[][] vectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f, 0.0f, 0.0f},
+ new float[] {0.9f, 0.1f, 0.0f, 0.0f},
+ new float[] {0.0f, 1.0f, 0.0f, 0.0f},
+ new float[] {0.0f, 0.0f, 1.0f, 0.0f},
+ new float[] {0.0f, 0.0f, 0.0f, 1.0f},
+ new float[] {0.95f, 0.05f, 0.0f, 0.0f}
+ };
+
+ BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
+ try (BatchTableWrite write = writeBuilder.newWrite();
+ BatchTableCommit commit = writeBuilder.newCommit()) {
+ for (int i = 0; i < vectors.length; i++) {
+ write.write(GenericRow.of(i, new GenericArray(vectors[i])));
+ }
+ commit.commit(write.prepareCommit());
+ }
+
+ Options indexOptions = table.coreOptions().toConfiguration();
+ Range rowRange = new Range(0, vectors.length - 1);
+
+ // Build Lumina vector index on "embedding".
+ DataField embeddingField = table.rowType().getField("embedding");
+ GlobalIndexSingletonWriter vectorWriter =
+ (GlobalIndexSingletonWriter)
+ GlobalIndexBuilderUtils.createIndexWriter(
+ table,
+ LuminaVectorGlobalIndexerFactory.IDENTIFIER,
+ embeddingField,
+ indexOptions);
+ for (float[] vec : vectors) {
+ vectorWriter.write(vec);
+ }
+ List<ResultEntry> vectorEntries = vectorWriter.finish();
+ List<IndexFileMeta> vectorIndexFiles =
+ GlobalIndexBuilderUtils.toIndexFileMetas(
+ table.fileIO(),
+ table.store().pathFactory().globalIndexFileFactory(),
+ table.coreOptions(),
+ rowRange,
+ embeddingField.id(),
+ LuminaVectorGlobalIndexerFactory.IDENTIFIER,
+ vectorEntries);
+
+ // Build BTree global index on "id".
+ DataField idField = table.rowType().getField("id");
+ GlobalIndexParallelWriter idWriter =
+ (GlobalIndexParallelWriter)
+ GlobalIndexBuilderUtils.createIndexWriter(
+ table, BTreeGlobalIndexerFactory.IDENTIFIER, idField, indexOptions);
+ for (int i = 0; i < vectors.length; i++) {
+ idWriter.write(i, i);
+ }
+ List<ResultEntry> idEntries = idWriter.finish();
+ List<IndexFileMeta> idIndexFiles =
+ GlobalIndexBuilderUtils.toIndexFileMetas(
+ table.fileIO(),
+ table.store().pathFactory().globalIndexFileFactory(),
+ table.coreOptions(),
+ rowRange,
+ idField.id(),
+ BTreeGlobalIndexerFactory.IDENTIFIER,
+ idEntries);
+
+ // Commit both index sets in a single index-only commit.
+ java.util.List<IndexFileMeta> allIndexFiles = new java.util.ArrayList<>();
+ allIndexFiles.addAll(vectorIndexFiles);
+ allIndexFiles.addAll(idIndexFiles);
+ DataIncrement dataIncrement = DataIncrement.indexIncrement(allIndexFiles);
+ CommitMessage message =
+ new CommitMessageImpl(
+ BinaryRow.EMPTY_ROW,
+ 0,
+ null,
+ dataIncrement,
+ CompactIncrement.emptyIncrement());
+ try (BatchTableCommit commit = writeBuilder.newCommit()) {
+ commit.commit(Collections.singletonList(message));
+ }
+
+ // Verify both index types are present.
+ List<org.apache.paimon.manifest.IndexManifestEntry> indexEntries =
+ table.indexManifestFileReader().read(table.latestSnapshot().get().indexManifest());
+ assertThat(indexEntries).hasSize(2);
+ assertThat(
+ indexEntries.stream()
+ .map(e -> e.indexFile().indexType())
+ .collect(java.util.stream.Collectors.toSet()))
+ .containsExactlyInAnyOrder(
+ LuminaVectorGlobalIndexerFactory.IDENTIFIER,
+ BTreeGlobalIndexerFactory.IDENTIFIER);
+ }
}
diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh
index 077b5af276..f1aae334a5 100755
--- a/paimon-python/dev/run_mixed_tests.sh
+++ b/paimon-python/dev/run_mixed_tests.sh
@@ -290,6 +290,30 @@ run_lumina_vector_test() {
fi
}
+# Function to run Lumina vector + BTree pre-filter test.
+run_lumina_vector_btree_test() {
+ echo -e "${YELLOW}=== Running Lumina Vector + BTree Pre-Filter Test (Java Write, Python Read) ===${NC}"
+
+ cd "$PROJECT_ROOT"
+
+ echo "Running Maven test for JavaPyLuminaE2ETest.testLuminaVectorWithBTreeIndexWrite..."
+ if mvn test -Dtest=org.apache.paimon.lumina.index.JavaPyLuminaE2ETest#testLuminaVectorWithBTreeIndexWrite -pl paimon-lumina -q -Drun.e2e.tests=true; then
+ echo -e "${GREEN}✓ Java test completed successfully${NC}"
+ else
+ echo -e "${RED}✗ Java test failed${NC}"
+ return 1
+ fi
+ cd "$PAIMON_PYTHON_DIR"
+ echo "Running Python test for JavaPyReadWriteTest.test_read_lumina_vector_with_btree_filter..."
+ if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest::test_read_lumina_vector_with_btree_filter -v; then
+ echo -e "${GREEN}✓ Python test completed successfully${NC}"
+ return 0
+ else
+ echo -e "${RED}✗ Python test failed${NC}"
+ return 1
+ fi
+}
+
run_compact_conflict_test() {
echo -e "${YELLOW}=== Running Compact Conflict Test (Java Write Base, Python Shard Update + Java Compact) ===${NC}"
@@ -350,6 +374,7 @@ main() {
local compressed_text_result=0
local tantivy_fulltext_result=0
local lumina_vector_result=0
+ local lumina_vector_btree_result=0
local compact_conflict_result=0
local blob_alter_compact_result=0
@@ -434,6 +459,13 @@ main() {
echo ""
+ # Run Lumina vector + BTree pre-filter test (Java write, Python read)
+ if ! run_lumina_vector_btree_test; then
+ lumina_vector_btree_result=1
+ fi
+
+ echo ""
+
# Run compact conflict test (Java write+compact, Python read)
if ! run_compact_conflict_test; then
compact_conflict_result=1
@@ -504,6 +536,12 @@ main() {
echo -e "${RED}✗ Lumina Vector Index Test (Java Write, Python Read): FAILED${NC}"
fi
+ if [[ $lumina_vector_btree_result -eq 0 ]]; then
+ echo -e "${GREEN}✓ Lumina Vector + BTree Pre-Filter Test (Java Write, Python Read): PASSED${NC}"
+ else
+ echo -e "${RED}✗ Lumina Vector + BTree Pre-Filter Test (Java Write, Python Read): FAILED${NC}"
+ fi
+
if [[ $compact_conflict_result -eq 0 ]]; then
echo -e "${GREEN}✓ Compact Conflict Test (Java Write+Compact, Python Read): PASSED${NC}"
else
@@ -521,7 +559,7 @@ main() {
# Clean up warehouse directory after all tests
cleanup_warehouse
- if [[ $java_write_result -eq 0 && $python_read_result -eq 0 && $python_write_result -eq 0 && $java_read_result -eq 0 && $pk_dv_result -eq 0 && $btree_index_result -eq 0 && $compressed_text_result -eq 0 && $tantivy_fulltext_result -eq 0 && $lumina_vector_result -eq 0 && $compact_conflict_result -eq 0 && $blob_alter_compact_result -eq 0 ]]; then
+ if [[ $java_write_result -eq 0 && $python_read_result -eq 0 && $python_write_result -eq 0 && $java_read_result -eq 0 && $pk_dv_result -eq 0 && $btree_index_result -eq 0 && $compressed_text_result -eq 0 && $tantivy_fulltext_result -eq 0 && $lumina_vector_result -eq 0 && $lumina_vector_btree_result -eq 0 && $compact_conflict_result -eq 0 && $blob_alter_compact_result -eq 0 ]]; then
echo -e "${GREEN}🎉 All tests passed! Java-Python interoperability verified.${NC}"
return 0
else
diff --git a/paimon-python/pypaimon/globalindex/btree/btree_index_reader.py b/paimon-python/pypaimon/globalindex/btree/btree_index_reader.py
index 65540f8c57..15aa637cc7 100644
--- a/paimon-python/pypaimon/globalindex/btree/btree_index_reader.py
+++ b/paimon-python/pypaimon/globalindex/btree/btree_index_reader.py
@@ -86,7 +86,10 @@ class BTreeIndexReader(GlobalIndexReader):
self.max_key = None
self.has_nulls = index_meta.has_nulls
- self.input_stream = file_io.new_input_stream(index_path + "/" + io_meta.file_name)
+ file_path = (io_meta.external_path
+ if io_meta.external_path
+ else index_path + "/" + io_meta.file_name)
+ self.input_stream = file_io.new_input_stream(file_path)
# Lazy-loaded null bitmap
self._null_bitmap: Optional[RoaringBitmap64] = None
diff --git a/paimon-python/pypaimon/globalindex/global_index_evaluator.py b/paimon-python/pypaimon/globalindex/global_index_evaluator.py
index bbab92e883..e7c0799c74 100644
--- a/paimon-python/pypaimon/globalindex/global_index_evaluator.py
+++ b/paimon-python/pypaimon/globalindex/global_index_evaluator.py
@@ -154,9 +154,11 @@ class GlobalIndexEvaluator:
return reader.visit_ends_with(field_ref, literals[0])
elif method == 'contains':
return reader.visit_contains(field_ref, literals[0])
+ elif method == 'like':
+ return reader.visit_like(field_ref, literals[0])
elif method == 'between':
return reader.visit_between(field_ref, literals[0], literals[1])
-
+
return None
def close(self) -> None:
diff --git a/paimon-python/pypaimon/globalindex/global_index_meta.py b/paimon-python/pypaimon/globalindex/global_index_meta.py
index 2c21f81a64..4e1c2f042f 100644
--- a/paimon-python/pypaimon/globalindex/global_index_meta.py
+++ b/paimon-python/pypaimon/globalindex/global_index_meta.py
@@ -60,6 +60,11 @@ class GlobalIndexIOMeta:
file_name: str
file_size: int
metadata: Optional[bytes] = None
+ # Optional full path, mirrors Java IndexFileMeta.externalPath(). When set,
+ # readers should open this path directly instead of joining the caller's
+ # index_path with file_name (Java FileStorePathFactory#toPath prefers
+ # externalPath over indexPath/fileName).
+ external_path: Optional[str] = None
def __eq__(self, other: object) -> bool:
if not isinstance(other, GlobalIndexIOMeta):
@@ -67,8 +72,9 @@ class GlobalIndexIOMeta:
return (
self.file_name == other.file_name and
self.file_size == other.file_size and
- self.metadata == other.metadata
+ self.metadata == other.metadata and
+ self.external_path == other.external_path
)
def __hash__(self) -> int:
- return hash((self.file_name, self.file_size))
+ return hash((self.file_name, self.file_size, self.external_path))
diff --git a/paimon-python/pypaimon/globalindex/global_index_scanner.py b/paimon-python/pypaimon/globalindex/global_index_scanner.py
index 19fba0cdeb..38408059a9 100644
--- a/paimon-python/pypaimon/globalindex/global_index_scanner.py
+++ b/paimon-python/pypaimon/globalindex/global_index_scanner.py
@@ -64,7 +64,8 @@ class GlobalIndexScanner:
io_meta = GlobalIndexIOMeta(
file_name=index_file.file_name,
file_size=index_file.file_size,
- metadata=global_index_meta.index_meta
+ metadata=global_index_meta.index_meta,
+ external_path=index_file.external_path,
)
index_metas[field_id][index_type][range_key].append(io_meta)
@@ -141,22 +142,63 @@ class GlobalIndexScanner:
def _create_readers(file_io, index_path, index_type_metas, field):
- """Create readers for a specific field from the index metadata."""
+ """Create readers for a specific field, dispatched by index_type.
+
+ Unknown indexTypes raise — a silent skip would make
+ ``VectorSearchReadImpl._pre_filter`` return ``None`` and the vector search
+ would then return unfiltered results that violate the user predicate.
+ """
if index_type_metas is None:
return []
+ from pypaimon.globalindex.offset_global_index_reader import (
+ OffsetGlobalIndexReader,
+ )
+ from pypaimon.globalindex.union_global_index_reader import (
+ UnionGlobalIndexReader,
+ )
+
readers = []
for index_type, range_metas in index_type_metas.items():
- if index_type == 'btree':
- from pypaimon.globalindex.btree import BTreeIndexReader
- from pypaimon.globalindex.btree.key_serializer import create_serializer
- for range_key, io_metas in range_metas.items():
- for metadata in io_metas:
- reader = BTreeIndexReader(
- key_serializer=create_serializer(field.type),
- file_io=file_io,
- index_path=index_path,
- io_meta=metadata
- )
- readers.append(reader)
+ offset_readers = []
+ for range_key, io_metas in range_metas.items():
+ inner_readers = _create_inner_readers(
+ index_type, file_io, index_path, field, io_metas)
+ for inner in inner_readers:
+ offset_readers.append(
+ OffsetGlobalIndexReader(
+ inner, range_key.from_, range_key.to))
+ if offset_readers:
+ readers.append(UnionGlobalIndexReader(offset_readers))
return readers
+
+
+def _create_inner_readers(index_type, file_io, index_path, field, io_metas):
+ """Build per-file (or per-shard) readers for a single indexType/range."""
+ if index_type == 'btree':
+ from pypaimon.globalindex.btree import BTreeIndexReader
+ from pypaimon.globalindex.btree.key_serializer import create_serializer
+ key_serializer = create_serializer(field.type)
+ return [
+ BTreeIndexReader(
+ key_serializer=key_serializer,
+ file_io=file_io,
+ index_path=index_path,
+ io_meta=io_meta,
+ )
+ for io_meta in io_metas
+ ]
+
+ from pypaimon.globalindex.tantivy import (
+ TANTIVY_FULLTEXT_IDENTIFIER,
+ TantivyFullTextGlobalIndexReader,
+ )
+ if index_type == TANTIVY_FULLTEXT_IDENTIFIER:
+ # Tantivy expects one file per shard; create one reader per io_meta.
+ return [
+ TantivyFullTextGlobalIndexReader(file_io, index_path, [io_meta])
+ for io_meta in io_metas
+ ]
+
+ raise ValueError(
+ "Unsupported global-index type in scanner: '%s'" % index_type)
diff --git a/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py b/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
index 2ec7acb9f5..b3c196d191 100644
--- a/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
+++ b/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
@@ -121,7 +121,9 @@ class LuminaVectorGlobalIndexReader(GlobalIndexReader):
searcher_options.update(self._index_meta.options)
self._search_options = searcher_options
- file_path = os.path.join(self._index_path, self._io_meta.file_name)
+ file_path = (self._io_meta.external_path
+ if self._io_meta.external_path
+ else os.path.join(self._index_path, self._io_meta.file_name))
stream = self._file_io.new_input_stream(file_path)
try:
self._searcher = LuminaSearcher(searcher_options)
diff --git a/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py b/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py
index c8b1711b59..46309b5348 100644
--- a/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py
+++ b/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py
@@ -168,8 +168,10 @@ class TantivyFullTextGlobalIndexReader(GlobalIndexReader):
import tantivy
- # Open the archive stream
- file_path = os.path.join(self._index_path, self._io_meta.file_name)
+ # Open the archive stream (prefer external_path if the manifest set it).
+ file_path = (self._io_meta.external_path
+ if self._io_meta.external_path
+ else os.path.join(self._index_path, self._io_meta.file_name))
stream = self._file_io.new_input_stream(file_path)
try:
# Parse archive header to get file layout
diff --git a/paimon-python/pypaimon/globalindex/union_global_index_reader.py b/paimon-python/pypaimon/globalindex/union_global_index_reader.py
new file mode 100644
index 0000000000..3dc0832e55
--- /dev/null
+++ b/paimon-python/pypaimon/globalindex/union_global_index_reader.py
@@ -0,0 +1,120 @@
+################################################################################
+# 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.
+################################################################################
+
+"""A GlobalIndexReader that unions results from multiple underlying readers.
+
+Each visit_* call is dispatched to every underlying reader and the results are
+OR-combined into a single ``GlobalIndexResult``. Readers returning ``None``
+("cannot answer") are skipped; an empty bitmap DOES contribute to the union
+and is NOT a short-circuit signal.
+"""
+
+from typing import Callable, List, Optional
+
+from pypaimon.globalindex.global_index_reader import FieldRef, GlobalIndexReader
+from pypaimon.globalindex.global_index_result import GlobalIndexResult
+
+
+class UnionGlobalIndexReader(GlobalIndexReader):
+
+ def __init__(self, readers: List[GlobalIndexReader]):
+ self._readers = readers
+
+ def _union(self, visitor: Callable[[GlobalIndexReader], Optional[GlobalIndexResult]]
+ ) -> Optional[GlobalIndexResult]:
+ result: Optional[GlobalIndexResult] = None
+ for reader in self._readers:
+ current = visitor(reader)
+ if current is None:
+ continue
+ if result is None:
+ result = current
+ else:
+ result = result.or_(current)
+ return result
+
+ # ---- vector / full-text search ----------------------------------------
+
+ def visit_vector_search(self, vector_search) -> Optional[GlobalIndexResult]:
+ from pypaimon.globalindex.vector_search_result import (
+ ScoredGlobalIndexResult,
+ )
+ result: Optional[ScoredGlobalIndexResult] = None
+ for reader in self._readers:
+ current = reader.visit_vector_search(vector_search)
+ if current is None:
+ continue
+ result = current if result is None else result.or_(current)
+ return result
+
+ def visit_full_text_search(self, full_text_search) -> Optional[GlobalIndexResult]:
+ return self._union(lambda r: r.visit_full_text_search(full_text_search))
+
+ # ---- scalar predicates (every reader sees the visit) ------------------
+
+ def visit_equal(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_equal(field_ref, literal))
+
+ def visit_not_equal(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_not_equal(field_ref, literal))
+
+ def visit_less_than(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_less_than(field_ref, literal))
+
+ def visit_less_or_equal(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_less_or_equal(field_ref, literal))
+
+ def visit_greater_than(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_greater_than(field_ref, literal))
+
+ def visit_greater_or_equal(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_greater_or_equal(field_ref, literal))
+
+ def visit_is_null(self, field_ref: FieldRef):
+ return self._union(lambda r: r.visit_is_null(field_ref))
+
+ def visit_is_not_null(self, field_ref: FieldRef):
+ return self._union(lambda r: r.visit_is_not_null(field_ref))
+
+ def visit_in(self, field_ref: FieldRef, literals):
+ return self._union(lambda r: r.visit_in(field_ref, literals))
+
+ def visit_not_in(self, field_ref: FieldRef, literals):
+ return self._union(lambda r: r.visit_not_in(field_ref, literals))
+
+ def visit_starts_with(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_starts_with(field_ref, literal))
+
+ def visit_ends_with(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_ends_with(field_ref, literal))
+
+ def visit_contains(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_contains(field_ref, literal))
+
+ def visit_like(self, field_ref: FieldRef, literal):
+ return self._union(lambda r: r.visit_like(field_ref, literal))
+
+ def visit_between(self, field_ref: FieldRef, from_v, to_v):
+ return self._union(lambda r: r.visit_between(field_ref, from_v, to_v))
+
+ def close(self) -> None:
+ for reader in self._readers:
+ try:
+ reader.close()
+ except Exception:
+ pass
diff --git a/paimon-python/pypaimon/table/source/vector_search_builder.py b/paimon-python/pypaimon/table/source/vector_search_builder.py
index 95f3cdfc67..05eb6708e3 100644
--- a/paimon-python/pypaimon/table/source/vector_search_builder.py
+++ b/paimon-python/pypaimon/table/source/vector_search_builder.py
@@ -20,6 +20,7 @@
from abc import ABC, abstractmethod
+from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.table.source.vector_search_read import VectorSearchReadImpl
from pypaimon.table.source.vector_search_scan import VectorSearchScanImpl
@@ -45,6 +46,18 @@ class VectorSearchBuilder(ABC):
"""The query vector (list of floats)."""
pass
+ @abstractmethod
+ def with_filter(self, predicate):
+ # type: (Predicate) -> VectorSearchBuilder
+ """Scalar predicate used to pre-filter rows before vector search."""
+ pass
+
+ @abstractmethod
+ def with_partition_filter(self, partition_filter):
+ # type: (Predicate) -> VectorSearchBuilder
+ """Partition predicate used to prune index manifest entries."""
+ pass
+
@abstractmethod
def new_vector_search_scan(self):
# type: () -> VectorSearchScan
@@ -73,6 +86,8 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
self._limit = 0
self._vector_column = None
self._query_vector = None
+ self._filter = None
+ self._partition_filter = None
def with_limit(self, limit):
# type: (int) -> VectorSearchBuilder
@@ -92,11 +107,97 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
self._query_vector = vector
return self
+ def with_filter(self, predicate):
+ # type: (Predicate) -> VectorSearchBuilder
+ if predicate is None:
+ return self
+ if self._filter is None:
+ self._filter = predicate
+ else:
+ self._filter = PredicateBuilder.and_predicates([self._filter, predicate])
+ # split out the partition-only conjuncts and store them as _partition_filter for
+ # manifest pruning. Non-partition conjuncts remain in self._filter;
+ # the silent drop of non-partition conjuncts *in the extracted copy*
+ # is intentional — nothing is lost overall.
+ extracted = self._extract_partition_only_conjuncts(predicate)
+ if extracted is not None:
+ if self._partition_filter is None:
+ self._partition_filter = extracted
+ else:
+ self._partition_filter = PredicateBuilder.and_predicates(
+ [self._partition_filter, extracted])
+ return self
+
+ def with_partition_filter(self, partition_filter):
+ # type: (Predicate) -> VectorSearchBuilder
+ if partition_filter is None:
+ self._partition_filter = None
+ return self
+ # Strict: every referenced field must be a partition key, otherwise a
+ # non-partition conjunct would be silently dropped (with_filter has
+ # the scalar fallback; with_partition_filter does not).
+ partition_keys = list(self._table.partition_keys or [])
+ if not partition_keys:
+ raise ValueError(
+ "with_partition_filter called on a non-partitioned table")
+ from pypaimon.read.push_down_utils import _get_all_fields
+ referenced = _get_all_fields(partition_filter)
+ extras = referenced - set(partition_keys)
+ if extras:
+ raise ValueError(
+ "Partition filter must reference only partition keys "
+ "(%s); got non-partition field(s): %s"
+ % (partition_keys, sorted(extras)))
+ self._partition_filter = self._rebuild_leaf_indices_by_name(
+ partition_filter,
+ {name: idx for idx, name in enumerate(partition_keys)},
+ )
+ return self
+
+ def _extract_partition_only_conjuncts(self, predicate):
+ """AND-split ``predicate``, keep conjuncts that reference ONLY
+ partition keys, and rebuild their leaf indices against the
+ partition-only row by field name (so the caller's PredicateBuilder
+ convention — full-row or partition-row — doesn't matter).
+ """
+ partition_keys = list(self._table.partition_keys or [])
+ if not partition_keys:
+ return None
+ from pypaimon.read.push_down_utils import _split_and, _get_all_fields
+ partition_key_set = set(partition_keys)
+ pk_to_idx = {name: idx for idx, name in enumerate(partition_keys)}
+ kept = [p for p in _split_and(predicate)
+ if _get_all_fields(p).issubset(partition_key_set)]
+ if not kept:
+ return None
+ rebuilt = [self._rebuild_leaf_indices_by_name(p, pk_to_idx)
+ for p in kept]
+ return PredicateBuilder.and_predicates(rebuilt)
+
+ @classmethod
+ def _rebuild_leaf_indices_by_name(cls, predicate, pk_to_idx):
+ """Return a copy of ``predicate`` with every leaf's ``index`` set to
+ its position in ``pk_to_idx`` (field-name lookup). Input predicate may
+ have been built against any schema — we key off ``Predicate.field``
+ rather than ``Predicate.index`` so positional convention doesn't
+ matter.
+ """
+ if predicate.method in ('and', 'or'):
+ new_children = [cls._rebuild_leaf_indices_by_name(c, pk_to_idx)
+ for c in (predicate.literals or [])]
+ return predicate.new_literals(new_children)
+ return predicate.new_index(pk_to_idx[predicate.field])
+
def new_vector_search_scan(self):
# type: () -> VectorSearchScan
if self._vector_column is None:
raise ValueError("Vector column must be set via with_vector_column()")
- return VectorSearchScanImpl(self._table, self._vector_column)
+ return VectorSearchScanImpl(
+ self._table,
+ self._vector_column,
+ filter_=self._filter,
+ partition_filter=self._partition_filter,
+ )
def new_vector_search_read(self):
# type: () -> VectorSearchRead
@@ -107,5 +208,9 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
if self._query_vector is None:
raise ValueError("Query vector must be set via with_query_vector()")
return VectorSearchReadImpl(
- self._table, self._limit, self._vector_column, self._query_vector
+ self._table,
+ self._limit,
+ self._vector_column,
+ self._query_vector,
+ filter_=self._filter,
)
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py
index 02776ae085..3cb17f2614 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -43,30 +43,65 @@ class VectorSearchRead(ABC):
class VectorSearchReadImpl(VectorSearchRead):
"""Implementation for VectorSearchRead."""
- def __init__(self, table, limit, vector_column, query_vector):
+ def __init__(self, table, limit, vector_column, query_vector, filter_=None):
self._table = table
self._limit = limit
self._vector_column = vector_column
self._query_vector = query_vector
+ self._filter = filter_
def read(self, splits):
# type: (List[VectorSearchSplit]) -> GlobalIndexResult
if not splits:
return GlobalIndexResult.create_empty()
+ pre_filter = self._pre_filter(splits)
+
result = ScoredGlobalIndexResult.create_empty()
for split in splits:
split_result = self._eval(
split.row_range_start, split.row_range_end,
- split.vector_index_files
+ split.vector_index_files, pre_filter
)
if split_result is not None:
result = result.or_(split_result)
return result.top_k(self._limit)
- def _eval(self, row_range_start, row_range_end, vector_index_files):
- # type: (int, int, list) -> Optional[ScoredGlobalIndexResult]
+ def _pre_filter(self, splits):
+ # type: (list) -> Optional[RoaringBitmap64]
+ """Evaluate the scalar filter against scalar global indexes to produce a row-id bitmap."""
+ if self._filter is None:
+ return None
+
+ # Collect scalar index files across splits, deduplicated by file name.
+ seen = set()
+ scalar_files = []
+ for split in splits:
+ for index_file in split.scalar_index_files:
+ if index_file.file_name in seen:
+ continue
+ seen.add(index_file.file_name)
+ scalar_files.append(index_file)
+
+ if not scalar_files:
+ return None
+
+ from pypaimon.globalindex.global_index_scanner import GlobalIndexScanner
+ scanner = GlobalIndexScanner.create(self._table, index_files=scalar_files)
+ if scanner is None:
+ return None
+ try:
+ result = scanner.scan(self._filter)
+ if result is None:
+ return None
+ return result.results()
+ finally:
+ scanner.close()
+
+ def _eval(self, row_range_start, row_range_end, vector_index_files,
+ include_row_ids):
+ # type: (int, int, list, Optional[RoaringBitmap64]) -> Optional[ScoredGlobalIndexResult]
if not vector_index_files:
return None
index_io_meta_list = []
@@ -77,7 +112,8 @@ class VectorSearchReadImpl(VectorSearchRead):
GlobalIndexIOMeta(
file_name=index_file.file_name,
file_size=index_file.file_size,
- metadata=meta.index_meta
+ metadata=meta.index_meta,
+ external_path=index_file.external_path,
)
)
@@ -91,6 +127,9 @@ class VectorSearchReadImpl(VectorSearchRead):
limit=self._limit,
field_name=self._vector_column.name
)
+ if include_row_ids is not None:
+ vector_search = vector_search.with_include_row_ids(include_row_ids)
+
with _create_vector_reader(
index_type, file_io, index_path,
index_io_meta_list, options
diff --git a/paimon-python/pypaimon/table/source/vector_search_scan.py b/paimon-python/pypaimon/table/source/vector_search_scan.py
index 450b334b13..5245d41f19 100644
--- a/paimon-python/pypaimon/table/source/vector_search_scan.py
+++ b/paimon-python/pypaimon/table/source/vector_search_scan.py
@@ -49,19 +49,33 @@ class VectorSearchScan(ABC):
class VectorSearchScanImpl(VectorSearchScan):
"""Implementation for VectorSearchScan."""
- def __init__(self, table, vector_column):
+ def __init__(self, table, vector_column, filter_=None, partition_filter=None):
self._table = table
self._vector_column = vector_column
+ self._filter = filter_
+ self._partition_filter = partition_filter
def scan(self):
# type: () -> VectorSearchScanPlan
+ from pypaimon.common.options.options import Options
from pypaimon.index.index_file_handler import IndexFileHandler
+ from pypaimon.read.push_down_utils import _get_all_fields
from pypaimon.snapshot.snapshot_manager import SnapshotManager
+ from pypaimon.snapshot.time_travel_util import TimeTravelUtil
vector_column = self._vector_column
- from pypaimon.snapshot.time_travel_util import TimeTravelUtil
- from pypaimon.common.options.options import Options
+ # Field IDs referenced by the scalar filter — used to pull in scalar
+ # index files for pre-filtering.
+ filter_field_ids = set()
+ if self._filter is not None:
+ referenced = _get_all_fields(self._filter)
+ field_by_name = {f.name: f for f in self._table.fields}
+ for name in referenced:
+ field = field_by_name.get(name)
+ if field is not None:
+ filter_field_ids.add(field.id)
+
snapshot = TimeTravelUtil.try_travel_to_snapshot(
Options(self._table.table_schema.options),
self._table.tag_manager()
@@ -71,25 +85,53 @@ class VectorSearchScanImpl(VectorSearchScan):
index_file_handler = IndexFileHandler(table=self._table)
+ partition_filter = self._partition_filter
+
def index_file_filter(entry):
+ if partition_filter is not None:
+ if not partition_filter.test(entry.partition):
+ return False
global_index_meta = entry.index_file.global_index_meta
if global_index_meta is None:
return False
- return vector_column.id == global_index_meta.index_field_id
+ field_id = global_index_meta.index_field_id
+ if vector_column.id == field_id:
+ return True
+ return field_id in filter_field_ids
entries = index_file_handler.scan(snapshot, index_file_filter)
all_index_files = [entry.index_file for entry in entries]
- # Group by (rowRangeStart, rowRangeEnd)
- by_range = defaultdict(list)
+ # Group vector index files by (rowRangeStart, rowRangeEnd).
+ vector_by_range = defaultdict(list)
for index_file in all_index_files:
meta = index_file.global_index_meta
assert meta is not None
+ if meta.index_field_id != vector_column.id:
+ continue
range_key = Range(meta.row_range_start, meta.row_range_end)
- by_range[range_key].append(index_file)
+ vector_by_range[range_key].append(index_file)
+ # For each vector range, attach matching scalar index files whose
+ # row range intersects the vector range.
splits = []
- for range_key, files in by_range.items():
- splits.append(VectorSearchSplit(range_key.from_, range_key.to, files))
+ for range_key, vector_files in vector_by_range.items():
+ scalar_files = []
+ for index_file in all_index_files:
+ meta = index_file.global_index_meta
+ assert meta is not None
+ if meta.index_field_id == vector_column.id:
+ continue
+ scalar_range = Range(meta.row_range_start, meta.row_range_end)
+ if range_key.overlaps(scalar_range):
+ scalar_files.append(index_file)
+ splits.append(
+ VectorSearchSplit(
+ range_key.from_,
+ range_key.to,
+ vector_files,
+ scalar_files,
+ )
+ )
return VectorSearchScanPlan(splits)
diff --git a/paimon-python/pypaimon/table/source/vector_search_split.py b/paimon-python/pypaimon/table/source/vector_search_split.py
index 80ea8d5e90..c80c8cb768 100644
--- a/paimon-python/pypaimon/table/source/vector_search_split.py
+++ b/paimon-python/pypaimon/table/source/vector_search_split.py
@@ -18,7 +18,7 @@
"""Split of vector search."""
-from dataclasses import dataclass
+from dataclasses import dataclass, field
from typing import List
from pypaimon.index.index_file_meta import IndexFileMeta
@@ -31,3 +31,4 @@ class VectorSearchSplit:
row_range_start: int
row_range_end: int
vector_index_files: List[IndexFileMeta]
+ scalar_index_files: List[IndexFileMeta] = field(default_factory=list)
diff --git a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py
index cf34ee27c0..6412c85320 100644
--- a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py
+++ b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py
@@ -627,6 +627,51 @@ class JavaPyReadWriteTest(unittest.TestCase):
print(f"Lumina vector search matched rows: ids={ids}")
self.assertIn(0, ids)
+ def test_read_lumina_vector_with_btree_filter(self):
+ """Vector search + btree scalar pre-filter, using a table that Java
+ populated with both a Lumina vector index on `embedding` and a BTree
+ global index on `id` (JavaPyLuminaE2ETest.testLuminaVectorWithBTreeIndexWrite)."""
+ from pypaimon.common.predicate_builder import PredicateBuilder
+
+ table = self.catalog.get_table('default.test_lumina_vector_btree_filter')
+
+ # Baseline search — same 6 vectors as test_read_lumina_vector_index,
+ # no filter, top 6 covers the whole table.
+ baseline = (table.new_vector_search_builder()
+ .with_vector_column('embedding')
+ .with_query_vector([1.0, 0.0, 0.0, 0.0])
+ .with_limit(6)
+ .execute_local())
+ baseline_ids = sorted(list(baseline.results()))
+ print(f"Baseline vector search ids={baseline_ids}")
+ self.assertEqual(baseline_ids, [0, 1, 2, 3, 4, 5])
+
+ # Filtered search — id >= 3 should restrict vector search to rows
+ # {3,4,5}, so top_k results must be a subset of that.
+ pb = PredicateBuilder(table.fields)
+ filter_pred = pb.greater_or_equal('id', 3)
+ filtered = (table.new_vector_search_builder()
+ .with_vector_column('embedding')
+ .with_query_vector([1.0, 0.0, 0.0, 0.0])
+ .with_limit(6)
+ .with_filter(filter_pred)
+ .execute_local())
+ filtered_ids = sorted(list(filtered.results()))
+ print(f"Filtered (id >= 3) vector search ids={filtered_ids}")
+ self.assertEqual(filtered_ids, [3, 4, 5])
+
+ # Narrower filter — id in {5} — only row 5 should survive.
+ filter_eq = pb.equal('id', 5)
+ eq_result = (table.new_vector_search_builder()
+ .with_vector_column('embedding')
+ .with_query_vector([1.0, 0.0, 0.0, 0.0])
+ .with_limit(6)
+ .with_filter(filter_eq)
+ .execute_local())
+ eq_ids = sorted(list(eq_result.results()))
+ print(f"Filtered (id == 5) vector search ids={eq_ids}")
+ self.assertEqual(eq_ids, [5])
+
def test_read_blob_after_alter_and_compact(self):
table = self.catalog.get_table('default.blob_alter_compact_test')
read_builder = table.new_read_builder()
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py
new file mode 100644
index 0000000000..454760d982
--- /dev/null
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -0,0 +1,525 @@
+################################################################################
+# 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.
+################################################################################
+
+"""Tests for VectorSearch + scalar predicate pre-filter wiring in pypaimon.
+
+Each test protects a distinct behavior introduced by this feature; no
+redundancy.
+"""
+
+import unittest
+from typing import List
+from unittest import mock
+
+from pypaimon.common.predicate import Predicate
+from pypaimon.common.predicate_builder import PredicateBuilder
+from pypaimon.globalindex.global_index_meta import GlobalIndexMeta
+from pypaimon.globalindex.global_index_result import GlobalIndexResult
+from pypaimon.globalindex.vector_search_result import ScoredGlobalIndexResult
+from pypaimon.index.index_file_meta import IndexFileMeta
+from pypaimon.manifest.index_manifest_entry import IndexManifestEntry
+from pypaimon.schema.data_types import AtomicType, DataField
+from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.table.source.vector_search_builder import VectorSearchBuilderImpl
+from pypaimon.utils.roaring_bitmap import RoaringBitmap64
+
+
+# ----------------------------- table stubs ---------------------------------
+
+
+class _StubSchema:
+ def __init__(self):
+ self.options = {}
+
+
+class _StubTable:
+ """Minimal FileStoreTable stand-in."""
+
+ def __init__(self, fields, entries, partition_fields=None):
+ self.fields = fields
+ self.partition_keys_fields = partition_fields or []
+ self.partition_keys: List[str] = [
+ f.name for f in self.partition_keys_fields]
+ self.table_schema = _StubSchema()
+ self.file_io = object()
+ self._entries = entries
+
+ def tag_manager(self):
+ return None
+
+ def snapshot_manager(self):
+ return None
+
+ def path_factory(self):
+ class _P:
+ def global_index_path_factory(self_inner):
+ class _G:
+ def index_path(self_g):
+ return "/tmp/unused"
+ return _G()
+ return _P()
+
+
+def _field(fid, name, dtype="INT"):
+ return DataField(id=fid, name=name,
+ type=AtomicType(dtype), description="")
+
+
+def _entry(partition_row, field_id, index_type, file_name,
+ row_range_start, row_range_end, external_path=None):
+ meta = GlobalIndexMeta(
+ row_range_start=row_range_start,
+ row_range_end=row_range_end,
+ index_field_id=field_id,
+ index_meta=b"",
+ )
+ index_file = IndexFileMeta(
+ index_type=index_type,
+ file_name=file_name,
+ file_size=1,
+ row_count=row_range_end - row_range_start + 1,
+ global_index_meta=meta,
+ external_path=external_path,
+ )
+ return IndexManifestEntry(kind=0, partition=partition_row, bucket=0,
+ index_file=index_file)
+
+
+def _patch_snapshot(testcase, entries):
+ """Stub IndexFileHandler.scan + snapshot resolution."""
+
+ def _scan(snapshot, entry_filter=None):
+ if entry_filter is None:
+ return list(entries)
+ return [e for e in entries if entry_filter(e)]
+
+ testcase._scan_patch = mock.patch(
+ "pypaimon.index.index_file_handler.IndexFileHandler.scan",
+ autospec=True, side_effect=lambda self_, s, f=None: _scan(s, f))
+ testcase._scan_patch.start()
+ testcase._travel_patch = mock.patch(
+ "pypaimon.snapshot.time_travel_util.TimeTravelUtil.try_travel_to_snapshot",
+ return_value=object())
+ testcase._travel_patch.start()
+
+
+# ----------------------------- tests ---------------------------------------
+
+
+class VectorSearchFilterTest(unittest.TestCase):
+ """Non-partitioned wiring: scan + read + external_path plumbing."""
+
+ def setUp(self):
+ self.id_field = _field(0, "id")
+ self.embedding_field = _field(1, "embedding", "FLOAT")
+ # 2 vector files ([0,4], [5,9]) + 1 btree on `id` covering [0,9] with
+ # an external_path so we can assert external_path is threaded through.
+ self.entries = [
+ _entry(None, field_id=1, index_type="lumina-vector-ann",
+ file_name="vec-0.index",
+ row_range_start=0, row_range_end=4,
+ external_path="oss://bucket/vec-0.index"),
+ _entry(None, field_id=1, index_type="lumina-vector-ann",
+ file_name="vec-1.index",
+ row_range_start=5, row_range_end=9,
+ external_path="oss://bucket/vec-1.index"),
+ _entry(None, field_id=0, index_type="btree",
+ file_name="id-btree-0.index",
+ row_range_start=0, row_range_end=9,
+ external_path="oss://bucket/id-btree-0.index"),
+ ]
+ self.table = _StubTable(fields=[self.id_field, self.embedding_field],
+ entries=self.entries)
+ _patch_snapshot(self, self.entries)
+
+ def tearDown(self):
+ mock.patch.stopall()
+
+ def _builder(self, filter_pred=None):
+ b = (VectorSearchBuilderImpl(self.table)
+ .with_vector_column("embedding")
+ .with_query_vector([1.0, 0.0, 0.0, 0.0])
+ .with_limit(3))
+ if filter_pred is not None:
+ b = b.with_filter(filter_pred)
+ return b
+
+ def test_scan_attaches_overlapping_scalar_index_files(self):
+ """``with_filter`` + scan: each vector-range split must carry the
+ scalar index files whose row range overlaps it."""
+ filter_pred = Predicate(method="greaterOrEqual", index=0, field="id",
+ literals=[5])
+ splits = self._builder(filter_pred).new_vector_search_scan().scan().splits()
+
+ self.assertEqual(2, len(splits))
+ splits_sorted = sorted(splits, key=lambda s: s.row_range_start)
+ for s in splits_sorted:
+ self.assertEqual(1, len(s.vector_index_files))
+ self.assertEqual(["id-btree-0.index"],
+ [f.file_name for f in s.scalar_index_files])
+ self.assertEqual((0, 4),
+ (splits_sorted[0].row_range_start,
+ splits_sorted[0].row_range_end))
+ self.assertEqual((5, 9),
+ (splits_sorted[1].row_range_start,
+ splits_sorted[1].row_range_end))
+
+ def test_read_threads_prefilter_bitmap_as_include_row_ids(self):
+ """preFilter bitmap from scanner.scan(filter) must reach each split's
+ VectorSearch, offset-rebased to local coords by OffsetGlobalIndexReader.
+ Also: the vector reader's io_meta carries external_path."""
+ filter_pred = Predicate(method="greaterOrEqual", index=0, field="id",
+ literals=[5])
+ scan_plan = self._builder(filter_pred).new_vector_search_scan().scan()
+
+ bitmap = RoaringBitmap64()
+ for rid in range(5, 10):
+ bitmap.add(rid)
+ scanner = mock.MagicMock()
+ scanner.scan.return_value = GlobalIndexResult.create(lambda: bitmap)
+
+ captured_searches = []
+ captured_io_metas = []
+
+ def _capture_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ captured_io_metas.append(list(index_io_meta_list))
+
+ class _FakeReader:
+ def visit_vector_search(self_inner, vs):
+ captured_searches.append(vs)
+ return ScoredGlobalIndexResult.create_empty()
+
+ def close(self_inner):
+ pass
+
+ def __enter__(self_inner):
+ return self_inner
+
+ def __exit__(self_inner, *a):
+ return False
+ return _FakeReader()
+
+ with mock.patch(
+ "pypaimon.globalindex.global_index_scanner.GlobalIndexScanner.create",
+ return_value=scanner), \
+ mock.patch(
+ "pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_capture_create):
+ self._builder(filter_pred).new_vector_search_read().read_plan(scan_plan)
+
+ # Pre-filter happened once with our filter.
+ self.assertEqual(1, scanner.scan.call_count)
+ self.assertIs(filter_pred, scanner.scan.call_args[0][0])
+
+ # [0,4] sees empty local bitmap; [5,9] sees {0..4}.
+ self.assertEqual(
+ [0, 5],
+ sorted(vs.include_row_ids.cardinality()
+ for vs in captured_searches))
+
+ # Vector reader io_meta carries external_path from IndexFileMeta.
+ seen_paths = {meta.external_path
+ for metas in captured_io_metas
+ for meta in metas}
+ self.assertEqual(
+ {"oss://bucket/vec-0.index", "oss://bucket/vec-1.index"},
+ seen_paths)
+
+ def test_scanner_threads_external_path_to_btree_reader(self):
+ """GlobalIndexScanner (backing _pre_filter) must thread external_path
+ onto the GlobalIndexIOMeta handed to the btree reader factory."""
+ from pypaimon.globalindex.global_index_scanner import GlobalIndexScanner
+
+ scalar_file = self.entries[2].index_file
+ scanner = GlobalIndexScanner(
+ fields=self.table.fields,
+ file_io=self.table.file_io,
+ index_path="/unused/index-path",
+ index_files=[scalar_file],
+ )
+ captured = []
+
+ class _FakeBTreeReader:
+ def __init__(self_inner, key_serializer, file_io, index_path,
+ io_meta):
+ captured.append(io_meta)
+
+ def close(self_inner):
+ pass
+
+ try:
+ with mock.patch(
+ "pypaimon.globalindex.btree.BTreeIndexReader",
+ _FakeBTreeReader):
+ list(scanner._evaluator._readers_function(self.id_field))
+ finally:
+ scanner.close()
+
+ self.assertEqual(1, len(captured))
+ self.assertEqual("oss://bucket/id-btree-0.index",
+ captured[0].external_path)
+
+
+class VectorSearchMultiShardScalarTest(unittest.TestCase):
+ """Scalar pre-filter across multiple btree shards of the same field.
+
+ Exercises the real GlobalIndexScanner reader-construction path (with
+ OffsetGlobalIndexReader + UnionGlobalIndexReader wrapping) so that:
+ - Local row ids from each shard are rebased to the global row-id space
+ before being unioned.
+ - An empty first shard does NOT short-circuit subsequent shards.
+ """
+
+ def test_hit_only_in_later_shard_returns_global_row_id(self):
+ from pypaimon.globalindex.global_index_reader import GlobalIndexReader
+ from pypaimon.globalindex.global_index_result import GlobalIndexResult
+ from pypaimon.globalindex.global_index_scanner import (
+ GlobalIndexScanner,
+ )
+
+ id_field = _field(0, "id")
+ emb_field = _field(1, "embedding", "FLOAT")
+
+ # Two btree shards: [0,4] and [5,9]. Predicate: id == 7, which only
+ # exists in shard [5,9] at local row id 2 (= 7 - 5).
+ shard_a = _entry(None, field_id=0, index_type="btree",
+ file_name="id-0.index",
+ row_range_start=0, row_range_end=4).index_file
+ shard_b = _entry(None, field_id=0, index_type="btree",
+ file_name="id-1.index",
+ row_range_start=5, row_range_end=9).index_file
+ table = _StubTable(fields=[id_field, emb_field], entries=[])
+
+ # Stub BTreeIndexReader: shard_a returns empty, shard_b returns {2}
+ # (local row id). After Offset wrapping the scanner should emit {7}.
+ class _StubBTreeReader(GlobalIndexReader):
+ def __init__(self_inner, key_serializer, file_io, index_path,
+ io_meta):
+ self_inner._file = io_meta.file_name
+
+ def visit_equal(self_inner, field_ref, literal):
+ bm = RoaringBitmap64()
+ if self_inner._file == "id-1.index":
+ bm.add(2) # local offset inside [5,9]
+ return GlobalIndexResult.create(lambda b=bm: b)
+
+ def close(self_inner):
+ pass
+
+ with mock.patch("pypaimon.globalindex.btree.BTreeIndexReader",
+ _StubBTreeReader):
+ scanner = GlobalIndexScanner(
+ fields=table.fields,
+ file_io=table.file_io,
+ index_path="/unused",
+ index_files=[shard_a, shard_b],
+ )
+ try:
+ result = scanner.scan(
+ Predicate(method="equal", index=0, field="id",
+ literals=[7]))
+ finally:
+ scanner.close()
+
+ self.assertIsNotNone(result)
+ hits = sorted(list(result.results()))
+ # Must be the GLOBAL row id (7 = 5 + 2), not the local (2).
+ # Must not be empty despite shard_a being empty (no short-circuit).
+ self.assertEqual([7], hits)
+
+ def test_tantivy_fulltext_index_is_dispatched_by_scanner(self):
+ """Non-btree scalar global indexes (tantivy-fulltext, etc.) must be
+ instantiated by GlobalIndexScanner — previously only 'btree' was
+ handled and everything else was silently dropped, making text-column
+ pre-filter a no-op."""
+ from pypaimon.globalindex.global_index_result import GlobalIndexResult
+ from pypaimon.globalindex.global_index_scanner import (
+ GlobalIndexScanner,
+ )
+
+ name_field = _field(0, "name", "STRING")
+ emb_field = _field(1, "embedding", "FLOAT")
+ tantivy_shard = _entry(
+ None, field_id=0, index_type="tantivy-fulltext",
+ file_name="name-ft.index",
+ row_range_start=0, row_range_end=9,
+ external_path="oss://bucket/name-ft.index").index_file
+ table = _StubTable(fields=[name_field, emb_field], entries=[])
+
+ captured_ctor_args = []
+ visit_calls = []
+
+ class _StubTantivyReader:
+ def __init__(self_inner, file_io, index_path, io_metas):
+ captured_ctor_args.append(
+ (file_io, index_path, list(io_metas)))
+
+ def visit_equal(self_inner, field_ref, literal):
+ visit_calls.append(("equal", literal))
+ bm = RoaringBitmap64()
+ bm.add(4)
+ return GlobalIndexResult.create(lambda b=bm: b)
+
+ def close(self_inner):
+ pass
+
+ with mock.patch(
+ "pypaimon.globalindex.tantivy.TantivyFullTextGlobalIndexReader",
+ _StubTantivyReader):
+ scanner = GlobalIndexScanner(
+ fields=table.fields,
+ file_io=table.file_io,
+ index_path="/unused",
+ index_files=[tantivy_shard],
+ )
+ try:
+ result = scanner.scan(
+ Predicate(method="equal", index=0, field="name",
+ literals=["x"]))
+ finally:
+ scanner.close()
+
+ # Tantivy reader was instantiated (it would NOT be before this fix).
+ self.assertEqual(1, len(captured_ctor_args))
+ _, _, io_metas = captured_ctor_args[0]
+ self.assertEqual("oss://bucket/name-ft.index",
+ io_metas[0].external_path)
+ # visit_equal was dispatched all the way through evaluator → union →
+ # offset → stub tantivy reader.
+ self.assertEqual([("equal", "x")], visit_calls)
+ # Row id 4 is inside [0,9] so offset rebase is a no-op.
+ self.assertEqual([4], sorted(list(result.results())))
+
+ def test_like_predicate_is_dispatched_to_reader(self):
+ """Evaluator must dispatch ``like`` to reader.visit_like — otherwise
+ the pre-filter is silently skipped and vector search returns rows
+ that violate the predicate."""
+ from pypaimon.globalindex.global_index_reader import GlobalIndexReader
+ from pypaimon.globalindex.global_index_result import GlobalIndexResult
+ from pypaimon.globalindex.global_index_scanner import (
+ GlobalIndexScanner,
+ )
+
+ name_field = _field(0, "name", "STRING")
+ emb_field = _field(1, "embedding", "FLOAT")
+ shard = _entry(None, field_id=0, index_type="btree",
+ file_name="name-0.index",
+ row_range_start=0, row_range_end=9).index_file
+ table = _StubTable(fields=[name_field, emb_field], entries=[])
+
+ observed_calls = []
+
+ class _StubBTreeReader(GlobalIndexReader):
+ def __init__(self_inner, key_serializer, file_io, index_path,
+ io_meta):
+ pass
+
+ def visit_like(self_inner, field_ref, literal):
+ observed_calls.append(("like", literal))
+ bm = RoaringBitmap64()
+ bm.add(3) # local, will be offset-rebased to 3 (range starts at 0)
+ return GlobalIndexResult.create(lambda b=bm: b)
+
+ def close(self_inner):
+ pass
+
+ with mock.patch("pypaimon.globalindex.btree.BTreeIndexReader",
+ _StubBTreeReader):
+ scanner = GlobalIndexScanner(
+ fields=table.fields,
+ file_io=table.file_io,
+ index_path="/unused",
+ index_files=[shard],
+ )
+ try:
+ result = scanner.scan(
+ Predicate(method="like", index=0, field="name",
+ literals=["abc%"]))
+ finally:
+ scanner.close()
+
+ self.assertEqual([("like", "abc%")], observed_calls)
+ self.assertIsNotNone(result)
+ self.assertEqual([3], sorted(list(result.results())))
+
+
+class VectorSearchPartitionedFilterTest(unittest.TestCase):
+ """Partitioned-table paths: with_filter auto-split + partition-filter
+ input validation."""
+
+ def setUp(self):
+ self.pt_field = _field(0, "pt")
+ self.id_field = _field(1, "id")
+ self.embedding_field = _field(2, "embedding", "FLOAT")
+
+ partition_pt1 = GenericRow([1], [self.pt_field])
+ partition_pt2 = GenericRow([2], [self.pt_field])
+ self.entries = [
+ _entry(partition_pt1, field_id=2,
+ index_type="lumina-vector-ann",
+ file_name="vec-pt1.index",
+ row_range_start=0, row_range_end=4),
+ _entry(partition_pt2, field_id=2,
+ index_type="lumina-vector-ann",
+ file_name="vec-pt2.index",
+ row_range_start=5, row_range_end=9),
+ ]
+ self.table = _StubTable(
+ fields=[self.pt_field, self.id_field, self.embedding_field],
+ partition_fields=[self.pt_field],
+ entries=self.entries)
+ _patch_snapshot(self, self.entries)
+
+ def tearDown(self):
+ mock.patch.stopall()
+
+ def test_with_filter_auto_splits_and_prunes_wrong_partition(self):
+ """A normal full-row predicate ``pt == 2`` must (a) be auto-split
+ into _partition_filter with indices re-based to the partition-only
+ row, and (b) drop pt=1 entries during manifest pruning."""
+ pb = PredicateBuilder(self.table.fields)
+ builder = (VectorSearchBuilderImpl(self.table)
+ .with_vector_column("embedding")
+ .with_query_vector([1.0, 0.0, 0.0, 0.0])
+ .with_limit(3)
+ .with_filter(pb.equal("pt", 2)))
+
+ # Partition filter's leaf index points into the partition-only row.
+ self.assertEqual("pt", builder._partition_filter.field)
+ self.assertEqual(0, builder._partition_filter.index)
+
+ splits = builder.new_vector_search_scan().scan().splits()
+ self.assertEqual(1, len(splits))
+ self.assertEqual(["vec-pt2.index"],
+ [f.file_name for f in splits[0].vector_index_files])
+
+ def test_with_partition_filter_rejects_non_partition_field(self):
+ """Non-partition conjuncts would be silently dropped by the extractor,
+ producing wrong results; the API must refuse them up front."""
+ pb = PredicateBuilder(self.table.fields)
+ builder = VectorSearchBuilderImpl(self.table)
+ with self.assertRaises(ValueError) as ctx:
+ builder.with_partition_filter(
+ PredicateBuilder.and_predicates(
+ [pb.equal("pt", 1), pb.equal("id", 5)]))
+ self.assertIn("non-partition", str(ctx.exception))
+
+
+if __name__ == "__main__":
+ unittest.main()