[python] Fix IVF-PQ searches with sparse filters

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),

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())
     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)
 
 

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

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