This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 1e9c720f9614175ae94f45a6dfd38ad2e40a34c7 Author: chaoyang <[email protected]> AuthorDate: Fri Sep 18 12:00:42 2026 +0800 [python] Share final lookups across batch vector queries (#9935) --- docs/docs/pypaimon/multimodal-search.md | 3 + paimon-python/pypaimon/multimodal/query.py | 40 ++++- .../pypaimon/tests/batch_vector_lookup_test.py | 181 +++++++++++++++++++++ 3 files changed, 220 insertions(+), 4 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 2a2940fe1c..034162016e 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -186,6 +186,9 @@ Dictionary vector routes also accept `anns_field`, `data`, and `param` aliases. Use `search_vectors` for multiple query vectors against one vector column. It returns one result set for each input vector, preserving input order. +Batch results share a lookup of their combined row IDs, so overlapping results +do not require reading the same row separately for each query. The limit and +post-filter still apply to each query's result set. ```python batch_neighbors = ( diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 9f5d92c19e..9cfea57e6f 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -634,10 +634,42 @@ class BatchVectorQuery(_PreFilterQuery): def to_arrow(self): query = self._for_execution() - return [ - query._read_global_index_result(result) - for result in query._execute_batch_vector(query) - ] + return query._read_batch_results(query._execute_batch_vector(query)) + + def _read_batch_results(self, results): + from pypaimon.globalindex.global_index_result import GlobalIndexResult + from pypaimon.utils.roaring_bitmap import RoaringBitmap64 + + if len(results) <= 1 or not self._configured_read_builder().read_type(): + return [self._read_global_index_result(result) for result in results] + + row_ids = RoaringBitmap64() + for result in results: + row_ids = RoaringBitmap64.or_(row_ids, result.results()) + + lookup = copy(self) + # Each result is already top-k. A shared read must not apply that limit + # to the union; where() still filters the selected rows during lookup. + lookup._limit = None + projection = self._effective_projection() + lookup._projection = list(projection) if projection else [f.name for f in self._table.fields] + added_row_id = SpecialFields.ROW_ID.name not in lookup._projection + if added_row_id: + lookup._projection.append(SpecialFields.ROW_ID.name) + fields = lookup._configured_read_builder().read_type() + row_id_column = next(i for i, field in enumerate(fields) if field.id == SpecialFields.ROW_ID.id) + table = lookup._read_global_index_result(GlobalIndexResult.create(row_ids)) + positions = {row_id: i for i, row_id in enumerate(table.column(row_id_column).to_pylist())} + if added_row_id: + table = table.select(list(range(table.num_columns - 1))) + output = [] + for result in results: + # Keep the physical read order, rather than imposing score or row-id order. + selected = sorted(positions[row_id] for row_id in result.results() if row_id in positions) + if self._limit is not None: + selected = selected[:self._limit] + output.append(table.take(pa.array(selected, type=pa.int64()))) + return output def to_pandas(self): return [table.to_pandas() for table in self.to_arrow()] diff --git a/paimon-python/pypaimon/tests/batch_vector_lookup_test.py b/paimon-python/pypaimon/tests/batch_vector_lookup_test.py new file mode 100644 index 0000000000..aa7877a749 --- /dev/null +++ b/paimon-python/pypaimon/tests/batch_vector_lookup_test.py @@ -0,0 +1,181 @@ +# 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. + +from unittest.mock import patch + +import pyarrow as pa +import pytest + +import pypaimon.multimodal as pm +from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.multimodal.query import ScanQuery +from pypaimon.read.table_read import TableRead + + [email protected](params=[1, 3]) +def docs(tmp_path, request): + schema = pa.schema([ + ("id", pa.int64()), ("category", pa.string()), ("embedding", pa.list_(pa.float32(), 2)), + ("info", pa.struct([("label", pa.string()), ("value", pa.int32())])), + ("attrs", pa.map_(pa.string(), pa.string())), ("payload", pa.large_binary()), + ]) + table = pm.connect(options={"warehouse": str(tmp_path)}).create_table( + "docs", schema=schema, partitioned=["category"], options={ + "file.format": "parquet", "vector.file.format": "parquet", + "vector-index.search-mode": "full", "read.parallelism": str(request.param), + "read.batch-size": "2", "blob-as-descriptor": "false", + }) + for start in (0, 4, 8): + table.add(pa.Table.from_pylist([{ + "id": i, "category": "a" if i % 2 else "b", "embedding": [float(i), 0.0], + "info": {"label": str(i), "value": i}, "attrs": {"x": str(i)}, + "payload": ("payload-%d" % i).encode(), + } for i in range(start, start + 4)], schema=schema)) + return table + + +def configure(query, projection, with_row_id=False, post_filter=None): + if projection is not None: + query.select(projection) + if with_row_id: + query.with_row_id() + if post_filter is not None: + query.where(post_filter) + return query.limit(3) + + [email protected]("projection", [ + None, [], ["id"], ["payload"], ["_ROW_ID", "id"], ["info", "attrs['x']"], ["missing"]]) [email protected]("with_row_id", [False, True]) +def test_batch_lookup_matches_individual_queries(docs, projection, with_row_id): + vectors = [[1.0, 0.0], [9.0, 0.0], [1.0, 0.0]] + expected = [configure(docs.search(vector), projection, with_row_id).to_arrow() for vector in vectors] + query = configure(docs.search_vectors(vectors), projection, with_row_id) + original = TableRead.to_arrow + reads = [] + + def read(table_read, splits, *args, **kwargs): + reads.append(len(splits)) + return original(table_read, splits, *args, **kwargs) + + with patch.object(TableRead, "to_arrow", read): + actual = query.to_arrow() + assert len(actual) == len(expected) + for left, right in zip(expected, actual): + assert left.schema == right.schema + assert left.to_pylist() == right.to_pylist() + # Preserve the existing zero-column behavior of invalid projections. + assert len(reads) == (3 if projection == ["missing"] and not with_row_id else 1) + + +def test_shared_lookup_keeps_pre_and_post_filters_and_per_query_limits(docs): + vectors = [[1.0, 0.0], [9.0, 0.0], [1.0, 0.0]] + options = {"pre_filter": "category = 'b'"} + expected = [configure(docs.search(vector, **options), ["id"], post_filter="id >= 5").to_arrow() + for vector in vectors] + actual = configure(docs.search_vectors(vectors, **options), ["id"], post_filter="id >= 5").to_arrow() + assert [result.to_pylist() for result in actual] == [result.to_pylist() for result in expected] + assert actual[0].num_rows == 0 + assert actual[1].num_rows == 3 + + +def test_empty_membership_missing_rows_and_scores_are_independent(docs): + row_ids = [row["_ROW_ID"] for row in docs.scan().select(["id"]).with_row_id().to_list()] + results = [DictBasedScoredIndexResult(values) for values in ( + {row_ids[0]: 0.9, row_ids[1]: 0.7}, {}, {row_ids[0]: 0.1}, {10000: 1.0})] + query = docs.search_vectors([[0.0, 0.0]] * len(results)).select(["id"]).limit(3)._for_execution() + expected = [query._read_global_index_result(result) for result in results] + actual = query._read_batch_results(results) + assert [table.to_pylist() for table in actual] == [table.to_pylist() for table in expected] + assert all(table.schema == expected[0].schema for table in actual) + assert results[0].score_getter()(row_ids[0]) == 0.9 + assert results[2].score_getter()(row_ids[0]) == 0.1 + + +def test_all_empty_results_preserve_projection_schema(docs): + query = docs.search_vectors([[0.0, 0.0], [11.0, 0.0]], pre_filter="category = 'missing'") + actual = query.select(["id", "payload"]).with_row_id().limit(2).to_arrow() + expected = docs.search([0.0, 0.0], pre_filter="category = 'missing'").select( + ["id", "payload"]).with_row_id().limit(2).to_arrow() + assert len(actual) == 2 + assert all(table.equals(expected) for table in actual) + + [email protected]("change", ["update", "delete"]) +def test_shared_lookup_uses_search_snapshot_and_query_can_be_reused(docs, change): + query = docs.search_vectors([[0.0, 0.0], [11.0, 0.0]]).select(["id", "info"]).limit(1) + expected = query.to_list() + lookup = ScanQuery._read_global_index_result + calls = [] + + def commit_before_lookup(execution, result): + calls.append(execution._table) + if change == "update": + docs.update("id = 0", {"info": {"label": "updated", "value": 99}}) + else: + docs.delete("id = 0") + return lookup(execution, result) + + with patch.object(ScanQuery, "_read_global_index_result", commit_before_lookup): + assert query.to_list() == expected + assert len(calls) == 1 + assert query.to_list() != expected + + [email protected]("selector", ["snapshot", "tag"]) +def test_shared_lookup_keeps_historical_views(docs, selector): + snapshot = docs.raw_table.snapshot_manager().get_latest_snapshot() + docs.raw_table.create_tag("before", snapshot_id=snapshot.id) + options = {"snapshot_id": snapshot.id} if selector == "snapshot" else {"tag_name": "before"} + query = docs.search_vectors([[0.0, 0.0], [11.0, 0.0]], **options).select(["id", "info"]).limit(1) + expected = query.to_list() + docs.update("id = 0", {"info": {"label": "new", "value": 99}}) + if selector == "tag": + docs.raw_table.file_io.delete(docs.raw_table.snapshot_manager().get_snapshot_path(snapshot.id)) + assert query.to_list() == expected + + +def test_single_query_keeps_single_lookup(docs): + expected = docs.search([1.0, 0.0]).select(["id"]).limit(3).to_arrow() + result = docs.search_vectors([[1.0, 0.0]]).select(["id"]).limit(3).to_arrow() + assert result[0].equals(expected) + + +def test_unsupported_nested_row_projection_still_raises(docs): + with pytest.raises(NotImplementedError, match="ROW nested-field projection"): + docs.search_vectors([[0.0, 0.0], [1.0, 0.0]]).select(["info.label"]).to_arrow() + + +def test_native_index_batch_uses_shared_lookup(docs): + pytest.importorskip("paimon_vindex") + docs.raw_table.copy({"deletion-vectors.enabled": "false"}).create_global_index( + "embedding", index_type="ivf-flat", options={ + "ivf-flat.dimension": "2", "ivf-flat.nlist": "1", "ivf-flat.distance.metric": "l2"}) + vectors = [[1.0, 0.0], [10.0, 0.0], [1.0, 0.0]] + expected = [docs.search(vector).select(["id"]).limit(3).to_list() for vector in vectors] + query = docs.search_vectors(vectors).select(["id"]).limit(3) + lookup = ScanQuery._read_global_index_result + calls = [] + + def read(execution, result): + calls.append(result.results().cardinality()) + return lookup(execution, result) + + with patch.object(ScanQuery, "_read_global_index_result", read): + actual = query.to_list() + assert actual == expected + assert calls == [6]
