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 eecae6d22d8bf8d8f664913775c373304ab05c95 Author: chaoyang <[email protected]> AuthorDate: Thu Sep 17 14:37:55 2026 +0800 [python] Search vector index shards concurrently (#9805) --- docs/docs/multimodal-table/global-index/vector.mdx | 10 ++ .../pypaimon/table/source/vector_search_read.py | 97 ++++++---- .../tests/vector_index_parallel_search_test.py | 197 +++++++++++++++++++++ 3 files changed, 269 insertions(+), 35 deletions(-) diff --git a/docs/docs/multimodal-table/global-index/vector.mdx b/docs/docs/multimodal-table/global-index/vector.mdx index afb49e4d4b..98ea0a5372 100644 --- a/docs/docs/multimodal-table/global-index/vector.mdx +++ b/docs/docs/multimodal-table/global-index/vector.mdx @@ -136,6 +136,16 @@ Every query vector must have that same dimension. ## Vector Search +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. `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 ef891f12b0..c473d36308 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -18,8 +18,10 @@ """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.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 @@ -94,7 +96,7 @@ class AbstractVectorSearchReadImpl: 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 +108,13 @@ class AbstractVectorSearchReadImpl: 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 +244,8 @@ class AbstractVectorSearchReadImpl: 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 +271,14 @@ class AbstractVectorSearchReadImpl: 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 +294,7 @@ class AbstractVectorSearchReadImpl: 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 +436,7 @@ class AbstractVectorSearchReadImpl: 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 +452,7 @@ class AbstractVectorSearchReadImpl: 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 +461,44 @@ class AbstractVectorSearchReadImpl: 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(). + 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): + 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: @@ -597,22 +645,11 @@ 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) - 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(): @@ -648,22 +685,12 @@ 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) - 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 0000000000..1a230d1c9b --- /dev/null +++ b/paimon-python/pypaimon/tests/vector_index_parallel_search_test.py @@ -0,0 +1,197 @@ +# 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 concurrent.futures import ThreadPoolExecutor +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, 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]) + 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 + + [email protected]("batch", [False, True]) +def test_parallel_open_search_close_and_ordered_results(batch): + 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(), {} + + 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" + + [email protected]("batch", [False, True]) [email protected]("failure", ["search", "metric"]) +def test_parallel_failure_closes_all_started_readers(batch, failure): + 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") + + 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} + + [email protected]("value", ["0", "-1", "1.5", "invalid", True, 1.5]) +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) + + [email protected]("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, 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 + + [email protected]("batch", [False, True]) [email protected]("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, {"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): + 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]
