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 2a5346f8e6 [python] Distribute batch vector refinement across Ray
workers (#9988)
2a5346f8e6 is described below
commit 2a5346f8e6649c5663926a6a6293d3ffe6a20343
Author: chaoyang <[email protected]>
AuthorDate: Sun Sep 20 13:44:34 2026 +0800
[python] Distribute batch vector refinement across Ray workers (#9988)
---
docs/docs/pypaimon/multimodal-search.md | 12 +-
paimon-python/pypaimon/multimodal/query.py | 4 +-
paimon-python/pypaimon/ray/batch_vector_search.py | 28 ++++
.../tests/ray_batch_vector_refinement_test.py | 145 +++++++++++++++++++++
.../pypaimon/tests/ray_batch_vector_search_test.py | 6 +-
5 files changed, 187 insertions(+), 8 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-search.md
b/docs/docs/pypaimon/multimodal-search.md
index 63cd120220..2431e48772 100644
--- a/docs/docs/pypaimon/multimodal-search.md
+++ b/docs/docs/pypaimon/multimodal-search.md
@@ -155,14 +155,16 @@ batch_neighbors = (
Each batch task handles all query vectors for one split: index workers reuse
an open shard across bounded query blocks, and raw workers stream each data
-split once with a separate top-k for each query. Batch refinement and the
shared
-final row lookup run on the driver. The entire batch uses one read snapshot.
+split once with a separate top-k for each query. For refinement, the driver
+selects each query's global candidates, then workers stream their union once
per
+read split, scoring each row only for the queries that selected it. Shared
final
+row lookup runs on the driver. The entire batch uses one read snapshot.
Candidate traffic and result memory grow with the number of query vectors;
split large query collections into smaller batches when necessary.
The driver fixes one read snapshot and plans the query. Workers search
individual
index shards and, when required by the table's search mode, scan unindexed
data.
-They return candidate row IDs and scores. For single-vector queries with
+They return candidate row IDs and scores. For vector queries with
refinement enabled, the driver first selects the global candidate set, then
workers read and score those candidate vectors using the index's persisted
metric. The driver merges their results and retains the local rules for
@@ -180,8 +182,8 @@ All workers must be able to access the table's storage.
Local filesystem paths
are suitable for a local Ray cluster; multiple nodes require shared storage.
Each index task searches one shard, while its native index I/O settings still
apply. Raw-scan parallelism is limited by the number of planned read splits,
-controlled by the table's `source.split.target-size` option. Single-vector
-refinement also uses planned read splits; workers return scores for merging.
+controlled by the table's `source.split.target-size` option. Refinement also
+uses planned read splits; workers return scores for merging.
Final row lookup runs on the driver. Candidate traffic grows
with the number of index shards and the configured refinement budget, so Ray
execution is most useful when shard search or raw scanning outweighs scheduling
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index 99a2a4993d..c3aa7b1bd2 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -662,8 +662,8 @@ class BatchVectorQuery(_PreFilterQuery):
Ray execution supports data-evolution tables. ``concurrency`` bounds
in-flight tasks (defaults to 4); ``ray_remote_args`` configures their
- resources and retries. Batch refinement and shared result lookup run
- on the driver, using the same snapshot as all workers.
+ resources and retries. Shared result lookup runs on the driver,
+ using the same snapshot as all workers, including refinement tasks.
"""
if execution == "local":
if concurrency is not None or ray_remote_args is not None:
diff --git a/paimon-python/pypaimon/ray/batch_vector_search.py
b/paimon-python/pypaimon/ray/batch_vector_search.py
index 9cd3939248..b7b4303a6f 100644
--- a/paimon-python/pypaimon/ray/batch_vector_search.py
+++ b/paimon-python/pypaimon/ray/batch_vector_search.py
@@ -76,6 +76,26 @@ class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl):
_offer_score(heap, self._limit, row_id, score)
return [_scored_result(heap) for heap in heaps]
+ def _stream_rerank_candidates(self, candidates, union_candidates,
query_vectors,
+ index_type, snapshot):
+ # Candidates have already been selected globally, separately per query.
+ queries_by_row = {}
+ for query_index, result in enumerate(candidates):
+ for row_id in result.results():
+ queries_by_row.setdefault(row_id, []).append(query_index)
+ table_read, splits = self._plan_raw_read(
+ union_candidates.to_range_list(), include_filter=False,
snapshot=snapshot)
+ context = (table_read, self._vector_column, query_vectors, self._limit,
+ self._search_metric(index_type), queries_by_row)
+ heaps = [[] for _ in query_vectors]
+ with closing(_map_tasks(
+ _search_batch_refine_split, context, splits,
self._concurrency, self._remote_args)) as tasks:
+ for _, results in tasks:
+ for heap, scores in zip(heaps, results):
+ for row_id, score in scores.items():
+ _offer_score(heap, self._limit, row_id, score)
+ return [_scored_result(heap) for heap in heaps]
+
def _search_batch_index_split(context, item):
table, column, queries, limit, options = context
@@ -108,6 +128,14 @@ def _search_batch_raw_split(context, split):
return [_scores(result) for result in
scorer._score_raw_batch_queries(batch_reader, metric)]
+def _search_batch_refine_split(context, split):
+ table_read, column, queries, limit, metric, queries_by_row = context
+ scorer = BatchVectorSearchReadImpl(table_read.table, limit, column,
queries)
+ # Read the candidate union once, scoring each row only for its own queries.
+ return [_scores(result) for result in scorer._score_refine_splits(
+ table_read, [split], queries_by_row, queries, metric)]
+
+
def _scores(result):
if result is None:
return {}
diff --git a/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py
b/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py
new file mode 100644
index 0000000000..f449eb3aad
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py
@@ -0,0 +1,145 @@
+# 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 types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pyarrow as pa
+import pytest
+
+pytest.importorskip("ray")
+
+from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
+from pypaimon.ray import batch_vector_search as search_module
+from pypaimon.read.table_read import TableRead
+from pypaimon.tests import ray_vector_search_test as fixtures
+from pypaimon.tests.ray_vector_search_test import VECTORS, add_rows,
build_index
+
+
+ray_cluster = fixtures.ray_cluster
+table = fixtures.table
+QUERIES = [[1., 1.], [-1., 2.], [1., 1.]]
+
+
[email protected]("metric", ["l2", "cosine", "inner_product"])
[email protected]("concurrency", [1, 2])
+def test_batch_refinement_reads_vectors_only_on_workers(table, ray_cluster,
metric, concurrency):
+ add_rows(table, VECTORS[:3])
+ add_rows(table, VECTORS[3:], 3)
+ build_index(table, metric)
+ query = table.search_vectors(QUERIES, options={"refine_factor":
"2"}).select(["id"]).limit(3)
+ expected = query.to_list()
+ original_read = TableRead._new_arrow_batch_reader
+ original_map = search_module._map_tasks
+ dispatched = []
+
+ def no_driver_vectors(read, *args, **kwargs):
+ assert "embedding" not in [field.name for field in read.read_type]
+ return original_read(read, *args, **kwargs)
+
+ def record(worker, context, items, *args):
+ if worker is search_module._search_batch_refine_split:
+ dispatched.append(len(items))
+ return original_map(worker, context, items, *args)
+
+ with patch.object(TableRead, "_new_arrow_batch_reader",
no_driver_vectors), \
+ patch.object(search_module, "_map_tasks", record):
+ actual = query.to_arrow(execution="ray", concurrency=concurrency)
+ assert [result.to_pylist() for result in actual] == expected
+ assert dispatched == [2]
+
+
[email protected]("fail_read", [False, True])
+def test_worker_shares_stream_and_keeps_per_query_membership(table, fail_read):
+ queries = [[1., 1.], [10., 1.], [1., 1.]]
+ query = table.search_vectors(queries)
+ reader =
query._batch_vector_search_builder(query).new_batch_vector_search_read()
+ schema = pa.schema([("_ROW_ID", pa.int64()), ("embedding",
pa.list_(pa.float32(), 2))])
+ closed = []
+
+ def batches():
+ try:
+ yield pa.record_batch([[0, 1], [[1., 1.], [10., 1.]]],
schema=schema)
+ if fail_read:
+ raise RuntimeError("injected refinement read failure")
+ yield pa.record_batch([[2], [None]], schema=schema)
+ finally:
+ closed.append(True)
+
+ source = batches()
+ arrow = pa.RecordBatchReader.from_batches(schema, source)
+ read = SimpleNamespace(table=reader._table,
_new_arrow_batch_reader=Mock(return_value=(arrow, source)))
+ # The closest row for query 1 belongs only to queries 0/2, and vice versa.
+ context = (read, reader._vector_column, queries, 1, "l2", {0: [1], 1: [0,
2], 2: [0, 1, 2]})
+ if fail_read:
+ with pytest.raises(RuntimeError, match="injected refinement read
failure"):
+ search_module._search_batch_refine_split(context, "split")
+ else:
+ results = search_module._search_batch_refine_split(context, "split")
+ assert [list(result) for result in results] == [[1], [0], [1]]
+ read._new_arrow_batch_reader.assert_called_once_with(["split"])
+ assert closed == [True]
+ assert source.gi_frame is None
+
+
+def test_empty_candidates_do_not_plan_or_dispatch(table):
+ query = table.search_vectors(QUERIES, options={"refine_factor": "2"})
+ reader =
query._batch_vector_search_builder(query).new_batch_vector_search_read()
+ distributed = search_module._RayBatchVectorSearchRead(reader, 2, {})
+ with patch.object(distributed, "_plan_raw_read",
side_effect=AssertionError("empty read")), \
+ patch.object(search_module, "_map_tasks",
side_effect=AssertionError("empty task")):
+ results = distributed._maybe_rerank_indexed_results(
+ [DictBasedScoredIndexResult({}) for _ in QUERIES], "ivf-flat",
QUERIES)
+ assert all(result.results().is_empty() for result in results)
+
+
+def test_commit_between_selection_and_refinement_keeps_snapshot(table,
ray_cluster):
+ add_rows(table, VECTORS)
+ build_index(table)
+ table.raw_table =
table.raw_table.copy({"global-index.column-update-action":
"DROP_PARTITION_INDEX"})
+ query = table.search_vectors(QUERIES, options={"refine_factor":
"2"}).select(["id"]).limit(1)
+ expected = query.to_list()
+ original = search_module._map_tasks
+ changed = []
+
+ def commit(worker, *args):
+ if worker is search_module._search_batch_refine_split and not changed:
+ table.update("id = 0", {"embedding": [100., 1.]})
+ changed.append(True)
+ return original(worker, *args)
+
+ with patch.object(search_module, "_map_tasks", commit):
+ assert [result.to_pylist() for result in
query.to_arrow(execution="ray")] == expected
+ assert changed == [True]
+ assert query.to_list() != expected
+ assert [result.to_pylist() for result in query.to_arrow(execution="ray")]
== query.to_list()
+
+
[email protected]("selector", ["snapshot", "tag"])
+def test_refinement_keeps_historical_deletions(table, ray_cluster, selector):
+ add_rows(table, VECTORS)
+ build_index(table)
+ table.delete("id = 0")
+ snapshot = table.raw_table.snapshot_manager().get_latest_snapshot()
+ table.raw_table.create_tag("saved", snapshot_id=snapshot.id)
+ options = {"snapshot_id": snapshot.id} if selector == "snapshot" else
{"tag_name": "saved"}
+ query = table.search_vectors(QUERIES, options={"refine_factor": "2"},
**options).select(["id"]).limit(3)
+ expected = query.to_list()
+ table.delete("id = 1")
+ if selector == "tag":
+
table.raw_table.file_io.delete(table.raw_table.snapshot_manager().get_snapshot_path(snapshot.id))
+ assert [result.to_pylist() for result in query.to_arrow(execution="ray")]
== expected
diff --git a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
index 5c401cc753..a28c82f90b 100644
--- a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
+++ b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
@@ -245,7 +245,11 @@ def test_global_candidates_and_duplicate_precedence(table):
index_file = IndexFileMeta("ivf-flat", "fake", 1, 2)
splits = [IndexVectorSearchSplit(0, 1, [index_file]),
IndexVectorSearchSplit(0, 3, [index_file])]
- def completed(*args):
+ def completed(worker, context, items, *args):
+ if worker is search_module._search_batch_refine_split:
+ for ordinal, split in enumerate(items):
+ yield ordinal, worker(context, split)
+ return
# First query must exclude exact nearest row 3; second query has its
own candidates.
# Duplicate row 0 must retain the earlier split's score for the second
query.
assert args[-1] is True
