From 80dd2c34ba0762c5258d95998027c236c3521bb1 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Mon, 14 Sep 2026 12:49:52 +0800 Subject: [PATCH 1/2] [python] Search vector index shards concurrently --- .../multimodal-table/global-index/vector.mdx | 9 + .../table/source/vector_search_read.py | 95 ++++++---- .../vector_index_parallel_search_test.py | 178 ++++++++++++++++++ 3 files changed, 247 insertions(+), 35 deletions(-) create mode 100644 paimon-python/pypaimon/tests/vector_index_parallel_search_test.py diff --git a/docs/docs/multimodal-table/global-index/vector.mdx b/docs/docs/multimodal-table/global-index/vector.mdx index afb49e4d4bc7..369b9185d5b6 100644 --- a/docs/docs/multimodal-table/global-index/vector.mdx +++ b/docs/docs/multimodal-table/global-index/vector.mdx @@ -136,6 +136,15 @@ Every query vector must have that same dimension. ## Vector Search +Python indexed vector searches support `vector.search.parallelism`, a positive +integer with default `1`. Set it in table options or override it on a single or +batch vector search builder with `.with_option("vector.search.parallelism", "4")` +to open and search up to four index shards concurrently. It applies to both +paimon-vindex and Lumina indexes. Results retain the same score and row-ID tie +ordering. Each shard can also use native search and I/O threads, so choose the +parallelism together with those settings and the available memory. Raw-vector +fallback continues to use the table read parallelism. + The examples search the three-dimensional IVF-flat index built above. A limit of `5` returns at most five matches; the three-row sample table returns fewer. `ivf.nprobe=1` matches the single cluster in this example. diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index b05f6bb8b4fe..0359805d51a5 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -18,7 +18,8 @@ """Vector search read to read index files.""" from abc import ABC, abstractmethod -from concurrent.futures import ThreadPoolExecutor, wait +from concurrent.futures import ThreadPoolExecutor +from threading import Lock from pypaimon.globalindex.batch_vector_search import BatchVectorSearch from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta @@ -94,7 +95,7 @@ def _search_metric(self, index_type=None): return _raw_search_metric( self._table, self._vector_column, self._options, index_type) - def _record_index_metric(self, reader, index_type): + def _record_index_metric(self, reader, index_type, metric_lock=None): """Keep one persisted metric for indexed scores, raw search and refinement.""" metric_getter = getattr(reader, "vector_metric", None) if metric_getter is None: @@ -106,6 +107,13 @@ def _record_index_metric(self, reader, index_type): raise ValueError( "Query vector metric '%s' does not match index metric '%s' for column '%s'." % (requested, metric, self._vector_column.name)) + if metric_lock is None: + self._set_index_metric(metric) + else: + with metric_lock: + self._set_index_metric(metric) + + def _set_index_metric(self, metric): if self._index_metric is not None and self._index_metric != metric: raise ValueError( "Cannot merge vector indexes with different metrics '%s' and '%s' for column '%s'." @@ -235,7 +243,8 @@ def _raw_pre_filter(self, splits, snapshot=None): finally: scanner.close() - def _open_offset_reader(self, vector_index_files, row_range_start, row_range_end): + def _open_offset_reader(self, vector_index_files, row_range_start, row_range_end, + metric_lock=None): """Open a vector index reader for the split, wrapped with the row-id offset. The caller must close the returned reader once its future completes. @@ -261,14 +270,14 @@ def _open_offset_reader(self, vector_index_files, row_range_start, row_range_end self._table.table_schema.options, ) try: - self._record_index_metric(reader, vector_index_files[0].index_type) + self._record_index_metric(reader, vector_index_files[0].index_type, metric_lock) return reader, OffsetGlobalIndexReader(reader, row_range_start, row_range_end) except Exception: reader.close() raise def _eval(self, row_range_start, row_range_end, vector_index_files, - query_vector, search_limit, include_row_ids): + query_vector, search_limit, include_row_ids, metric_lock=None): from pypaimon.globalindex.global_index_reader import _completed_future if not vector_index_files: @@ -284,7 +293,7 @@ def _eval(self, row_range_start, row_range_end, vector_index_files, vector_search = vector_search.with_include_row_ids(include_row_ids) reader, offset_reader = self._open_offset_reader( - vector_index_files, row_range_start, row_range_end) + vector_index_files, row_range_start, row_range_end, metric_lock) try: future = offset_reader.visit_vector_search(vector_search) except BaseException: @@ -426,7 +435,7 @@ def _raw_search_projection(self, include_filter): return projection def _eval_batch(self, row_range_start, row_range_end, vector_index_files, - query_vectors, search_limit, include_row_ids): + query_vectors, search_limit, include_row_ids, metric_lock=None): from pypaimon.globalindex.global_index_reader import _completed_future if not vector_index_files: @@ -442,7 +451,7 @@ def _eval_batch(self, row_range_start, row_range_end, vector_index_files, batch_vector_search = batch_vector_search.with_include_row_ids(include_row_ids) reader, offset_reader = self._open_offset_reader( - vector_index_files, row_range_start, row_range_end) + vector_index_files, row_range_start, row_range_end, metric_lock) try: future = offset_reader.visit_batch_vector_search(batch_vector_search) except BaseException: @@ -451,6 +460,43 @@ def _eval_batch(self, row_range_start, row_range_end, vector_index_files, future.add_done_callback(lambda _: reader.close()) return future + def _search_index_splits(self, splits, query, search_limit, pre_filters, batch=False): + # Native readers finish their search before returning a completed Future. + # Schedule the entire open/search/close operation, not just Future.result(). + key = "vector.search.parallelism" + value = self._options.get(key, _table_options_map(self._table).get(key, "1")) + try: + parallelism = int(str(value)) + except (ValueError, TypeError): + parallelism = 0 + if parallelism < 1: + raise ValueError("'%s' must be a positive integer, got: %s" % (key, value)) + + metric_lock = Lock() + evaluate = self._eval_batch if batch else self._eval + + def search(i): + split = splits[i] + return evaluate( + split.row_range_start, split.row_range_end, split.vector_index_files, + query, search_limit, None if not pre_filters else pre_filters[i], + metric_lock, + ).result() + + workers = min(parallelism, len(splits)) + if workers <= 1: + return [search(i) for i in range(len(splits))] + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = [] + try: + for i in range(len(splits)): + futures.append(executor.submit(search, i)) + return [future.result() for future in futures] + finally: + for future in futures: + future.cancel() + # Executor shutdown waits for started readers to close on failure. + def _indexed_search_limit(self, index_type): refine_factor = self._configured_refine_factor(index_type) if refine_factor == 0: @@ -552,22 +598,11 @@ def _read_indexed(self, splits, query_vector, snapshot): index_type = _vector_index_type(splits) search_limit = self._indexed_search_limit(index_type) pre_filters = self._pre_filters(splits, snapshot) - futures = [ - self._eval( - split.row_range_start, split.row_range_end, - split.vector_index_files, - query_vector, - search_limit, - None if not pre_filters else pre_filters[i] - ) - for i, split in enumerate(splits) - ] - - wait(futures) + results = self._search_index_splits( + splits, query_vector, search_limit, pre_filters) merged_scores = {} - for future in futures: - split_result = future.result() + for split_result in results: if split_result is not None: score_getter = split_result.score_getter() for row_id in split_result.results(): @@ -603,22 +638,12 @@ def _read_batch(self, splits, snapshot): index_type = _vector_index_type(index_splits) search_limit = self._indexed_search_limit(index_type) pre_filters = self._pre_filters(index_splits, snapshot) - futures = [ - self._eval_batch( - split.row_range_start, split.row_range_end, - split.vector_index_files, self._query_vectors, - search_limit, - None if not pre_filters else pre_filters[i], - ) - for i, split in enumerate(index_splits) - ] - - wait(futures) + 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 future in futures: - split_results = future.result() + for split_results in results: for i in range(n): split_result = split_results[i] if split_result is None: diff --git a/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py b/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py new file mode 100644 index 000000000000..295fbea91f5e --- /dev/null +++ b/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py @@ -0,0 +1,178 @@ +# 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. + +import threading +from unittest import mock + +import pytest + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.common.options.options import Options +from pypaimon.globalindex.global_index_reader import _completed_future +from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.table.source.vector_search_read import ( + BatchVectorSearchReadImpl, DataEvolutionVectorRead, +) +from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit +from pypaimon.tests.vector_search_filter_test import _StubTable, _field, _entry, _bitmap + +MODULE = "pypaimon.table.source.vector_search_read" + + +def make_read(batch, options=None, count=4): + field = _field(1, "embedding", "FLOAT") + table = _StubTable([field], []) + cls = BatchVectorSearchReadImpl if batch else DataEvolutionVectorRead + read = cls(table, 2, field, [[1.0], [2.0]] if batch else [1.0], options=options) + splits = [] + for i in range(count): + entry = _entry(None, 1, "ivf-flat", str(i), i * 10, i * 10 + 9) + splits.append(IndexVectorSearchSplit(i * 10, i * 10 + 9, [entry.index_file])) + return read, splits + + +@pytest.mark.parametrize("batch", [False, True]) +def test_parallel_open_search_close_and_ordered_results(batch): + read, splits = make_read(batch, {"vector.search.parallelism": "2"}) + barrier = threading.Barrier(2, timeout=5) + lock = threading.Lock() + active, peak, closed, filters = set(), [0], set(), {} + + class Reader: + def __init__(self, index): + self.index = index + with lock: + active.add(index) + peak[0] = max(peak[0], len(active)) + + def vector_metric(self): + # Opening/loading must overlap too, not only native search. + barrier.wait() + return "l2" + + def visit_vector_search(self, query): + filters[self.index] = list(query.include_row_ids) + return _completed_future(DictBasedScoredIndexResult({1: 0.5, 2: 0.5})) + + def visit_batch_vector_search(self, query): + filters[self.index] = list(query.include_row_ids) + return _completed_future([ + DictBasedScoredIndexResult({1: 0.5, 2: 0.5}), + DictBasedScoredIndexResult({3: 1.0}), + ]) + + def close(self): + with lock: + active.remove(self.index) + closed.add(self.index) + + def factory(index_type, file_io, path, metas, options): + return Reader(int(metas[0].file_name)) + + with mock.patch(MODULE + "._create_vector_reader", side_effect=factory), \ + mock.patch.object(read, "_pre_filters", + return_value=[_bitmap(i * 10 + 1) for i in range(4)]): + if batch: + result = read._read_batch(splits, None) + assert list(result[0].results()) == [1, 2] + assert list(result[1].results()) == [3, 13] + else: + result = read._read_indexed(splits, [1.0], None) + assert list(result.results()) == [1, 2] + assert peak == [2] + assert not active + assert closed == set(range(4)) + assert filters == {i: [1] for i in range(4)} + assert read._index_metric == "l2" + + +@pytest.mark.parametrize("batch", [False, True]) +@pytest.mark.parametrize("failure", ["search", "metric"]) +def test_parallel_failure_closes_all_started_readers(batch, failure): + read, splits = make_read(batch, {"vector.search.parallelism": "2"}, count=2) + barrier = threading.Barrier(2, timeout=5) + closed = set() + original = ValueError("native search failed") + + class Reader: + def __init__(self, i): + self.i = i + + def vector_metric(self): + barrier.wait() + return "cosine" if failure == "metric" and self.i else "l2" + + def search(self, query): + barrier.wait() + if self.i == 0: + raise original + return _completed_future([None, None] if batch else None) + + visit_vector_search = search + visit_batch_vector_search = search + + def close(self): + closed.add(self.i) + + # Metric mismatch fails before search; no search barrier is needed in that case. + if failure == "metric": + Reader.visit_vector_search = lambda self, q: _completed_future(None) + Reader.visit_batch_vector_search = lambda self, q: _completed_future([None, None]) + + 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) + if failure == "search": + assert exc.value is original + else: + assert "different metrics" in str(exc.value) + assert closed == {0, 1} + + +@pytest.mark.parametrize("value", ["0", "-1", "1.5", "invalid", True]) +def test_invalid_parallelism(value): + read, splits = make_read(False, {"vector.search.parallelism": value}) + with pytest.raises(ValueError, match="positive integer"): + read._search_index_splits(splits, [1.0], 2, None) + + +def test_serial_default_query_override_and_single_split_fast_path(): + for count, table_value, query_options in ( + (3, None, {}), (3, "4", {"vector.search.parallelism": "1"}), + (1, "4", {}), (0, "4", {})): + read, splits = make_read(False, query_options, count) + read._table.options = CoreOptions(Options( + {} if table_value is None else {"vector.search.parallelism": table_value})) + with mock.patch(MODULE + ".ThreadPoolExecutor", side_effect=AssertionError("pool")), \ + mock.patch.object(read, "_eval", return_value=_completed_future(None)) as evaluate: + assert read._search_index_splits(splits, [1.0], 2, None) == [None] * count + assert evaluate.call_count == count + + +def test_parallelism_from_table_options(): + read, splits = make_read(False, count=2) + read._table.options = CoreOptions(Options({"vector.search.parallelism": "2"})) + barrier = threading.Barrier(2, timeout=5) + + def evaluate(*args): + barrier.wait() + 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] From d0753103c9d77e9762365aa4af04bcab9709bfef Mon Sep 17 00:00:00 2001 From: chaoyang Date: Tue, 15 Sep 2026 12:02:44 +0800 Subject: [PATCH 2/2] [python] Reuse global index thread count for vector searches --- .../multimodal-table/global-index/vector.mdx | 17 +++--- .../table/source/vector_search_read.py | 6 +- .../vector_index_parallel_search_test.py | 55 +++++++++++++------ 3 files changed, 50 insertions(+), 28 deletions(-) diff --git a/docs/docs/multimodal-table/global-index/vector.mdx b/docs/docs/multimodal-table/global-index/vector.mdx index 369b9185d5b6..98ea0a5372a2 100644 --- a/docs/docs/multimodal-table/global-index/vector.mdx +++ b/docs/docs/multimodal-table/global-index/vector.mdx @@ -136,14 +136,15 @@ Every query vector must have that same dimension. ## Vector Search -Python indexed vector searches support `vector.search.parallelism`, a positive -integer with default `1`. Set it in table options or override it on a single or -batch vector search builder with `.with_option("vector.search.parallelism", "4")` -to open and search up to four index shards concurrently. It applies to both -paimon-vindex and Lumina indexes. Results retain the same score and row-ID tie -ordering. Each shard can also use native search and I/O threads, so choose the -parallelism together with those settings and the available memory. Raw-vector -fallback continues to use the table read parallelism. +Python indexed vector searches use the `global-index.thread-num` table option, +the same as Java. It must be a positive integer and defaults to `32`. Single and +batch queries open and search at most that many index shards concurrently, +bounded by the number of splits. Set it to `1` for serial shard searches. It +applies to both paimon-vindex and Lumina indexes. Results retain the same score +and row-ID tie ordering. Each shard can also use native search and I/O threads, +so choose the thread count together with those settings and the available +memory. In Python, raw-vector fallback continues to use the table read +parallelism. The examples search the three-dimensional IVF-flat index built above. A limit of `5` returns at most five matches; the three-row sample table returns fewer. diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 0359805d51a5..4a034a3af4a6 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -21,6 +21,7 @@ from concurrent.futures import ThreadPoolExecutor from threading import Lock +from pypaimon.common.options.core_options import CoreOptions from pypaimon.globalindex.batch_vector_search import BatchVectorSearch from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta from pypaimon.globalindex.global_index_result import GlobalIndexResult @@ -463,8 +464,9 @@ def _eval_batch(self, row_range_start, row_range_end, vector_index_files, def _search_index_splits(self, splits, query, search_limit, pre_filters, batch=False): # Native readers finish their search before returning a completed Future. # Schedule the entire open/search/close operation, not just Future.result(). - key = "vector.search.parallelism" - value = self._options.get(key, _table_options_map(self._table).get(key, "1")) + option = CoreOptions.GLOBAL_INDEX_THREAD_NUM + key = option.key() + value = _table_options_map(self._table).get(key, option.default_value()) try: parallelism = int(str(value)) except (ValueError, TypeError): 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 295fbea91f5e..1a230d1c9b41 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 +from concurrent.futures import ThreadPoolExecutor from unittest import mock import pytest @@ -33,11 +34,12 @@ MODULE = "pypaimon.table.source.vector_search_read" -def make_read(batch, options=None, count=4): +def make_read(batch, table_options=None, count=4): field = _field(1, "embedding", "FLOAT") table = _StubTable([field], []) + table.options = CoreOptions(Options(table_options or {})) cls = BatchVectorSearchReadImpl if batch else DataEvolutionVectorRead - read = cls(table, 2, field, [[1.0], [2.0]] if batch else [1.0], options=options) + read = cls(table, 2, field, [[1.0], [2.0]] if batch else [1.0]) splits = [] for i in range(count): entry = _entry(None, 1, "ivf-flat", str(i), i * 10, i * 10 + 9) @@ -47,7 +49,7 @@ def make_read(batch, options=None, count=4): @pytest.mark.parametrize("batch", [False, True]) def test_parallel_open_search_close_and_ordered_results(batch): - read, splits = make_read(batch, {"vector.search.parallelism": "2"}) + read, splits = make_read(batch, {"global-index.thread-num": "2"}) barrier = threading.Barrier(2, timeout=5) lock = threading.Lock() active, peak, closed, filters = set(), [0], set(), {} @@ -103,7 +105,7 @@ def factory(index_type, file_io, path, metas, options): @pytest.mark.parametrize("batch", [False, True]) @pytest.mark.parametrize("failure", ["search", "metric"]) def test_parallel_failure_closes_all_started_readers(batch, failure): - read, splits = make_read(batch, {"vector.search.parallelism": "2"}, count=2) + read, splits = make_read(batch, {"global-index.thread-num": "2"}, count=2) barrier = threading.Barrier(2, timeout=5) closed = set() original = ValueError("native search failed") @@ -145,29 +147,46 @@ def close(self): assert closed == {0, 1} -@pytest.mark.parametrize("value", ["0", "-1", "1.5", "invalid", True]) +@pytest.mark.parametrize("value", ["0", "-1", "1.5", "invalid", True, 1.5]) def test_invalid_parallelism(value): - read, splits = make_read(False, {"vector.search.parallelism": value}) - with pytest.raises(ValueError, match="positive integer"): + 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) -def test_serial_default_query_override_and_single_split_fast_path(): - for count, table_value, query_options in ( - (3, None, {}), (3, "4", {"vector.search.parallelism": "1"}), - (1, "4", {}), (0, "4", {})): - read, splits = make_read(False, query_options, count) - read._table.options = CoreOptions(Options( - {} if table_value is None else {"vector.search.parallelism": table_value})) +@pytest.mark.parametrize("batch", [False, True]) +def test_explicit_serial_and_single_split_fast_path(batch): + for count, table_value in ((3, "1"), (1, None), (0, None), (1, "4"), (0, "4")): + read, splits = make_read( + batch, {} if table_value is None else {"global-index.thread-num": table_value}, count) + query = [[1.0], [2.0]] if batch else [1.0] + method = "_eval_batch" if batch else "_eval" with mock.patch(MODULE + ".ThreadPoolExecutor", side_effect=AssertionError("pool")), \ - mock.patch.object(read, "_eval", return_value=_completed_future(None)) as evaluate: - assert read._search_index_splits(splits, [1.0], 2, None) == [None] * count + 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 evaluate.call_count == count +@pytest.mark.parametrize("batch", [False, True]) +@pytest.mark.parametrize("count, table_value, workers", [(2, None, 2), (33, None, 32), (4, "3", 3)]) +def test_default_and_configured_worker_limits(batch, count, table_value, workers): + read, splits = make_read( + batch, {} if table_value is None else {"global-index.thread-num": table_value}, count) + query = [[1.0], [2.0]] if batch else [1.0] + 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 + executor.assert_called_once_with(max_workers=workers) + assert evaluate.call_count == count + + def test_parallelism_from_table_options(): - read, splits = make_read(False, count=2) - read._table.options = CoreOptions(Options({"vector.search.parallelism": "2"})) + read, splits = make_read(False, {"global-index.thread-num": "2"}, count=2) + # Like Java, shard concurrency comes from the table, not native query options. + read._options["global-index.thread-num"] = "1" barrier = threading.Barrier(2, timeout=5) def evaluate(*args):