From 446700ad17118b1d89eba1ca76ed3cc7787b04e2 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Wed, 21 Oct 2026 22:16:25 +0800 Subject: [PATCH] [python] Fix IVF-PQ searches with sparse filters Use the existing Vindex batch API for single queries and singleton batches so filters containing live NULL vectors can return fewer hits than requested. Keep Python execution independent of the Native backend and preserve the existing score conversion and invalid-ID filtering. Add real SDK and deletion-vector regressions for sparse and empty results. Signed-off-by: QuakeWang --- .../vindex_vector_global_index_reader.py | 7 +- .../tests/native_vector_index_build_test.py | 3 +- .../tests/vector_index_dv_build_test.py | 25 ++++++ .../tests/vindex_sparse_filter_test.py | 76 +++++++++++++++++++ 4 files changed, 104 insertions(+), 7 deletions(-) create mode 100644 paimon-python/pypaimon/tests/vindex_sparse_filter_test.py diff --git a/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py b/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py index ee97f685796d..1b7f8a88d867 100644 --- a/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py +++ b/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py @@ -59,11 +59,8 @@ def _run_search(self, vectors, limit, include_row_ids, query_options): params = _search_params(self._search_params_type, options, effective_k) filter_bytes = _filter_bytes(include_row_ids) - if n == 1: - ids, distances = self._reader.search( - queries[0], params, filter_bytes=filter_bytes) - return [_result_from_scores(ids, distances, self._metadata.metric)] - + # The batch API pads missing hits with -1, including for one query. + # The scalar FFI rejects short IVF-PQ results from sparse filters. if hasattr(self._reader, "search_batch"): batch_result = self._reader.search_batch( np.ascontiguousarray(queries), diff --git a/paimon-python/pypaimon/tests/native_vector_index_build_test.py b/paimon-python/pypaimon/tests/native_vector_index_build_test.py index 38b56791a6e2..5bd5c43025b4 100644 --- a/paimon-python/pypaimon/tests/native_vector_index_build_test.py +++ b/paimon-python/pypaimon/tests/native_vector_index_build_test.py @@ -137,8 +137,7 @@ def test_native_search_filters_live_null_vectors_with_deletion_vectors(rest_cata result = (_native_table(table).new_vector_search_builder().with_vector_column('embedding') .with_query_vector([1., 1.]).with_limit(64).with_options(options).execute_local()) assert sorted(result.row_ids()) == [1, 4] - # Python table-level planning also uses the core low-level index reader, - # avoiding the third-party scalar FFI's fixed-result-length restriction. + # Python search must also allow fewer hits than the live-row filter's size. _assert_ids(table, [1, 4], kind) 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 da06652e56d9..87325c7e08fa 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,31 @@ def test_native_search_excludes_deletions_before_and_after_build(table, metric, assert build(table, metric) == 0 +@pytest.mark.python_plan +@pytest.mark.python_write +@pytest.mark.python_commit +@pytest.mark.parametrize('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()] + 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 000000000000..1a47407d144d --- /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 + + +@pytest.fixture(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() + + +@pytest.mark.parametrize('batch_size', [None, 1, 2], ids=['single', 'batch-one', 'batch-two']) +@pytest.mark.parametrize('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