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 9a070f31b9 [python] Search vector index shards concurrently (#9805)
9a070f31b9 is described below
commit 9a070f31b945144e31aeaf57fae3f299054f3c91
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]
