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 54c4e66d0b [python] Stream vector index candidates during global 
merging (#9989)
54c4e66d0b is described below

commit 54c4e66d0b0f7347af27e3f9b3a7762b8bae0b03
Author: chaoyang <[email protected]>
AuthorDate: Sun Sep 20 13:44:23 2026 +0800

    [python] Stream vector index candidates during global merging (#9989)
---
 docs/docs/pypaimon/multimodal-search.md            |  6 ++
 paimon-python/pypaimon/ray/batch_vector_search.py  |  9 ++-
 paimon-python/pypaimon/ray/vector_search.py        | 28 +++++---
 .../pypaimon/table/source/vector_search_read.py    | 64 ++++++++++--------
 .../pypaimon/tests/ray_batch_vector_search_test.py |  7 +-
 .../pypaimon/tests/ray_vector_search_test.py       | 76 ++++++++++++++++++----
 .../tests/vector_index_parallel_search_test.py     | 75 ++++++++++++++++++---
 7 files changed, 200 insertions(+), 65 deletions(-)

diff --git a/docs/docs/pypaimon/multimodal-search.md 
b/docs/docs/pypaimon/multimodal-search.md
index f7ed76cf01..63cd120220 100644
--- a/docs/docs/pypaimon/multimodal-search.md
+++ b/docs/docs/pypaimon/multimodal-search.md
@@ -187,6 +187,12 @@ 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.
 
+Local and Ray index searches merge shard results incrementally in plan order.
+The concurrency limit bounds running tasks plus completed results waiting for
+earlier shards, so a slow shard can delay further submissions. Global top-k
+selection still uses all merged candidate scores, whose storage grows with the
+number of unique candidate rows per query.
+
 ## Search Hybrid
 
 Use `search_hybrid` to combine vector and full-text routes, then rerank the
diff --git a/paimon-python/pypaimon/ray/batch_vector_search.py 
b/paimon-python/pypaimon/ray/batch_vector_search.py
index 6cf1f6593d..9cd3939248 100644
--- a/paimon-python/pypaimon/ray/batch_vector_search.py
+++ b/paimon-python/pypaimon/ray/batch_vector_search.py
@@ -51,15 +51,14 @@ class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl):
         items = [(split, None if not pre_filters or pre_filters[i] is None
                   else pre_filters[i].serialize()) for i, split in 
enumerate(splits)]
         context = (self._table, self._vector_column, query, search_limit, 
self._options)
-        results = [None] * len(splits)
         with closing(_map_tasks(
-                _search_batch_index_split, context, items, self._concurrency, 
self._remote_args)) as tasks:
-            for ordinal, (metric, scores) in tasks:
+                _search_batch_index_split, context, items, self._concurrency, 
self._remote_args, True)) as tasks:
+            for _, (metric, scores) in tasks:
                 if metric is not None:
                     self._set_index_metric(metric)
                 # Retain plan order for duplicate IDs and global per-query 
selection.
-                results[ordinal] = [DictBasedScoredIndexResult(values) for 
values in scores]
-        return results
+                yield [DictBasedScoredIndexResult(values) for values in scores]
+                del scores
 
     def _read_raw_batch_search(self, raw_row_ranges, pre_filter, 
index_type=None, snapshot=None):
         heaps = [[] for _ in self._query_vectors]
diff --git a/paimon-python/pypaimon/ray/vector_search.py 
b/paimon-python/pypaimon/ray/vector_search.py
index 13e0fb8695..9e99425827 100644
--- a/paimon-python/pypaimon/ray/vector_search.py
+++ b/paimon-python/pypaimon/ray/vector_search.py
@@ -80,16 +80,15 @@ class _RayVectorSearchRead(DataEvolutionVectorRead):
         items = [(split, None if not pre_filters or pre_filters[i] is None
                   else pre_filters[i].serialize()) for i, split in 
enumerate(splits)]
         context = (self._table, self._vector_column, query, search_limit, 
self._options)
-        results = [None] * len(splits)
         with closing(_map_tasks(
-                _search_index_split, context, items, self._concurrency, 
self._remote_args)) as tasks:
-            for i, (metric, scores) in tasks:
+                _search_index_split, context, items, self._concurrency, 
self._remote_args, True)) as tasks:
+            for _, (metric, scores) in tasks:
                 if metric is not None:
                     self._set_index_metric(metric)
                 # Keep plan order, including duplicate-row precedence, 
regardless
                 # of worker completion order. Do not refine or truncate per 
worker.
-                results[i] = DictBasedScoredIndexResult(scores)
-        return results
+                yield DictBasedScoredIndexResult(scores)
+                del scores
 
     def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector,
                          index_type=None, include_filter=True,
@@ -149,8 +148,8 @@ def _scores(result):
     return scores
 
 
-def _map_tasks(worker, context, items, concurrency, remote_args):
-    """Yield completed tasks with their plan ordinal, bounding in-flight 
work."""
+def _map_tasks(worker, context, items, concurrency, remote_args, 
ordered=False):
+    """Bound in-flight work and, optionally, completed results awaiting plan 
order."""
     import ray
 
     if not items:
@@ -160,9 +159,11 @@ def _map_tasks(worker, context, items, concurrency, 
remote_args):
     remote = ray.remote(worker).options(**remote_args)
     remaining = iter(enumerate(items))
     pending = {}
+    buffered = {}
+    next_ordinal = 0
     try:
         while True:
-            while len(pending) < concurrency:
+            while len(pending) + len(buffered) < concurrency:
                 item = next(remaining, None)
                 if item is None:
                     break
@@ -173,7 +174,16 @@ def _map_tasks(worker, context, items, concurrency, 
remote_args):
             ready, _ = ray.wait(list(pending), num_returns=1)
             ref = ready[0]
             result = ray.get(ref)
-            yield pending.pop(ref), result
+            ordinal = pending.pop(ref)
+            if ordered:
+                buffered[ordinal] = result
+                del result
+                while next_ordinal in buffered:
+                    yield next_ordinal, buffered.pop(next_ordinal)
+                    next_ordinal += 1
+            else:
+                yield ordinal, result
+                del result
     finally:
         for ref in pending:
             try:
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py 
b/paimon-python/pypaimon/table/source/vector_search_read.py
index 6a0fbf16e1..f48f1b10f4 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -18,7 +18,9 @@
 """Vector search read to read index files."""
 
 from abc import ABC, abstractmethod
+from collections import deque
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import closing
 from threading import Lock
 
 from pypaimon.common.options.core_options import CoreOptions
@@ -494,15 +496,24 @@ class AbstractVectorSearchReadImpl:
 
         workers = min(parallelism, len(splits))
         if workers <= 1:
-            return [search(i) for i in range(len(splits))]
+            for i in range(len(splits)):
+                yield search(i)
+            return
         with ThreadPoolExecutor(max_workers=workers) as executor:
-            futures = []
+            pending = deque()
             try:
-                for i in range(len(splits)):
-                    futures.append(executor.submit(search, i))
-                return [future.result() for future in futures]
+                for i in range(workers):
+                    pending.append(executor.submit(search, i))
+                remaining = iter(range(workers, len(splits)))
+                while pending:
+                    # Consume in plan order to retain duplicate-row precedence.
+                    # Bound completed results as well as running searches.
+                    yield pending.popleft().result()
+                    i = next(remaining, None)
+                    if i is not None:
+                        pending.append(executor.submit(search, i))
             finally:
-                for future in futures:
+                for future in pending:
                     future.cancel()
                 # Executor shutdown waits for started readers to close on 
failure.
 
@@ -652,16 +663,12 @@ class 
DataEvolutionVectorRead(AbstractVectorSearchReadImpl, VectorSearchRead):
         index_type = _vector_index_type(splits)
         search_limit = self._indexed_search_limit(index_type)
         pre_filters = self._pre_filters(splits, snapshot)
-        results = self._search_index_splits(
-            splits, query_vector, search_limit, pre_filters)
-
         merged_scores = {}
-        for split_result in results:
-            if split_result is not None:
-                score_getter = split_result.score_getter()
-                for row_id in split_result.results():
-                    if row_id not in merged_scores:
-                        merged_scores[row_id] = score_getter(row_id)
+        with closing(self._search_index_splits(
+                splits, query_vector, search_limit, pre_filters)) as results:
+            for split_result in results:
+                _merge_index_scores(merged_scores, split_result)
+                del split_result
 
         indexed = DictBasedScoredIndexResult(merged_scores).top_k(search_limit)
         return self._maybe_rerank_indexed_result(
@@ -692,20 +699,14 @@ class 
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
         index_type = _vector_index_type(index_splits)
         search_limit = self._indexed_search_limit(index_type)
         pre_filters = self._pre_filters(index_splits, snapshot)
-        results = self._search_index_splits(
-            index_splits, self._query_vectors, search_limit, pre_filters, 
batch=True)
-
         # Merge each query vector's indexed results across index splits.
         merged_scores = [{} for _ in range(n)]
-        for split_results in results:
-            for i in range(n):
-                split_result = split_results[i]
-                if split_result is None:
-                    continue
-                score_getter = split_result.score_getter()
-                for row_id in split_result.results():
-                    if row_id not in merged_scores[i]:
-                        merged_scores[i][row_id] = score_getter(row_id)
+        with closing(self._search_index_splits(
+                index_splits, self._query_vectors, search_limit, pre_filters, 
batch=True)) as results:
+            for split_results in results:
+                for i in range(n):
+                    _merge_index_scores(merged_scores[i], split_results[i])
+                del split_results
 
         indexed_results = [
             DictBasedScoredIndexResult(merged_scores[i]).top_k(search_limit)
@@ -777,6 +778,15 @@ class 
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
         return [_scored_result(heap) for heap in heaps]
 
 
+def _merge_index_scores(merged_scores, result):
+    if result is None:
+        return
+    score_getter = result.score_getter()
+    for row_id in result.results():
+        if row_id not in merged_scores:
+            merged_scores[row_id] = score_getter(row_id)
+
+
 def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list, 
options=None):
     """Create a global index reader for vector search."""
     from pypaimon.globalindex.lumina.lumina_vector_global_index_reader import (
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 78ab13c9b2..5c401cc753 100644
--- a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
+++ b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
@@ -237,7 +237,7 @@ def test_empty_snapshot_does_not_follow_first_commit(table, 
ray_cluster):
     assert [ids(result) for result in query.to_arrow(execution="ray")] == 
[[0]] * 3
 
 
-def 
test_global_candidates_and_duplicate_precedence_ignore_completion_order(table):
+def test_global_candidates_and_duplicate_precedence(table):
     add_rows(table, [[0., 1.], [10., 1.], [20., 1.], [1., 1.]])
     query = table.search_vectors([[1., 1.], [20., 1.]], 
options={"refine_factor": "2"}).limit(1)._for_execution()
     reader = 
query._batch_vector_search_builder(query).new_batch_vector_search_read()
@@ -248,8 +248,9 @@ def 
test_global_candidates_and_duplicate_precedence_ignore_completion_order(tabl
     def completed(*args):
         # 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.
-        yield 1, ("l2", [{2: 8., 3: 7.}, {0: 100., 2: 8., 3: 7.}])
+        assert args[-1] is True
         yield 0, ("l2", [{0: 10., 1: 9.}, {0: 1., 1: 9.}])
+        yield 1, ("l2", [{2: 8., 3: 7.}, {0: 100., 2: 8., 3: 7.}])
 
     with patch.object(search_module, "_map_tasks", completed):
         results = distributed._read_batch(splits, query._table._read_snapshot)
@@ -272,7 +273,7 @@ def test_metric_mismatch_in_empty_shard_closes_tasks(table):
 
     with patch.object(search_module, "_map_tasks", completed), \
             pytest.raises(ValueError, match="different metrics"):
-        distributed._search_index_splits([None, None], QUERIES, 1, [], 
batch=True)
+        list(distributed._search_index_splits([None, None], QUERIES, 1, [], 
batch=True))
     assert closed == [True]
 
 
diff --git a/paimon-python/pypaimon/tests/ray_vector_search_test.py 
b/paimon-python/pypaimon/tests/ray_vector_search_test.py
index 6c251b039a..986e766706 100644
--- a/paimon-python/pypaimon/tests/ray_vector_search_test.py
+++ b/paimon-python/pypaimon/tests/ray_vector_search_test.py
@@ -17,6 +17,7 @@
 
 import os
 from contextlib import closing
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import numpy as np
@@ -255,7 +256,8 @@ def 
test_warmed_read_context_and_distinct_worker_processes(table, ray_cluster):
     assert ids(query.to_arrow(execution="ray")) == [0, 1]
 
 
-def test_task_retry_keeps_original_snapshot(table, ray_cluster, tmp_path):
[email protected]("ordered", [False, True])
+def test_task_retry_keeps_original_snapshot(ordered, table, ray_cluster, 
tmp_path):
     add_rows(table, [[1., 1.]])
     pinned = table.search([1., 1.])._for_execution()._table
     add_rows(table, [[2., 1.]], 1)
@@ -272,10 +274,12 @@ def test_task_retry_keeps_original_snapshot(table, 
ray_cluster, tmp_path):
         raise RuntimeError("retry this read")
 
     assert list(search_module._map_tasks(
-        retry_once, pinned, [marker], 1, {"max_retries": 1, 
"retry_exceptions": True})) == [(0, (1, [0]))]
+        retry_once, pinned, [marker], 1,
+        {"max_retries": 1, "retry_exceptions": True}, ordered=ordered)) == 
[(0, (1, [0]))]
 
 
-def test_failure_cancels_outstanding_tasks(ray_cluster):
[email protected]("ordered", [False, True])
+def test_failure_cancels_outstanding_tasks(ordered, ray_cluster):
     def fail(context, item):
         if item == 0:
             raise ValueError("injected search failure")
@@ -284,11 +288,12 @@ def test_failure_cancels_outstanding_tasks(ray_cluster):
 
     with patch.object(ray, "cancel", wraps=ray.cancel) as cancel:
         with pytest.raises(ValueError, match="injected search failure"):
-            list(search_module._map_tasks(fail, None, [0, 1, 2], 2, {}))
+            list(search_module._map_tasks(fail, None, [0, 1, 2], 2, {}, 
ordered=ordered))
         assert cancel.call_count == 2
 
 
-def test_map_tasks_bounds_submissions_and_cancels_on_early_close(ray_cluster):
[email protected]("ordered", [False, True])
+def test_map_tasks_bounds_submissions_and_cancels_on_early_close(ordered, 
ray_cluster):
     def echo(context, item):
         return item
 
@@ -300,13 +305,13 @@ def 
test_map_tasks_bounds_submissions_and_cancels_on_early_close(ray_cluster):
         return wait(refs, **kwargs)
 
     with patch.object(ray, "wait", record_wait), patch.object(ray, "cancel", 
wraps=ray.cancel) as cancel:
-        with closing(search_module._map_tasks(echo, None, list(range(10)), 2, 
{})) as tasks:
+        with closing(search_module._map_tasks(echo, None, list(range(10)), 2, 
{}, ordered=ordered)) as tasks:
             next(tasks)
-        assert counts == [2]
-        assert cancel.call_count == 1
+        assert counts == [2] or (ordered and counts == [2, 1])
+        assert cancel.call_count == 2 - len(counts)
 
 
-def 
test_refinement_uses_global_candidates_despite_reverse_completion_order(table):
+def test_refinement_uses_global_candidates(table):
     from pypaimon.table.source.vector_search_split import 
IndexVectorSearchSplit
     from pypaimon.index.index_file_meta import IndexFileMeta
 
@@ -324,8 +329,9 @@ def 
test_refinement_uses_global_candidates_despite_reverse_completion_order(tabl
             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.})
+        assert args[-1] is True  # Index tasks request plan-ordered results.
         yield 0, ("l2", {0: 10., 1: 9.})
+        yield 1, ("l2", {2: 8., 3: 7.})
 
     with patch.object(search_module, "_map_tasks", completed):
         result = distributed._read_indexed(splits, [1., 1.], 
query._table._read_snapshot)
@@ -344,8 +350,9 @@ def test_duplicate_scores_keep_plan_order(table):
     splits = [IndexVectorSearchSplit(0, 1, [index_file]), 
IndexVectorSearchSplit(0, 3, [index_file])]
 
     def completed(*args):
-        yield 1, ("l2", {0: 100.})
+        assert args[-1] is True
         yield 0, ("l2", {0: 1., 1: 2.})
+        yield 1, ("l2", {0: 100.})
 
     with patch.object(search_module, "_map_tasks", completed):
         result = distributed._read_indexed(splits, [1., 1.], None)
@@ -370,8 +377,8 @@ def 
test_metric_mismatch_fails_even_when_one_shard_has_no_hits(table):
 
     with patch.object(search_module, "_map_tasks", completed), \
             pytest.raises(ValueError, match="different metrics"):
-        distributed._search_index_splits(
-            [IndexVectorSearchSplit(0, 1, []), IndexVectorSearchSplit(2, 3, 
[])], [1., 1.], 1, [])
+        list(distributed._search_index_splits(
+            [IndexVectorSearchSplit(0, 1, []), IndexVectorSearchSplit(2, 3, 
[])], [1., 1.], 1, []))
     assert closed == [True]
 
 
@@ -407,3 +414,46 @@ def test_nan_scores_fail_before_worker_top_k(table, 
ray_cluster):
     # otherwise a different split layout can silently change the answer.
     with pytest.raises(ValueError, match="cannot rank NaN"):
         table.search([1., 1.]).limit(1).to_arrow(execution="ray")
+
+
[email protected]("ordered", [False, True])
[email protected]("early_close", [False, True])
+def test_slow_first_task_bounds_completed_results(ordered, early_close):
+    submitted, completed, consumed, cancelled = [], [], [], []
+
+    def submit(context, item):
+        # A slow prefix must not allow all later tasks/results to accumulate.
+        assert len(submitted) - len(consumed) < 3
+        submitted.append(item)
+        return item
+
+    def wait(refs, **kwargs):
+        # Complete the newest task first, delaying the first submitted task.
+        ref = max(refs)
+        completed.append(ref)
+        return [ref], [item for item in refs if item != ref]
+
+    remote = SimpleNamespace(remote=submit)
+    remote.options = lambda **kwargs: remote
+    with patch.object(ray, "put", return_value=None), \
+            patch.object(ray, "remote", return_value=remote), \
+            patch.object(ray, "wait", side_effect=wait), \
+            patch.object(ray, "get", side_effect=lambda ref: ref), \
+            patch.object(ray, "cancel", side_effect=cancelled.append):
+        with closing(search_module._map_tasks(None, None, list(range(10)), 3, 
{}, ordered)) as tasks:
+            for ordinal, result in tasks:
+                assert ordinal == result
+                consumed.append(ordinal)
+                if early_close:
+                    break
+    if ordered:
+        assert completed[:3] == [2, 1, 0]
+        assert consumed == ([0] if early_close else list(range(10)))
+    else:
+        assert consumed[0] == 2
+    if early_close:
+        assert submitted == [0, 1, 2]
+        assert cancelled == ([] if ordered else [0, 1])
+    else:
+        assert sorted(consumed) == list(range(10))
+        assert cancelled == []
diff --git a/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py 
b/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py
index 1a230d1c9b..12e0a3ce14 100644
--- a/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py
+++ b/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py
@@ -16,6 +16,7 @@
 # under the License.
 
 import threading
+import weakref
 from concurrent.futures import ThreadPoolExecutor
 from unittest import mock
 
@@ -138,8 +139,8 @@ def test_parallel_failure_closes_all_started_readers(batch, 
failure):
     with mock.patch(MODULE + "._create_vector_reader", side_effect=lambda t, 
f, p, m, o:
                     Reader(int(m[0].file_name))):
         with pytest.raises(ValueError) as exc:
-            read._search_index_splits(
-                splits, [[1.0], [2.0]] if batch else [1.0], 2, None, 
batch=batch)
+            list(read._search_index_splits(
+                splits, [[1.0], [2.0]] if batch else [1.0], 2, None, 
batch=batch))
     if failure == "search":
         assert exc.value is original
     else:
@@ -151,7 +152,7 @@ def test_parallel_failure_closes_all_started_readers(batch, 
failure):
 def test_invalid_parallelism(value):
     read, splits = make_read(False, {"global-index.thread-num": value})
     with pytest.raises(ValueError, match="'global-index.thread-num' must be a 
positive integer"):
-        read._search_index_splits(splits, [1.0], 2, None)
+        list(read._search_index_splits(splits, [1.0], 2, None))
 
 
 @pytest.mark.parametrize("batch", [False, True])
@@ -163,8 +164,8 @@ def test_explicit_serial_and_single_split_fast_path(batch):
         method = "_eval_batch" if batch else "_eval"
         with mock.patch(MODULE + ".ThreadPoolExecutor", 
side_effect=AssertionError("pool")), \
                 mock.patch.object(read, method, 
return_value=_completed_future(None)) as evaluate:
-            assert read._search_index_splits(
-                splits, query, 2, None, batch=batch) == [None] * count
+            assert list(read._search_index_splits(
+                splits, query, 2, None, batch=batch)) == [None] * count
         assert evaluate.call_count == count
 
 
@@ -177,8 +178,8 @@ def test_default_and_configured_worker_limits(batch, count, 
table_value, workers
     method = "_eval_batch" if batch else "_eval"
     with mock.patch(MODULE + ".ThreadPoolExecutor", wraps=ThreadPoolExecutor) 
as executor, \
             mock.patch.object(read, method, 
return_value=_completed_future(None)) as evaluate:
-        assert read._search_index_splits(
-            splits, query, 2, None, batch=batch) == [None] * count
+        assert list(read._search_index_splits(
+            splits, query, 2, None, batch=batch)) == [None] * count
     executor.assert_called_once_with(max_workers=workers)
     assert evaluate.call_count == count
 
@@ -194,4 +195,62 @@ def test_parallelism_from_table_options():
         return _completed_future(None)
 
     with mock.patch.object(read, "_eval", side_effect=evaluate):
-        assert read._search_index_splits(splits, [1.0], 2, None) == [None, 
None]
+        assert list(read._search_index_splits(splits, [1.0], 2, None)) == 
[None, None]
+
+
[email protected]("batch", [False, True])
[email protected]("fail_merge", [False, True])
+def test_merge_releases_results_and_bounds_submissions(batch, fail_merge):
+    read, splits = make_read(batch, {"global-index.thread-num": "2"}, count=20)
+    references, submitted, cancelled = [], [], []
+
+    class Result(DictBasedScoredIndexResult):
+        def score_getter(self):
+            if fail_merge:
+                raise ValueError("injected merge failure")
+            return super().score_getter()
+
+    def evaluate(start, *args):
+        # Completed results from earlier splits must be released before 
refilling.
+        assert sum(ref() is not None for ref in references) <= (2 if batch 
else 1)
+        results = [Result({0: float(start), start + 1: 1.}),
+                   Result({0: -float(start), start + 2: 2.})] if batch else [
+            Result({0: float(start), start + 1: 1.})]
+        references.extend(weakref.ref(result) for result in results)
+        return _completed_future(results if batch else results[0])
+
+    class Executor:
+        def __init__(self, max_workers):
+            assert max_workers == 2
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            pass
+
+        def submit(self, fn, i):
+            submitted.append(i)
+            future = _completed_future(fn(i))
+            future.cancel = lambda: cancelled.append(i)
+            return future
+
+    method = "_eval_batch" if batch else "_eval"
+    with mock.patch(MODULE + ".ThreadPoolExecutor", Executor), \
+            mock.patch.object(read, method, side_effect=evaluate):
+        if fail_merge:
+            with pytest.raises(ValueError, match="injected merge failure"):
+                if batch:
+                    read._read_batch(splits, None)
+                else:
+                    read._read_indexed(splits, [1.0], None)
+            assert submitted == [0, 1]
+            assert cancelled == [1]
+        else:
+            results = read._read_batch(splits, None) if batch else 
[read._read_indexed(splits, [1.0], None)]
+            assert list(results[0].results()) == [1, 11]
+            if batch:
+                assert list(results[1].results()) == [2, 12]
+            assert submitted == list(range(20))
+            assert cancelled == []
+            assert all(ref() is None for ref in references)

Reply via email to