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 c2b6a0435ce36befb95fce4397db8d9ef9a58dea Author: chaoyang <[email protected]> AuthorDate: Fri Sep 18 19:38:47 2026 +0800 [python] Distribute vector search refinement with Ray (#9951) --- docs/docs/pypaimon/multimodal-search.md | 12 +- paimon-python/pypaimon/multimodal/query.py | 2 +- paimon-python/pypaimon/ray/vector_search.py | 14 +-- .../pypaimon/tests/ray_vector_refinement_test.py | 134 +++++++++++++++++++++ .../pypaimon/tests/ray_vector_search_test.py | 6 +- 5 files changed, 152 insertions(+), 16 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 9ec110a42c..73f0280b0e 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -143,8 +143,11 @@ Ray execution fails because NaN cannot be ranked consistently across tasks. 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. The driver retains the local search -rules for global candidate selection, refinement, filtering, and result lookup. +They return candidate row IDs and scores. For single-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 +filtering and final row lookup. In particular, `pre_filter` filters candidates before ranking; `where()` filters the selected rows and can return fewer than the requested number of results. @@ -158,8 +161,9 @@ 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. Refinement and final -row lookup run on the driver. Candidate traffic grows +controlled by the table's `source.split.target-size` option. Single-vector +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 and transfer costs. Small queries can be faster locally. diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 60c3452dc5..da6d712713 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -537,7 +537,7 @@ class VectorQuery(_PreFilterQuery): Ray execution supports data-evolution tables. ``concurrency`` bounds in-flight tasks (defaults to 4); ``ray_remote_args`` configures their - resources and retries. Refinement and result lookup run on the driver. + resources and retries. Result lookup runs on the driver. """ if execution == "local": if concurrency is not None or ray_remote_args is not None: diff --git a/paimon-python/pypaimon/ray/vector_search.py b/paimon-python/pypaimon/ray/vector_search.py index 348df0f93c..dfd51f1f35 100644 --- a/paimon-python/pypaimon/ray/vector_search.py +++ b/paimon-python/pypaimon/ray/vector_search.py @@ -67,10 +67,6 @@ class _RayVectorSearchRead(DataEvolutionVectorRead): self._concurrency = concurrency self._remote_args = remote_args - def _score_raw_batches(self, batches, query_vector, metric, score_candidates=None, reject_nan=True): - # Refinement also runs through this scorer, before its local top-k. - return super()._score_raw_batches(batches, query_vector, metric, score_candidates, reject_nan=True) - def _search_index_splits(self, splits, query, search_limit, pre_filters, batch=False): # Filters are planned once on the driver, at the query snapshot. Each # worker receives only its own include-row bitmap, not a table-wide set. @@ -91,13 +87,11 @@ class _RayVectorSearchRead(DataEvolutionVectorRead): def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, index_type=None, include_filter=True, score_candidates=None, snapshot=None): - if score_candidates is not None: - # Refinement uses the globally selected indexed candidates. - return super()._read_raw_search( - raw_row_ranges, pre_filter, query_vector, index_type, - include_filter, score_candidates, snapshot) - ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) + if score_candidates is not None: + # Refine only the globally selected candidates. Intersect before + # planning so workers neither read nor score other rows. + ranges = _filtered_raw_row_ranges(ranges, score_candidates.to_range_list()) if not ranges: return DictBasedScoredIndexResult({}) table_read, splits = self._plan_raw_read(ranges, include_filter, snapshot) diff --git a/paimon-python/pypaimon/tests/ray_vector_refinement_test.py b/paimon-python/pypaimon/tests/ray_vector_refinement_test.py new file mode 100644 index 0000000000..386dbaafbc --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_vector_refinement_test.py @@ -0,0 +1,134 @@ +# 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 pytest + +pytest.importorskip("ray") + +from pypaimon.ray import 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, ids +from pypaimon.utils.range import Range +from pypaimon.utils.roaring_bitmap import RoaringBitmap64 + + +ray_cluster = fixtures.ray_cluster +table = fixtures.table + + [email protected]("metric", ["l2", "cosine", "inner_product"]) [email protected]("concurrency", [1, 2]) +def test_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([1., 1.], options={"refine_factor": "2"}).select(["id"]).limit(3) + expected = ids(query.to_arrow()) + 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_raw_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): + assert ids(query.to_arrow(execution="ray", concurrency=concurrency)) == expected + assert dispatched == [2] + + +def test_candidate_intersection_excludes_nan_and_streams_real_splits(table, ray_cluster): + add_rows(table, [[0., 1.], [float("nan"), 1.]]) + add_rows(table, [[float("nan"), 1.], [1., 1.]], 2) + query = table.search([1., 1.], options={"metric": "l2"}).limit(1)._for_execution() + reader = query._vector_search_builder(query).new_vector_search_read() + distributed = search_module._RayVectorSearchRead(reader, 2, {}) + candidates = RoaringBitmap64() + candidates.add(0) + candidates.add(3) + result = distributed._read_raw_search( + [Range(0, 3)], None, [1., 1.], score_candidates=candidates, + snapshot=query._table._read_snapshot) + assert list(result.results()) == [3] + assert result.score_getter()(3) == 1.0 + candidates.add(1) + with pytest.raises(ValueError, match="cannot rank NaN"): + distributed._read_raw_search( + [Range(0, 3)], None, [1., 1.], score_candidates=candidates, + snapshot=query._table._read_snapshot) + + +def test_empty_refinement_does_not_plan_or_dispatch(table): + query = table.search([1., 1.]).limit(1)._for_execution() + reader = query._vector_search_builder(query).new_vector_search_read() + distributed = search_module._RayVectorSearchRead(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")): + result = distributed._read_raw_search( + [Range(0, 100)], None, [1., 1.], score_candidates=RoaringBitmap64()) + assert result.results().is_empty() + + [email protected]("selector", ["snapshot", "tag"]) +def test_refinement_keeps_historical_filters_and_deletions(table, ray_cluster, selector): + add_rows(table, VECTORS) + build_index(table) + table.update("id = 2", {"category": "no"}) + table.delete("id = 0") + snapshot = table.raw_table.snapshot_manager().get_latest_snapshot() + table.raw_table.create_tag("saved", snapshot_id=snapshot.id) + table.delete("id = 4") + if selector == "tag": + table.raw_table.file_io.delete(table.raw_table.snapshot_manager().get_snapshot_path(snapshot.id)) + options = {"tag_name": "saved"} + else: + options = {"snapshot_id": snapshot.id} + query = table.search([1., 1.], pre_filter="category = 'yes'", + options={"refine_factor": "2"}, **options).select(["id"]).limit(3) + assert ids(query.to_arrow()) == [4] + assert ids(query.to_arrow(execution="ray")) == [4] + + +def test_commit_between_candidate_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([1., 1.], options={"refine_factor": "2"}).select(["id"]).limit(1) + expected = ids(query.to_arrow()) + original = search_module._map_tasks + changed = [] + + def commit_before_refine(worker, *args): + if worker is search_module._search_raw_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_before_refine): + assert ids(query.to_arrow(execution="ray")) == expected + assert changed == [True] + assert ids(query.to_arrow()) != expected + assert ids(query.to_arrow(execution="ray")) == ids(query.to_arrow()) diff --git a/paimon-python/pypaimon/tests/ray_vector_search_test.py b/paimon-python/pypaimon/tests/ray_vector_search_test.py index 06a7446863..6c251b039a 100644 --- a/paimon-python/pypaimon/tests/ray_vector_search_test.py +++ b/paimon-python/pypaimon/tests/ray_vector_search_test.py @@ -317,7 +317,11 @@ def test_refinement_uses_global_candidates_despite_reverse_completion_order(tabl index_file = IndexFileMeta("ivf-flat", "fake", 1, 2, None) splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(2, 3, [index_file])] - def completed(*args): + def completed(worker, context, items, *args): + if worker is search_module._search_raw_split: + for ordinal, split in enumerate(items): + yield ordinal, worker(context, split) + return # Row 3 is the true nearest, but is outside the GLOBAL approximate top-2. # Refining independently per shard would incorrectly bring it back in. yield 1, ("l2", {2: 8., 3: 7.})
