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 eb43b419ce88e9f491fa1289fb017e2077c4dedd Author: chaoyang <[email protected]> AuthorDate: Tue Sep 15 11:15:17 2026 +0800 [python] Execute hybrid search routes concurrently (#9797) --- .../pypaimon/table/source/hybrid_search_builder.py | 35 ++++- .../pypaimon/tests/hybrid_search_execution_test.py | 142 +++++++++++++++++++++ 2 files changed, 173 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/table/source/hybrid_search_builder.py b/paimon-python/pypaimon/table/source/hybrid_search_builder.py index 196acbb8aa..e05c248a63 100644 --- a/paimon-python/pypaimon/table/source/hybrid_search_builder.py +++ b/paimon-python/pypaimon/table/source/hybrid_search_builder.py @@ -20,6 +20,7 @@ import heapq import math from abc import ABC, abstractmethod +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from dataclasses import dataclass, field from typing import Dict, List, Optional @@ -34,6 +35,7 @@ RRF_RANKER = "rrf" WEIGHTED_SCORE_RANKER = "weighted_score" MRR_RANKER = "mrr" _RRF_K = 60.0 +_MAX_ROUTE_WORKERS = 4 def _check_full_text_options(options: Dict[str, str]): @@ -230,11 +232,36 @@ class HybridSearchBuilder(ABC): def execute_local(self) -> ScoredGlobalIndexResult: """Execute hybrid index search locally.""" route_builders = self.route_builders() - route_results = [] - for route_builder in route_builders: - route_results.append( + if len(route_builders) <= 1: + return self.rank([ self.to_route_result( - route_builder, route_builder.execute_local())) + route_builder, route_builder.execute_local()) + for route_builder in route_builders + ]) + + workers = min(len(route_builders), _MAX_ROUTE_WORKERS) + with ThreadPoolExecutor( + max_workers=workers, + thread_name_prefix="paimon-hybrid-search") as executor: + futures = [ + executor.submit(route_builder.execute_local) + for route_builder in route_builders + ] + done, pending = wait(futures, return_when=FIRST_EXCEPTION) + failed = next( + (future for future in futures + if future in done and future.exception() is not None), + None, + ) + if failed is not None: + for future in pending: + future.cancel() + failed.result() + + route_results = [ + self.to_route_result(route_builder, future.result()) + for route_builder, future in zip(route_builders, futures) + ] return self.rank(route_results) diff --git a/paimon-python/pypaimon/tests/hybrid_search_execution_test.py b/paimon-python/pypaimon/tests/hybrid_search_execution_test.py new file mode 100644 index 0000000000..22a195be8f --- /dev/null +++ b/paimon-python/pypaimon/tests/hybrid_search_execution_test.py @@ -0,0 +1,142 @@ +# 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. + +"""Tests for local hybrid-search route execution.""" + +import threading +import time +import unittest +from concurrent.futures import ThreadPoolExecutor +from unittest import mock + +from pypaimon.table.source.hybrid_search_builder import ( + HybridSearchBuilderImpl, + HybridSearchRouteBuilder, +) + + +class _CallableSearchBuilder: + + def __init__(self, execute): + self._execute = execute + + def execute_local(self): + return self._execute() + + +def _execution_builder(executions): + route_builders = [ + HybridSearchRouteBuilder( + "route-%d" % index, _CallableSearchBuilder(execute)) + for index, execute in enumerate(executions) + ] + builder = HybridSearchBuilderImpl(table=None) + builder.route_builders = lambda: route_builders + builder.to_route_result = ( + lambda route_builder, result: (route_builder.route, result)) + builder.rank = lambda route_results: route_results + return builder + + +class HybridSearchExecutionTest(unittest.TestCase): + + def test_executes_routes_concurrently_and_preserves_order(self): + started = threading.Barrier(2) + + def execute(index, delay): + started.wait(timeout=2.0) + time.sleep(delay) + return "result-%d" % index + + builder = _execution_builder([ + lambda: execute(0, 0.03), + lambda: execute(1, 0.0), + ]) + + self.assertEqual( + [("route-0", "result-0"), ("route-1", "result-1")], + builder.execute_local(), + ) + + def test_single_route_avoids_executor(self): + builder = _execution_builder([lambda: "result"]) + + with mock.patch( + "pypaimon.table.source.hybrid_search_builder." + "ThreadPoolExecutor") as executor: + self.assertEqual( + [("route-0", "result")], builder.execute_local()) + + executor.assert_not_called() + + def test_caps_route_workers(self): + builder = _execution_builder([ + lambda index=index: index for index in range(6) + ]) + worker_counts = [] + + def new_executor(*args, **kwargs): + worker_counts.append(kwargs["max_workers"]) + return ThreadPoolExecutor(*args, **kwargs) + + with mock.patch( + "pypaimon.table.source.hybrid_search_builder." + "ThreadPoolExecutor", side_effect=new_executor): + builder.execute_local() + + self.assertEqual([4], worker_counts) + + def test_propagates_failure_after_started_routes_finish(self): + started = threading.Barrier(2) + failed = threading.Event() + release = threading.Event() + finished = threading.Event() + outcome = {} + + def fail(): + started.wait(timeout=2.0) + failed.set() + raise RuntimeError("route failed") + + def block(): + started.wait(timeout=2.0) + release.wait(timeout=2.0) + finished.set() + return "result" + + builder = _execution_builder([fail, block]) + + def execute(): + try: + builder.execute_local() + except BaseException as error: + outcome["error"] = error + + caller = threading.Thread(target=execute) + caller.start() + self.assertTrue(failed.wait(timeout=2.0)) + self.assertTrue(caller.is_alive()) + release.set() + caller.join(timeout=2.0) + + self.assertFalse(caller.is_alive()) + self.assertTrue(finished.is_set()) + self.assertIsInstance(outcome.get("error"), RuntimeError) + self.assertEqual("route failed", str(outcome["error"])) + + +if __name__ == "__main__": + unittest.main()
