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 110985f3bf [core][docs] Parallelize primary-key vector search (#8668)
110985f3bf is described below
commit 110985f3bf6aac2d80f537f28789e4f150477cf4
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Jul 15 23:42:27 2026 +0800
[core][docs] Parallelize primary-key vector search (#8668)
Primary-key vector reads currently wait for each bucket and ANN segment
before starting the next one, so local search latency grows serially
with the number of buckets and segments.
---
...026-07-15-primary-key-vector-parallel-search.md | 75 +++++++++++
docs/docs/primary-key-table/global-index.mdx | 24 ++--
.../index/pkvector/PkVectorAnnSegmentSearcher.java | 127 ++++++++++++-------
.../pkvector/PrimaryKeyVectorBucketSearch.java | 118 ++++++++++++-----
.../paimon/table/source/PrimaryKeyVectorRead.java | 49 ++++---
.../pkvector/PrimaryKeyVectorBucketSearchTest.java | 141 ++++++++++++++++++++-
.../table/source/PrimaryKeyVectorReadTest.java | 105 +++++++++++++++
7 files changed, 533 insertions(+), 106 deletions(-)
diff --git a/docs/designs/2026-07-15-primary-key-vector-parallel-search.md
b/docs/designs/2026-07-15-primary-key-vector-parallel-search.md
new file mode 100644
index 0000000000..8135f63310
--- /dev/null
+++ b/docs/designs/2026-07-15-primary-key-vector-parallel-search.md
@@ -0,0 +1,75 @@
+<!--
+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.
+-->
+
+# Primary-Key Vector Parallel Search Design
+
+## Problem Statement
+
+`PrimaryKeyVectorRead` searches bucket splits serially, and
+`PrimaryKeyVectorBucketSearch` waits for each ANN segment before starting the
next one. A query
+over multiple buckets or segments therefore accumulates their local-search
latency.
+
+## Chosen Approach
+
+Use the existing global-index read executor and `global-index.thread-num` to
compose bucket,
+segment, and uncovered-file searches with `CompletableFuture`. Submit only
leaf searches to the
+executor and wait once after all searches have been started. This avoids
blocking an executor task
+while it waits for more work submitted to the same executor.
+
+## Design Details
+
+### Asynchronous Search
+
+- Add asynchronous ANN segment search while retaining synchronous wrappers for
compatibility.
+- Compose all segment futures within a bucket instead of joining each segment
immediately.
+- Compose all bucket futures in `PrimaryKeyVectorRead` and preserve the
existing deterministic
+ global Top-K merge.
+- Submit FULL and DETAIL exact searches for uncovered files as independent
leaf tasks.
+
+### Resource and Error Handling
+
+- Attach index-reader closure directly to the future returned by the index
reader.
+- Close exact-search readers within their leaf tasks.
+- Do not submit additional work from completion callbacks.
+- Preserve the existing top-level `Failed to search primary-key vector index.`
error contract.
+
+### Compatibility
+
+- Keep JDK 8 syntax and existing synchronous search entry points.
+- Do not add a new option or executor.
+- Keep bucket preparation, deletion-vector loading, residual-filter
evaluation, reranking, and
+ result ordering unchanged.
+
+### Verification
+
+- Prove multiple ANN segment searches can be in flight together.
+- Prove multiple bucket searches are composed concurrently and merge
deterministically.
+- Prove FULL and DETAIL exact fallback searches uncovered files concurrently.
+- Verify exceptional completion closes readers and is propagated.
+- Verify `global-index.thread-num=1` completes without deadlock.
+
+## Open Questions
+
+None.
+
+## Out of Scope
+
+- Fair scheduling across buckets when a semaphore-limited executor blocks
submissions.
+- Parallel bucket metadata, deletion-vector, residual-filter, or rerank reads.
+- Changes to `GlobalIndexReadThreadPool` or its executor wrappers.
diff --git a/docs/docs/primary-key-table/global-index.mdx
b/docs/docs/primary-key-table/global-index.mdx
index a8f3dea2f0..ba59a1e972 100644
--- a/docs/docs/primary-key-table/global-index.mdx
+++ b/docs/docs/primary-key-table/global-index.mdx
@@ -236,6 +236,7 @@ schema validation.
| `fields.<column>.pk-bitmap.index.options` | Not set | JSON object containing
Bitmap build options. Unqualified keys are scoped to `bitmap-index`. |
| `fields.<column>.pk-index.compaction.level-fanout` | `5` | Number of
similarly sized index groups which triggers a rebuild and maximum row-count
ratio within one size tier. Shared by all four families. Must be greater than
`1`. |
| `fields.<column>.pk-index.compaction.stale-ratio-threshold` | `0.2` | Ratio
of rows from inactive source files which triggers a rebuild. Shared by all four
families. Must be in `(0, 1]`. |
+| `global-index.search-mode` | `fast` | Search mode for primary-key Vector and
Full Text queries. `fast` searches indexed data only, so uncovered files are
omitted. For Vector, `full` and `detail` search uncovered files exactly.
Primary-key Full Text supports only `fast`. |
For algorithm-specific options, see the corresponding
[BTree](../multimodal-table/global-index/btree),
@@ -274,11 +275,16 @@ Index construction can execute asynchronously inside the
writer. A writer which
compaction also waits for active index maintenance; a non-blocking writer can
complete maintenance
in a later commit. Coverage can therefore be temporarily partial.
-For scalar and vector indexes, partial coverage affects acceleration, not
correctness:
+Coverage behavior depends on the index family and search mode.
`global-index.search-mode` defaults
+to `fast`. For primary-key Vector and Full Text queries, `fast` searches only
data covered by an
+active index group. Uncovered files are not searched, so partial coverage can
make results
+incomplete.
-- BTree and Bitmap scans read uncovered files through the ordinary data path.
-- Vector search evaluates files without an active ANN group exactly.
-- The original scalar predicate and deletion vectors are applied after index
pruning.
+- BTree and Bitmap scans always read uncovered files through the ordinary data
path. Partial
+ coverage affects their acceleration, not result completeness. The original
scalar predicate and
+ deletion vectors are applied after index pruning.
+- Vector search in `full` or `detail` mode evaluates files without an active
ANN group exactly and
+ merges those results with ANN candidates. In the default `fast` mode, those
files are omitted.
Primary-key Full Text currently supports only `global-index.search-mode =
fast`. It searches
persistent archives and ignores uncovered files; `full` and `detail` are
rejected because a
@@ -388,8 +394,9 @@ factor must be a positive integer. A factor of `1` performs
exact reranking with
additional candidates.
Only ANN candidates can win the rerank, so a larger factor can improve recall
but does not
-guarantee the exact global Top-K. It also increases ANN work and data-file
I/O. Uncovered files are
-searched exactly and merged separately with the ANN candidates.
+guarantee the exact global Top-K. It also increases ANN work and data-file
I/O. In `full` or
+`detail` mode, uncovered files are searched exactly and merged separately with
the ANN candidates;
+the default `fast` mode does not search them.
## Full-Text Search
@@ -514,8 +521,9 @@ create a table with the desired definition and migrate the
data.
- Index acceleration and vector search are snapshot-scoped batch operations;
continuous streaming
and lateral Vector or Full Text search are not supported.
- Flink vector search returns rows but does not expose the ANN score as a
separate column.
-- Primary-key Full Text supports only FAST search mode and excludes uncovered
files until
- compaction creates persistent archives.
+- FAST is the default search mode and excludes uncovered files from
primary-key Vector and Full
+ Text results. Vector supports exact fallback for uncovered files in FULL and
DETAIL modes;
+ primary-key Full Text supports only FAST until compaction creates persistent
archives.
- Full Text routes support partition pruning but not arbitrary row predicates
before Top-K.
- Hybrid search cannot mix source-backed physical routes with global row-ID
routes.
- Online replacement between two definitions on the same column is not
supported.
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
index 797afd104c..58696e5a0b 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
@@ -48,6 +48,7 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import static org.apache.paimon.utils.Preconditions.checkArgument;
@@ -148,6 +149,27 @@ public class PkVectorAnnSegmentSearcher {
Set<String> activeSourceFiles,
Map<String, List<Range>> rowRangesByFile,
Map<String, String> searchOptions) {
+ return searchAsync(
+ segment,
+ sourceMeta,
+ query,
+ limit,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ searchOptions)
+ .join();
+ }
+
+ CompletableFuture<List<PkVectorSearchResult>> searchAsync(
+ IndexFileMeta segment,
+ PrimaryKeyIndexSourceMeta sourceMeta,
+ float[] query,
+ int limit,
+ Map<String, DeletionVector> deletionVectors,
+ Set<String> activeSourceFiles,
+ Map<String, List<Range>> rowRangesByFile,
+ Map<String, String> searchOptions) {
checkArgument(limit > 0, "Vector search limit must be positive: %s.",
limit);
GlobalIndexMeta globalIndexMeta = segment.globalIndexMeta();
checkArgument(
@@ -189,51 +211,70 @@ public class PkVectorAnnSegmentSearcher {
if (liveRows != null) {
search.withIncludeRowIds(liveRows);
}
- Optional<ScoredGlobalIndexResult> result =
reader.visitVectorSearch(search).join();
- if (!result.isPresent()) {
- return Collections.emptyList();
- }
-
- long sourceRowCount = totalRowCount(sourceMeta.sourceFiles());
- List<PkVectorSearchResult> candidates = new ArrayList<>();
- ScoredGlobalIndexResult scored = result.get();
- for (long ordinal : scored.results()) {
- checkArgument(
- ordinal >= 0 && ordinal < sourceRowCount,
- "ANN segment %s returned ordinal %s outside [0, %s).",
- segment.fileName(),
- ordinal,
- sourceRowCount);
- FilePosition filePosition =
filePosition(sourceMeta.sourceFiles(), ordinal);
- checkArgument(
- activeSourceFiles.contains(filePosition.dataFileName),
- "ANN segment %s returned inactive source %s.",
- segment.fileName(),
- filePosition.dataFileName);
- DeletionVector deletionVector =
deletionVectors.get(filePosition.dataFileName);
- checkArgument(
- deletionVector == null
- ||
!deletionVector.isDeleted(filePosition.rowPosition),
- "ANN segment %s returned snapshot-deleted row position
%s.",
- segment.fileName(),
- filePosition.rowPosition);
- List<Range> rowRanges =
rowRangesByFile.get(filePosition.dataFileName);
- checkArgument(
- rowRanges == null || contains(rowRanges,
filePosition.rowPosition),
- "ANN segment %s returned a row outside the
pre-filter.",
- segment.fileName());
- candidates.add(
- new PkVectorSearchResult(
- filePosition.dataFileName,
- filePosition.rowPosition,
- VectorSearchMetric.scoreToDistance(
- scored.scoreGetter().score(ordinal),
metric)));
- }
- Collections.sort(candidates, BEST_FIRST);
- return Collections.unmodifiableList(candidates);
- } finally {
+ return reader.visitVectorSearch(search)
+ .whenComplete((ignored, error) ->
IOUtils.closeQuietly(reader))
+ .thenApply(
+ result ->
+ mapResults(
+ segment,
+ sourceMeta,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ result));
+ } catch (RuntimeException | Error t) {
IOUtils.closeQuietly(reader);
+ throw t;
+ }
+ }
+
+ private List<PkVectorSearchResult> mapResults(
+ IndexFileMeta segment,
+ PrimaryKeyIndexSourceMeta sourceMeta,
+ Map<String, DeletionVector> deletionVectors,
+ Set<String> activeSourceFiles,
+ Map<String, List<Range>> rowRangesByFile,
+ Optional<ScoredGlobalIndexResult> result) {
+ if (!result.isPresent()) {
+ return Collections.emptyList();
+ }
+
+ long sourceRowCount = totalRowCount(sourceMeta.sourceFiles());
+ List<PkVectorSearchResult> candidates = new ArrayList<>();
+ ScoredGlobalIndexResult scored = result.get();
+ for (long ordinal : scored.results()) {
+ checkArgument(
+ ordinal >= 0 && ordinal < sourceRowCount,
+ "ANN segment %s returned ordinal %s outside [0, %s).",
+ segment.fileName(),
+ ordinal,
+ sourceRowCount);
+ FilePosition filePosition = filePosition(sourceMeta.sourceFiles(),
ordinal);
+ checkArgument(
+ activeSourceFiles.contains(filePosition.dataFileName),
+ "ANN segment %s returned inactive source %s.",
+ segment.fileName(),
+ filePosition.dataFileName);
+ DeletionVector deletionVector =
deletionVectors.get(filePosition.dataFileName);
+ checkArgument(
+ deletionVector == null ||
!deletionVector.isDeleted(filePosition.rowPosition),
+ "ANN segment %s returned snapshot-deleted row position
%s.",
+ segment.fileName(),
+ filePosition.rowPosition);
+ List<Range> rowRanges =
rowRangesByFile.get(filePosition.dataFileName);
+ checkArgument(
+ rowRanges == null || contains(rowRanges,
filePosition.rowPosition),
+ "ANN segment %s returned a row outside the pre-filter.",
+ segment.fileName());
+ candidates.add(
+ new PkVectorSearchResult(
+ filePosition.dataFileName,
+ filePosition.rowPosition,
+ VectorSearchMetric.scoreToDistance(
+ scored.scoreGetter().score(ordinal),
metric)));
}
+ Collections.sort(candidates, BEST_FIRST);
+ return Collections.unmodifiableList(candidates);
}
@Nullable
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
index d864a3a3fb..ed2e2c4114 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
@@ -38,6 +38,9 @@ import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.Executor;
import java.util.function.LongPredicate;
import static org.apache.paimon.utils.Preconditions.checkArgument;
@@ -115,18 +118,36 @@ public class PrimaryKeyVectorBucketSearch {
int indexedLimit,
int exactLimit)
throws IOException {
+ return join(
+ searchAsync(
+ state,
+ activeFiles,
+ deletionVectors,
+ rowRangesByFile,
+ query,
+ indexedLimit,
+ exactLimit,
+ Runnable::run));
+ }
+
+ public CompletableFuture<Result> searchAsync(
+ PkVectorBucketIndexState state,
+ List<DataFileMeta> activeFiles,
+ Map<String, DeletionVector> deletionVectors,
+ Map<String, List<Range>> rowRangesByFile,
+ float[] query,
+ int indexedLimit,
+ int exactLimit,
+ Executor executor) {
checkArgument(indexedLimit > 0, "Vector indexed search limit must be
positive.");
checkArgument(exactLimit > 0, "Vector exact search limit must be
positive.");
Map<String, DataFileMeta> filesByName = new HashMap<>();
for (DataFileMeta file : activeFiles) {
checkArgument(filesByName.put(file.fileName(), file) == null,
"Duplicate data file.");
}
- PriorityQueue<PkVectorSearchResult> indexedNearest =
- new PriorityQueue<>(indexedLimit, BEST_FIRST.reversed());
- PriorityQueue<PkVectorSearchResult> exactNearest =
- new PriorityQueue<>(exactLimit, BEST_FIRST.reversed());
Set<String> activeSourceFiles = new HashSet<>(filesByName.keySet());
Set<String> covered = new HashSet<>();
+ List<CompletableFuture<List<PkVectorSearchResult>>> indexedFutures =
new ArrayList<>();
for (IndexFileMeta ann : state.annSegments()) {
PrimaryKeyIndexSourceMeta sourceMeta =
PrimaryKeyIndexSourceMeta.fromIndexFile(ann);
for (PrimaryKeyIndexSourceFile source : sourceMeta.sourceFiles()) {
@@ -141,30 +162,19 @@ public class PrimaryKeyVectorBucketSearch {
covered.add(source.fileName());
}
checkArgument(annSearcher != null, "ANN search is not
configured.");
- List<PkVectorSearchResult> annResults =
- rowRangesByFile.isEmpty()
- ? annSearcher.search(
- ann,
- sourceMeta,
- query,
- indexedLimit,
- deletionVectors,
- activeSourceFiles,
- searchOptions)
- : annSearcher.search(
- ann,
- sourceMeta,
- query,
- indexedLimit,
- deletionVectors,
- activeSourceFiles,
- rowRangesByFile,
- searchOptions);
- for (PkVectorSearchResult result : annResults) {
- add(indexedNearest, result, indexedLimit);
- }
+ indexedFutures.add(
+ annSearcher.searchAsync(
+ ann,
+ sourceMeta,
+ query,
+ indexedLimit,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ searchOptions));
}
+ List<CompletableFuture<List<PkVectorSearchResult>>> exactFutures = new
ArrayList<>();
if (searchMode != GlobalIndexSearchMode.FAST) {
for (DataFileMeta file : activeFiles) {
if (covered.contains(file.fileName())) {
@@ -179,16 +189,56 @@ public class PrimaryKeyVectorBucketSearch {
position ->
(dv != null && dv.isDeleted(position))
|| (rowRanges != null &&
!contains(rowRanges, position));
- try (PkVectorReader reader = vectorReaderFactory.create(file))
{
- for (PkVectorSearchResult result :
- PkVectorExactSearcher.search(
- file.fileName(), reader, query, metric,
exactLimit, excluded)) {
- add(exactNearest, result, exactLimit);
- }
- }
+ exactFutures.add(
+ CompletableFuture.supplyAsync(
+ () -> exactSearch(file, query, exactLimit,
excluded), executor));
+ }
+ }
+ List<CompletableFuture<?>> futures = new ArrayList<>();
+ futures.addAll(indexedFutures);
+ futures.addAll(exactFutures);
+ return CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0]))
+ .thenApply(
+ ignored -> {
+ PriorityQueue<PkVectorSearchResult> indexedNearest
=
+ new PriorityQueue<>(indexedLimit,
BEST_FIRST.reversed());
+ for (CompletableFuture<List<PkVectorSearchResult>>
future :
+ indexedFutures) {
+ for (PkVectorSearchResult result :
future.join()) {
+ add(indexedNearest, result, indexedLimit);
+ }
+ }
+ PriorityQueue<PkVectorSearchResult> exactNearest =
+ new PriorityQueue<>(exactLimit,
BEST_FIRST.reversed());
+ for (CompletableFuture<List<PkVectorSearchResult>>
future :
+ exactFutures) {
+ for (PkVectorSearchResult result :
future.join()) {
+ add(exactNearest, result, exactLimit);
+ }
+ }
+ return new Result(sorted(indexedNearest),
sorted(exactNearest));
+ });
+ }
+
+ private List<PkVectorSearchResult> exactSearch(
+ DataFileMeta file, float[] query, int limit, LongPredicate
excluded) {
+ try (PkVectorReader reader = vectorReaderFactory.create(file)) {
+ return PkVectorExactSearcher.search(
+ file.fileName(), reader, query, metric, limit, excluded);
+ } catch (IOException e) {
+ throw new CompletionException(e);
+ }
+ }
+
+ private static <T> T join(CompletableFuture<T> future) throws IOException {
+ try {
+ return future.join();
+ } catch (CompletionException e) {
+ if (e.getCause() instanceof IOException) {
+ throw (IOException) e.getCause();
}
+ throw e;
}
- return new Result(sorted(indexedNearest), sorted(exactNearest));
}
private static boolean contains(List<Range> ranges, long position) {
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
index 315050bb59..ad1a4a2dca 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
@@ -65,6 +65,8 @@ import java.util.Objects;
import java.util.Optional;
import java.util.PriorityQueue;
import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutorService;
import java.util.stream.Collectors;
@@ -188,21 +190,31 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
protected SearchResult searchBuckets(List<BucketVectorSearchSplit> splits)
{
try {
- SearchContext context = new SearchContext(table);
- List<Candidate> indexedCandidates = new ArrayList<>();
- List<Candidate> exactCandidates = new ArrayList<>();
+ SearchContext context = createSearchContext();
+ List<CompletableFuture<SearchResult>> futures = new
ArrayList<>(splits.size());
+ // Start from the caller because each bucket submits leaf searches
to the same executor.
+ // Wrapping the whole bucket in that executor and waiting inside
it can starve leaf
+ // work.
for (BucketVectorSearchSplit split : splits) {
- SearchResult result = search(split, context);
- indexedCandidates.addAll(result.indexedCandidates());
- exactCandidates.addAll(result.exactCandidates());
+ futures.add(searchAsync(split, context));
}
- return new SearchResult(
- topK(indexedCandidates, indexedLimit),
topK(exactCandidates, limit));
+ CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
+ List<SearchResult> results = new ArrayList<>(futures.size());
+ for (CompletableFuture<SearchResult> future : futures) {
+ results.add(future.join());
+ }
+ return mergeSearchResults(results);
} catch (IOException e) {
throw new RuntimeException("Failed to search primary-key vector
index.", e);
+ } catch (CompletionException e) {
+ throw new RuntimeException("Failed to search primary-key vector
index.", e.getCause());
}
}
+ SearchContext createSearchContext() {
+ return new SearchContext(table);
+ }
+
protected GlobalIndexResult createResult(
PrimaryKeyVectorScan.Plan plan, SearchResult searchResult) {
List<Candidate> indexedCandidates =
topK(searchResult.indexedCandidates(), indexedLimit);
@@ -227,8 +239,8 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
topK(indexedCandidates, indexedLimit), topK(exactCandidates,
limit));
}
- private SearchResult search(BucketVectorSearchSplit split, SearchContext
context)
- throws IOException {
+ CompletableFuture<SearchResult> searchAsync(
+ BucketVectorSearchSplit split, SearchContext context) throws
IOException {
DataSplit dataSplit = split.dataSplit();
List<DataFileMeta> activeFiles =
dataSplit.dataFiles().stream()
@@ -261,18 +273,21 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
searchOptions,
metric,
table.coreOptions().globalIndexSearchMode());
- PrimaryKeyVectorBucketSearch.Result result =
- bucketSearch.search(
+ return bucketSearch
+ .searchAsync(
state,
activeFiles,
deletionVectors,
rowRangesByFile(split),
query,
indexedLimit,
- limit);
- return new SearchResult(
- candidates(dataSplit, result.indexedCandidates()),
- candidates(dataSplit, result.exactCandidates()));
+ limit,
+ context.executor)
+ .thenApply(
+ result ->
+ new SearchResult(
+ candidates(dataSplit,
result.indexedCandidates()),
+ candidates(dataSplit,
result.exactCandidates())));
}
private Map<String, List<Range>> rowRangesByFile(BucketVectorSearchSplit
split)
@@ -621,7 +636,7 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
}
}
- private static class SearchContext {
+ static class SearchContext {
private final FileIO fileIO;
private final IndexFileHandler indexFileHandler;
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
index c4c53d6fe4..a80cb69e0f 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
@@ -38,6 +38,11 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
import static org.assertj.core.api.Assertions.assertThat;
import static
org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
@@ -49,6 +54,108 @@ import static org.mockito.Mockito.when;
/** Tests for {@link PrimaryKeyVectorBucketSearch}. */
class PrimaryKeyVectorBucketSearchTest {
+ @Test
+ void testSearchesUncoveredFilesConcurrently() throws Exception {
+ DataFileMeta data1 = dataFile("data-1");
+ DataFileMeta data2 = dataFile("data-2");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ CountDownLatch started = new CountDownLatch(2);
+ CountDownLatch release = new CountDownLatch(1);
+ PkVectorDataFileReader reader1 = blockingReader(started, release);
+ PkVectorDataFileReader reader2 = blockingReader(started, release);
+ when(readerFactory.create(data1)).thenReturn(reader1);
+ when(readerFactory.create(data2)).thenReturn(reader2);
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ CompletableFuture<PrimaryKeyVectorBucketSearch.Result> future =
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ null,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.FULL)
+ .searchAsync(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Collections.emptyList()),
+ Arrays.asList(data1, data2),
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ new float[] {0, 0},
+ 1,
+ 1,
+ executor);
+
+ try {
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+ } finally {
+ release.countDown();
+ }
+ assertThat(future.get(5,
TimeUnit.SECONDS).exactCandidates()).isEmpty();
+ } finally {
+ release.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void testStartsAnnSegmentsBeforeWaitingForResults() throws Exception {
+ DataFileMeta data1 = dataFile("data-1");
+ DataFileMeta data2 = dataFile("data-2");
+ IndexFileMeta ann1 = segment("ann-1", data1);
+ IndexFileMeta ann2 = segment("ann-2", data2);
+ PkVectorAnnSegmentSearcher annSearcher =
mock(PkVectorAnnSegmentSearcher.class);
+ CompletableFuture<List<PkVectorSearchResult>> future1 = new
CompletableFuture<>();
+ CompletableFuture<List<PkVectorSearchResult>> future2 = new
CompletableFuture<>();
+ CountDownLatch started = new CountDownLatch(2);
+ when(annSearcher.searchAsync(
+ org.mockito.ArgumentMatchers.any(IndexFileMeta.class),
+
org.mockito.ArgumentMatchers.any(PrimaryKeyIndexSourceMeta.class),
+ org.mockito.ArgumentMatchers.any(float[].class),
+ org.mockito.ArgumentMatchers.eq(1),
+
org.mockito.ArgumentMatchers.eq(Collections.emptyMap()),
+ org.mockito.ArgumentMatchers.eq(
+ new
java.util.HashSet<>(Arrays.asList("data-1", "data-2"))),
+
org.mockito.ArgumentMatchers.eq(Collections.emptyMap()),
+
org.mockito.ArgumentMatchers.eq(Collections.emptyMap())))
+ .thenAnswer(
+ invocation -> {
+ started.countDown();
+ return invocation
+ .<IndexFileMeta>getArgument(0)
+ .fileName()
+ .equals("ann-1")
+ ? future1
+ : future2;
+ });
+
+ CompletableFuture<PrimaryKeyVectorBucketSearch.Result> resultFuture =
+ new PrimaryKeyVectorBucketSearch(
+ mock(PkVectorDataFileReader.Factory.class),
+ annSearcher,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.FAST)
+ .searchAsync(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Arrays.asList(ann1, ann2)),
+ Arrays.asList(data1, data2),
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ new float[] {0, 0},
+ 1,
+ 1,
+ Runnable::run);
+
+ assertThat(started.getCount()).isZero();
+ future1.complete(Collections.singletonList(new
PkVectorSearchResult("data-1", 0, 2F)));
+ future2.complete(Collections.singletonList(new
PkVectorSearchResult("data-2", 0, 1F)));
+
+ assertThat(resultFuture.get(5, TimeUnit.SECONDS).indexedCandidates())
+ .extracting(PkVectorSearchResult::dataFileName)
+ .containsExactly("data-2");
+ }
+
@Test
void testFastModeSkipsExactFallback() throws Exception {
DataFileMeta data = dataFile("data");
@@ -112,7 +219,7 @@ class PrimaryKeyVectorBucketSearchTest {
PkVectorAnnSegmentSearcher annSearcher =
mock(PkVectorAnnSegmentSearcher.class);
Map<String, DeletionVector> deletionVectors = Collections.emptyMap();
Map<String, String> searchOptions =
Collections.singletonMap("nprobes", "8");
- when(annSearcher.search(
+ when(annSearcher.searchAsync(
org.mockito.ArgumentMatchers.eq(ann),
org.mockito.ArgumentMatchers.any(PrimaryKeyIndexSourceMeta.class),
org.mockito.ArgumentMatchers.any(float[].class),
@@ -120,8 +227,12 @@ class PrimaryKeyVectorBucketSearchTest {
org.mockito.ArgumentMatchers.eq(deletionVectors),
org.mockito.ArgumentMatchers.eq(
new
java.util.HashSet<>(Arrays.asList("data-1", "data-2"))),
+
org.mockito.ArgumentMatchers.eq(Collections.emptyMap()),
org.mockito.ArgumentMatchers.eq(searchOptions)))
- .thenReturn(Collections.singletonList(new
PkVectorSearchResult("data-1", 1, 0.5F)));
+ .thenReturn(
+ CompletableFuture.completedFuture(
+ Collections.singletonList(
+ new PkVectorSearchResult("data-1", 1,
0.5F))));
PrimaryKeyVectorBucketSearch.Result results =
new PrimaryKeyVectorBucketSearch(
@@ -162,15 +273,19 @@ class PrimaryKeyVectorBucketSearchTest {
IndexFileMeta ann = segment("ann", Arrays.asList(retired, active));
PkVectorAnnSegmentSearcher annSearcher =
mock(PkVectorAnnSegmentSearcher.class);
Map<String, DeletionVector> deletionVectors = Collections.emptyMap();
- when(annSearcher.search(
+ when(annSearcher.searchAsync(
org.mockito.ArgumentMatchers.eq(ann),
org.mockito.ArgumentMatchers.any(PrimaryKeyIndexSourceMeta.class),
org.mockito.ArgumentMatchers.any(float[].class),
org.mockito.ArgumentMatchers.eq(1),
org.mockito.ArgumentMatchers.eq(deletionVectors),
org.mockito.ArgumentMatchers.eq(Collections.singleton("active")),
+
org.mockito.ArgumentMatchers.eq(Collections.emptyMap()),
org.mockito.ArgumentMatchers.eq(Collections.emptyMap())))
- .thenReturn(Collections.singletonList(new
PkVectorSearchResult("active", 0, 1F)));
+ .thenReturn(
+ CompletableFuture.completedFuture(
+ Collections.singletonList(
+ new PkVectorSearchResult("active", 0,
1F))));
PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
List<PkVectorSearchResult> results =
@@ -255,6 +370,24 @@ class PrimaryKeyVectorBucketSearchTest {
return reader;
}
+ private static PkVectorDataFileReader blockingReader(
+ CountDownLatch started, CountDownLatch release) throws Exception {
+ PkVectorDataFileReader reader = mock(PkVectorDataFileReader.class);
+ when(reader.dimension()).thenReturn(2);
+ when(reader.rowCount()).thenReturn(2L);
+
when(reader.readNextVector(org.mockito.ArgumentMatchers.any(float[].class)))
+ .thenAnswer(
+ invocation -> {
+ started.countDown();
+ if (!release.await(5, TimeUnit.SECONDS)) {
+ throw new AssertionError(
+ "Timed out waiting for concurrent
search.");
+ }
+ return false;
+ });
+ return reader;
+ }
+
private static DataFileMeta dataFile(String fileName) {
return DataFileMeta.forAppend(
fileName,
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java
index cf7b447732..0164c38cbd 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java
@@ -18,18 +18,69 @@
package org.apache.paimon.table.source;
+import org.apache.paimon.CoreOptions;
import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.types.FloatType;
+import org.apache.paimon.types.VectorType;
import org.junit.jupiter.api.Test;
import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
/** Tests global candidate merging for primary-key vector search. */
class PrimaryKeyVectorReadTest {
+ @Test
+ void testStartsBucketsBeforeWaitingForResults() throws Exception {
+ BucketVectorSearchSplit split1 = mock(BucketVectorSearchSplit.class);
+ BucketVectorSearchSplit split2 = mock(BucketVectorSearchSplit.class);
+ CompletableFuture<PrimaryKeyVectorRead.SearchResult> future1 = new
CompletableFuture<>();
+ CompletableFuture<PrimaryKeyVectorRead.SearchResult> future2 = new
CompletableFuture<>();
+ CountDownLatch started = new CountDownLatch(2);
+ TestingPrimaryKeyVectorRead read =
+ new TestingPrimaryKeyVectorRead(split1, future1, split2,
future2, started);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try {
+ Future<PrimaryKeyVectorRead.SearchResult> resultFuture =
+ executor.submit(() ->
read.searchBuckets(Arrays.asList(split1, split2)));
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+
+ future1.complete(
+ new PrimaryKeyVectorRead.SearchResult(
+ Collections.singletonList(candidate(0, "data-1",
0, 2F)),
+ Collections.emptyList()));
+ future2.complete(
+ new PrimaryKeyVectorRead.SearchResult(
+ Collections.singletonList(candidate(1, "data-2",
0, 1F)),
+ Collections.emptyList()));
+
+ assertThat(resultFuture.get(5,
TimeUnit.SECONDS).indexedCandidates())
+ .extracting(PrimaryKeyVectorRead.Candidate::dataFileName)
+ .containsExactly("data-2");
+ } finally {
+ future1.completeExceptionally(new RuntimeException("Test
cleanup."));
+ future2.completeExceptionally(new RuntimeException("Test
cleanup."));
+ executor.shutdownNow();
+ }
+ }
+
@Test
void testMergesGlobalTopKWithDeterministicTies() {
List<PrimaryKeyVectorRead.Candidate> candidates =
@@ -55,4 +106,58 @@ class PrimaryKeyVectorReadTest {
return new PrimaryKeyVectorRead.Candidate(
BinaryRow.EMPTY_ROW, bucket, fileName, position, distance);
}
+
+ private static FileStoreTable table() {
+ Map<String, String> options = new HashMap<>();
+ options.put("fields.vector.pk-vector.index.type", "test-vector-ann");
+ options.put("fields.vector.pk-vector.distance.metric", "l2");
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.coreOptions()).thenReturn(new CoreOptions(options));
+ when(table.options()).thenReturn(options);
+ return table;
+ }
+
+ private static class TestingPrimaryKeyVectorRead extends
PrimaryKeyVectorRead {
+
+ private final BucketVectorSearchSplit split1;
+ private final CompletableFuture<SearchResult> future1;
+ private final BucketVectorSearchSplit split2;
+ private final CompletableFuture<SearchResult> future2;
+ private final CountDownLatch started;
+
+ private TestingPrimaryKeyVectorRead(
+ BucketVectorSearchSplit split1,
+ CompletableFuture<SearchResult> future1,
+ BucketVectorSearchSplit split2,
+ CompletableFuture<SearchResult> future2,
+ CountDownLatch started) {
+ super(
+ table(),
+ new DataField(1, "vector", new VectorType(2, new
FloatType())),
+ new float[] {0, 0},
+ 1,
+ Collections.emptyMap());
+ this.split1 = split1;
+ this.future1 = future1;
+ this.split2 = split2;
+ this.future2 = future2;
+ this.started = started;
+ }
+
+ @Override
+ protected SearchContext createSearchContext() {
+ return null;
+ }
+
+ @Override
+ protected CompletableFuture<SearchResult> searchAsync(
+ BucketVectorSearchSplit split, SearchContext context) {
+ started.countDown();
+ if (split == split1) {
+ return future1;
+ }
+ assertThat(split).isSameAs(split2);
+ return future2;
+ }
+ }
}
