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 ae2ea66a43 [python] Add regression coverage for sparse IVF-PQ searches
(#10726)
ae2ea66a43 is described below
commit ae2ea66a433630ff0dec8253594ee2980546aa3a
Author: QuakeWang <[email protected]>
AuthorDate: Mon Oct 26 12:20:15 2026 +0800
[python] Add regression coverage for sparse IVF-PQ searches (#10726)
---
.../pypaimon/tests/vector_index_dv_build_test.py | 26 ++++++++
.../pypaimon/tests/vindex_sparse_filter_test.py | 76 ++++++++++++++++++++++
2 files changed, 102 insertions(+)
diff --git a/paimon-python/pypaimon/tests/vector_index_dv_build_test.py
b/paimon-python/pypaimon/tests/vector_index_dv_build_test.py
index da06652e56..f8766fbc9f 100644
--- a/paimon-python/pypaimon/tests/vector_index_dv_build_test.py
+++ b/paimon-python/pypaimon/tests/vector_index_dv_build_test.py
@@ -114,6 +114,32 @@ def
test_native_search_excludes_deletions_before_and_after_build(table, metric,
assert build(table, metric) == 0
[email protected]_plan
[email protected]_write
[email protected]_commit
[email protected]('batch', [False, True])
+def test_python_ivfpq_search_keeps_live_null_vectors_out_of_results(tmp_path,
batch):
+ pytest.importorskip('paimon_vindex')
+ schema = pa.schema([('id', pa.int64()), ('embedding',
pa.list_(pa.float32()))])
+ table = pm.connect(options={'warehouse': str(tmp_path)}).create_table(
+ 'vectors', schema=schema, options={
+ 'file.format': 'parquet', 'vector.file.format': 'parquet',
'vector-index.search-mode': 'fast'})
+ table.add(pa.table({
+ 'id': list(range(5)), 'embedding': [[1., 0.], [0., 1.], None, [1.,
1.], [2., 2.]]}, schema=schema))
+ table.delete('id = 0 OR id = 3')
+ assert table.create_index('embedding', 'ivf-pq', options={
+ 'ivf-pq.dimension': '2', 'ivf-pq.nlist': '1', 'ivf-pq.pq.m': '1',
+ 'ivf-pq.distance.metric': 'l2'}) == 1
+ raw_table = table.raw_table
+ query =
(raw_table.new_batch_vector_search_builder().with_query_vectors([[1., 1.], [1.,
1.]]) if batch
+ else raw_table.new_vector_search_builder().with_query_vector([1.,
1.]))
+
query.with_vector_column('embedding').with_limit(64).with_option('ivf.nprobe',
'1')
+ results = query.execute_batch_local() if batch else [query.execute_local()]
+ assert len(results) == (2 if batch else 1)
+ for result in results:
+ assert sorted(result.results()) == [1, 4]
+
+
def test_native_build_and_query_keep_snapshot_semantics(table):
table.delete('id = 1 OR id = 6')
assert build(table) == 2
diff --git a/paimon-python/pypaimon/tests/vindex_sparse_filter_test.py
b/paimon-python/pypaimon/tests/vindex_sparse_filter_test.py
new file mode 100644
index 0000000000..1a47407d14
--- /dev/null
+++ b/paimon-python/pypaimon/tests/vindex_sparse_filter_test.py
@@ -0,0 +1,76 @@
+# 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.
+
+"""Python Vindex searches allow fewer valid hits than the requested limit."""
+
+import io
+
+import numpy as np
+import pytest
+
+from pypaimon.execution.python.read.search.vector_search_read import
_create_vector_reader
+from pypaimon.globalindex.batch_vector_search import BatchVectorSearch
+from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta
+from pypaimon.globalindex.vector_search import VectorSearch
+from pypaimon.utils.roaring_bitmap import RoaringBitmap64
+
+
[email protected](params=['l2', 'cosine', 'inner_product'])
+def index_bytes(request):
+ vindex = pytest.importorskip('paimon_vindex')
+ vectors = np.array([[1., 0.], [0., 1.], [1., 1.], [2., 2.]],
dtype=np.float32)
+ options = {'index.type': 'ivf_pq', 'dimension': '2', 'nlist': '1',
+ 'pq.m': '1', 'metric': request.param}
+ output = io.BytesIO()
+ with vindex.VectorIndexTrainer.train(options, vectors) as training:
+ with vindex.VectorIndexWriter(training) as writer:
+ writer.add_vectors(np.array([0, 1, 3, 4], dtype=np.int64), vectors)
+ writer.write(output)
+ return output.getvalue()
+
+
[email protected]('batch_size', [None, 1, 2], ids=['single',
'batch-one', 'batch-two'])
[email protected]('selected, expected', [([1, 2, 4], [1, 4]), ([2],
[]), ([], [])],
+ ids=['live-null', 'null-only', 'empty'])
+def test_python_ivfpq_sparse_filter(index_bytes, batch_size, selected,
expected):
+ class FileIO:
+ def new_input_stream(self, path):
+ self.stream = io.BytesIO(index_bytes)
+ return self.stream
+
+ file_io = FileIO()
+ include = RoaringBitmap64()
+ for row_id in selected:
+ include.add(row_id)
+ # Row 2 is a live NULL vector; rows 0 and 3 are deleted by the query
filter.
+ query_type = BatchVectorSearch if batch_size else VectorSearch
+ vectors = [[1., 1.], [0., 1.]][:batch_size] if batch_size else [1., 1.]
+ query = query_type(vectors, 64, 'embedding', options={'ivf.nprobe': '1'})
+ reader = _create_vector_reader('ivf-pq', file_io, '/indexes',
[GlobalIndexIOMeta('file', len(index_bytes))])
+ with reader:
+ search = reader.visit_batch_vector_search if batch_size else
reader.visit_vector_search
+ reference = search(query).result()
+ results = search(query.with_include_row_ids(include)).result()
+ if not batch_size:
+ reference, results = [reference], [results]
+ assert len(results) == len(reference) == (batch_size or 1)
+ for result, unfiltered in zip(results, reference):
+ assert ([] if result is None else sorted(result.results())) ==
expected
+ if result is not None:
+ assert {row_id: result.score_getter()(row_id) for row_id in
expected} == pytest.approx(
+ {row_id: unfiltered.score_getter()(row_id) for row_id in
expected})
+ assert file_io.stream.closed
