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 bb47db10f6 [python] Execute Ray hybrid routes within a shared
concurrency budget (#10585)
bb47db10f6 is described below
commit bb47db10f651d59067144a3eebd5a7ca82c97679
Author: chaoyang <[email protected]>
AuthorDate: Fri Oct 16 13:01:03 2026 +0800
[python] Execute Ray hybrid routes within a shared concurrency budget
(#10585)
---
paimon-python/pypaimon/ray/full_text_search.py | 42 +++++--
paimon-python/pypaimon/ray/vector_search.py | 26 +++-
.../pypaimon/tests/ray_hybrid_parallel_test.py | 132 +++++++++++++++++++++
3 files changed, 192 insertions(+), 8 deletions(-)
diff --git a/paimon-python/pypaimon/ray/full_text_search.py
b/paimon-python/pypaimon/ray/full_text_search.py
index f4028202ac..a9d54e9e85 100644
--- a/paimon-python/pypaimon/ray/full_text_search.py
+++ b/paimon-python/pypaimon/ray/full_text_search.py
@@ -17,7 +17,9 @@
"""Ray full-text shard execution, preserving local corpus and merge
semantics."""
+from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import closing
+from threading import Event
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
from pypaimon.ray.vector_search import _execution_options, _map_tasks,
_require_ray, _scores
@@ -93,14 +95,40 @@ def _search_raw(reader, item):
def _execute_hybrid_search(builder, *, concurrency=None, ray_remote_args=None):
- from pypaimon.ray.vector_search import _execute_vector_search
+ from pypaimon.ray.vector_search import _execute_vector_search,
_task_cancellation
concurrency, remote_args = _execution_options(concurrency, ray_remote_args)
- route_results = []
- # Route builders share the resolved snapshot. Running routes successively
- # makes concurrency a query-wide bound, rather than multiplying it per
route.
- for route in builder.route_builders():
+ # Resolve the shared snapshot before any route starts planning on a thread.
+ routes = builder.route_builders()
+ workers = min(concurrency, len(routes))
+ route_concurrency = concurrency // max(1, workers)
+
+ def execute_route(route):
execute = _execute_vector_search if route.route.is_vector() else
_execute_full_text_search
- result = execute(route.search_builder, concurrency=concurrency,
ray_remote_args=remote_args)
- route_results.append(builder.to_route_result(route, result))
+ return execute(route.search_builder, concurrency=route_concurrency,
ray_remote_args=remote_args)
+
+ if workers <= 1:
+ return builder.rank([builder.to_route_result(route,
execute_route(route)) for route in routes])
+
+ cancelled = Event()
+
+ def run(route):
+ with _task_cancellation(cancelled):
+ return execute_route(route)
+
+ route_results = [None] * len(routes)
+ # Each concurrent route gets a share of the existing search-task budget.
+ # Completion order must not change route ordering supplied to the ranker.
+ with ThreadPoolExecutor(max_workers=workers) as executor:
+ futures = {}
+ try:
+ for ordinal, route in enumerate(routes):
+ futures[executor.submit(run, route)] = ordinal
+ for future in as_completed(futures):
+ ordinal = futures[future]
+ route_results[ordinal] =
builder.to_route_result(routes[ordinal], future.result())
+ finally:
+ cancelled.set()
+ for future in futures:
+ future.cancel()
return builder.rank(route_results)
diff --git a/paimon-python/pypaimon/ray/vector_search.py
b/paimon-python/pypaimon/ray/vector_search.py
index 7752da972e..4e7239bfd6 100644
--- a/paimon-python/pypaimon/ray/vector_search.py
+++ b/paimon-python/pypaimon/ray/vector_search.py
@@ -18,8 +18,10 @@
"""Ray execution of snapshot-pinned, single-vector queries."""
from collections import Counter
+from concurrent.futures import CancelledError
from contextlib import closing, contextmanager
import math
+import threading
from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
@@ -216,6 +218,22 @@ def _scores(result):
return scores
+_task_state = threading.local()
+
+
+@contextmanager
+def _task_cancellation(event):
+ """Keep hybrid cancellation on driver threads, out of serialized worker
state."""
+ previous = getattr(_task_state, "cancelled", None)
+ _task_state.cancelled = event
+ try:
+ if event.is_set():
+ raise CancelledError()
+ yield
+ finally:
+ _task_state.cancelled = previous
+
+
def _map_tasks(worker, context, items, concurrency, remote_args,
ordered=False, prepare_item=None):
"""Bound in-flight work and, optionally, completed results awaiting plan
order."""
import ray
@@ -229,8 +247,11 @@ def _map_tasks(worker, context, items, concurrency,
remote_args, ordered=False,
pending = {}
buffered = {}
next_ordinal = 0
+ cancelled = getattr(_task_state, "cancelled", None)
try:
while True:
+ if cancelled is not None and cancelled.is_set():
+ raise CancelledError()
while len(pending) + len(buffered) < concurrency:
item = next(remaining, None)
if item is None:
@@ -240,7 +261,10 @@ def _map_tasks(worker, context, items, concurrency,
remote_args, ordered=False,
pending[remote.remote(context_ref, *args)] = ordinal
if not pending:
break
- ready, _ = ray.wait(list(pending), num_returns=1)
+ ready, _ = ray.wait(list(pending), num_returns=1,
+ **({"timeout": 0.1} if cancelled is not None
else {}))
+ if not ready:
+ continue
ref = ready[0]
result = ray.get(ref)
ordinal = pending.pop(ref)
diff --git a/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py
b/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py
new file mode 100644
index 0000000000..755e6a2e53
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py
@@ -0,0 +1,132 @@
+# 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 concurrent.futures import CancelledError
+from threading import Barrier, Event, Lock
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+ray = pytest.importorskip("ray")
+
+from pypaimon.ray import full_text_search, vector_search
+
+
+def builder(count):
+ routes = [SimpleNamespace(route=SimpleNamespace(is_vector=lambda: False),
search_builder=i)
+ for i in range(count)]
+ return SimpleNamespace(route_builders=lambda: routes,
+ to_route_result=lambda route, result:
(route.search_builder, result),
+ rank=lambda results: results)
+
+
[email protected]("concurrency, count", [(2, 2), (3, 2), (4, 2), (2,
4)])
+def test_routes_overlap_within_budget_and_preserve_order(concurrency, count):
+ barrier = Barrier(2, timeout=10)
+ lock = Lock()
+ second_completed = Event()
+ active = 0
+ peak = 0
+ completed = []
+
+ def execute(ordinal, *, concurrency, ray_remote_args):
+ nonlocal active, peak
+ assert ray_remote_args == {"num_cpus": 0.5}
+ with lock:
+ active += concurrency
+ peak = max(peak, active)
+ barrier.wait()
+ if ordinal == 0:
+ assert second_completed.wait(10)
+ with lock:
+ completed.append(ordinal)
+ active -= concurrency
+ if ordinal == 1:
+ second_completed.set()
+ return ordinal
+
+ with patch.object(full_text_search, "_execute_full_text_search", execute):
+ result = full_text_search._execute_hybrid_search(
+ builder(count), concurrency=concurrency,
ray_remote_args={"num_cpus": 0.5})
+ assert result == [(i, i) for i in range(count)]
+ assert completed.index(1) < completed.index(0)
+ assert 2 <= peak <= concurrency
+
+
[email protected]("concurrency, count", [(1, 3), (4, 1)])
+def test_serial_routes_keep_full_budget(concurrency, count):
+ calls = []
+
+ def execute(ordinal, **kwargs):
+ assert kwargs["concurrency"] == concurrency
+ calls.append(ordinal)
+ return ordinal
+
+ with patch.object(full_text_search, "_execute_full_text_search", execute):
+ assert full_text_search._execute_hybrid_search(builder(count),
concurrency=concurrency) == [
+ (i, i) for i in range(count)]
+ assert calls == list(range(count))
+
+
+def test_route_failure_cancels_sibling_ray_tasks():
+ waiting = Event()
+ submitted, cancelled = [], []
+
+ def submit(context, item):
+ submitted.append(item)
+ return item
+
+ remote = SimpleNamespace(remote=submit)
+ remote.options = lambda **kwargs: remote
+
+ def wait(refs, **kwargs):
+ assert kwargs["timeout"] == 0.1
+ waiting.set()
+ # The task never completes: cancellation must interrupt polling.
+ Event().wait(0.01)
+ return [], refs
+
+ def execute(ordinal, **kwargs):
+ if ordinal == 0:
+ assert waiting.wait(10)
+ raise ValueError("route failed")
+ return list(vector_search._map_tasks(None, None, [0, 1, 2], 1, {}))
+
+ with patch.object(full_text_search, "_execute_full_text_search", execute),
\
+ patch.object(ray, "put", return_value=None), \
+ patch.object(ray, "remote", return_value=remote), \
+ patch.object(ray, "wait", side_effect=wait), \
+ patch.object(ray, "cancel", side_effect=cancelled.append), \
+ pytest.raises(ValueError, match="route failed"):
+ full_text_search._execute_hybrid_search(builder(2), concurrency=2)
+ assert submitted == [0]
+ assert cancelled == [0]
+
+
+def test_cancellation_scope_restores_state_and_skips_pending_routes():
+ cancelled = Event()
+ cancelled.set()
+ with pytest.raises(CancelledError),
vector_search._task_cancellation(cancelled):
+ pytest.fail("cancelled route started")
+ assert getattr(vector_search._task_state, "cancelled", None) is None
+ outer, inner = Event(), Event()
+ with vector_search._task_cancellation(outer):
+ with vector_search._task_cancellation(inner):
+ assert vector_search._task_state.cancelled is inner
+ assert vector_search._task_state.cancelled is outer
+ assert vector_search._task_state.cancelled is None
