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 3c4be51aac [python] Share final lookups across batch vector queries
(#9935)
3c4be51aac is described below
commit 3c4be51aacf405a32803b8d66a10515e81017caa
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]
