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 b7df2b487e [core] Support batch primary-key vector search (#8666)
b7df2b487e is described below
commit b7df2b487edccb6f683165890e580c08332fe5b3
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jul 16 00:29:25 2026 +0800
[core] Support batch primary-key vector search (#8666)
---
.../index/pkvector/PkVectorAnnSegmentSearcher.java | 114 ++++++++++++-
.../index/pkvector/PkVectorExactSearcher.java | 79 ++++++---
.../pkvector/PrimaryKeyVectorBucketSearch.java | 161 ++++++++++++++----
.../table/source/BatchVectorSearchBuilderImpl.java | 17 ++
.../table/source/PrimaryKeyBatchVectorRead.java | 80 +++++++++
.../paimon/table/source/PrimaryKeyVectorRead.java | 59 ++++++-
.../index/pkvector/PkVectorAnnSegmentFileTest.java | 55 +++++++
.../index/pkvector/PkVectorExactSearcherTest.java | 26 +++
.../pkvector/PrimaryKeyVectorBucketSearchTest.java | 33 ++++
.../table/source/PrimaryKeyVectorSearchTest.java | 39 +++++
.../paimon/spark/execution/PaimonStrategy.scala | 179 ++++++++++++++++++++-
.../spark/sql/PrimaryKeyVectorSearchTest.scala | 59 +++++++
12 files changed, 837 insertions(+), 64 deletions(-)
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 788ff56121..75274336a9 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
@@ -31,6 +31,7 @@ import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile;
import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta;
import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.types.DataField;
import org.apache.paimon.utils.IOUtils;
@@ -177,7 +178,7 @@ public class PkVectorAnnSegmentSearcher {
"Vector segment %s has no source metadata.",
segment.fileName());
if (segment.rowCount() == 0) {
- return Collections.emptyList();
+ return CompletableFuture.completedFuture(Collections.emptyList());
}
GlobalIndexer indexer =
GlobalIndexer.create(segment.indexType(), vectorField,
indexOptions);
@@ -231,6 +232,117 @@ public class PkVectorAnnSegmentSearcher {
}
}
+ public List<List<PkVectorSearchResult>> searchBatch(
+ IndexFileMeta segment,
+ PrimaryKeyIndexSourceMeta sourceMeta,
+ float[][] queries,
+ int limit,
+ Map<String, DeletionVector> deletionVectors,
+ Set<String> activeSourceFiles,
+ Map<String, List<Range>> rowRangesByFile,
+ Map<String, String> searchOptions) {
+ return searchBatchAsync(
+ segment,
+ sourceMeta,
+ queries,
+ limit,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ searchOptions)
+ .join();
+ }
+
+ CompletableFuture<List<List<PkVectorSearchResult>>> searchBatchAsync(
+ IndexFileMeta segment,
+ PrimaryKeyIndexSourceMeta sourceMeta,
+ float[][] queries,
+ int limit,
+ Map<String, DeletionVector> deletionVectors,
+ Set<String> activeSourceFiles,
+ Map<String, List<Range>> rowRangesByFile,
+ Map<String, String> searchOptions) {
+ checkArgument(queries != null && queries.length > 0, "Query vectors
cannot be empty.");
+ checkArgument(limit > 0, "Vector search limit must be positive: %s.",
limit);
+ GlobalIndexMeta globalIndexMeta = segment.globalIndexMeta();
+ checkArgument(
+ globalIndexMeta != null && globalIndexMeta.sourceMeta() !=
null,
+ "Vector segment %s has no source metadata.",
+ segment.fileName());
+ if (segment.rowCount() == 0) {
+ List<List<PkVectorSearchResult>> results = new
ArrayList<>(queries.length);
+ for (int i = 0; i < queries.length; i++) {
+ results.add(Collections.emptyList());
+ }
+ return
CompletableFuture.completedFuture(Collections.unmodifiableList(results));
+ }
+ GlobalIndexer indexer =
+ GlobalIndexer.create(segment.indexType(), vectorField,
indexOptions);
+ checkArgument(
+ indexer instanceof VectorGlobalIndexer,
+ "Index algorithm %s does not implement VectorGlobalIndexer.",
+ segment.indexType());
+ String readerMetric =
+ VectorSearchMetric.normalize(((VectorGlobalIndexer)
indexer).metric());
+ checkArgument(
+ metric.equals(readerMetric),
+ "ANN segment metric %s does not match index reader metric %s.",
+ metric,
+ readerMetric);
+
+ GlobalIndexIOMeta ioMeta =
+ new GlobalIndexIOMeta(
+ annSegmentFile.path(segment),
+ segment.fileSize(),
+ globalIndexMeta.indexMeta());
+ GlobalIndexReader reader =
+ indexer.createReader(
+ meta -> fileIO.newInputStream(meta.filePath()),
+ Collections.singletonList(ioMeta),
+ executor);
+ try {
+ BatchVectorSearch search =
+ new BatchVectorSearch(queries, limit, vectorField.name(),
searchOptions);
+ RoaringNavigableMap64 liveRows =
+ liveRowPositions(
+ sourceMeta.sourceFiles(),
+ activeSourceFiles,
+ deletionVectors,
+ rowRangesByFile);
+ if (liveRows != null) {
+ search.withIncludeRowIds(liveRows);
+ }
+ return reader.visitBatchVectorSearch(search)
+ .whenComplete((ignored, error) ->
IOUtils.closeQuietly(reader))
+ .thenApply(
+ scoredResults -> {
+ checkArgument(
+ scoredResults.size() == queries.length,
+ "ANN segment %s returned %s batch
results for %s queries.",
+ segment.fileName(),
+ scoredResults.size(),
+ queries.length);
+ List<List<PkVectorSearchResult>> results =
+ new ArrayList<>(queries.length);
+ for (Optional<ScoredGlobalIndexResult>
scoredResult :
+ scoredResults) {
+ results.add(
+ mapResults(
+ segment,
+ sourceMeta,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ scoredResult));
+ }
+ return Collections.unmodifiableList(results);
+ });
+ } catch (RuntimeException | Error t) {
+ IOUtils.closeQuietly(reader);
+ throw t;
+ }
+ }
+
private List<PkVectorSearchResult> mapResults(
IndexFileMeta segment,
PrimaryKeyIndexSourceMeta sourceMeta,
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorExactSearcher.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorExactSearcher.java
index e63fa700b8..781cde4068 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorExactSearcher.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorExactSearcher.java
@@ -43,44 +43,79 @@ public final class PkVectorExactSearcher {
int limit,
LongPredicate excludedPosition)
throws IOException {
- checkArgument(query.length == reader.dimension(), "Query vector
dimension does not match.");
+ return searchBatch(
+ dataFileName,
+ reader,
+ new float[][] {query},
+ metric,
+ limit,
+ excludedPosition)
+ .get(0);
+ }
+
+ public static List<List<PkVectorSearchResult>> searchBatch(
+ String dataFileName,
+ PkVectorReader reader,
+ float[][] queries,
+ String metric,
+ int limit,
+ LongPredicate excludedPosition)
+ throws IOException {
+ checkArgument(queries != null && queries.length > 0, "Query vectors
cannot be empty.");
checkArgument(limit > 0, "Vector search limit must be positive.");
checkArgument(
VectorSearchMetric.isSupported(metric),
"Unsupported vector distance metric: %s.",
metric);
metric = VectorSearchMetric.normalize(metric);
- for (int i = 0; i < query.length; i++) {
- checkArgument(
- Float.isFinite(query[i]),
- "Query vector element at position %s must be finite.",
- i);
- }
Comparator<PkVectorSearchResult> bestFirst =
Comparator.comparingDouble(PkVectorSearchResult::distance)
.thenComparingLong(PkVectorSearchResult::rowPosition);
- PriorityQueue<PkVectorSearchResult> nearest =
- new PriorityQueue<>(limit, bestFirst.reversed());
+ List<PriorityQueue<PkVectorSearchResult>> nearest = new
ArrayList<>(queries.length);
+ for (float[] query : queries) {
+ validateQuery(query, reader.dimension());
+ nearest.add(new PriorityQueue<>(limit, bestFirst.reversed()));
+ }
+
float[] vector = new float[reader.dimension()];
for (long position = 0; position < reader.rowCount(); position++) {
if (!reader.readNextVector(vector) ||
excludedPosition.test(position)) {
continue;
}
- PkVectorSearchResult candidate =
- new PkVectorSearchResult(
- dataFileName,
- position,
- VectorSearchMetric.computeDistance(query, vector,
metric));
- if (nearest.size() < limit) {
- nearest.add(candidate);
- } else if (bestFirst.compare(candidate, nearest.peek()) < 0) {
- nearest.poll();
- nearest.add(candidate);
+ for (int i = 0; i < queries.length; i++) {
+ PkVectorSearchResult candidate =
+ new PkVectorSearchResult(
+ dataFileName,
+ position,
+ VectorSearchMetric.computeDistance(queries[i],
vector, metric));
+ PriorityQueue<PkVectorSearchResult> queryNearest =
nearest.get(i);
+ if (queryNearest.size() < limit) {
+ queryNearest.add(candidate);
+ } else if (bestFirst.compare(candidate, queryNearest.peek()) <
0) {
+ queryNearest.poll();
+ queryNearest.add(candidate);
+ }
}
}
- List<PkVectorSearchResult> result = new ArrayList<>(nearest);
- Collections.sort(result, bestFirst);
- return Collections.unmodifiableList(result);
+
+ List<List<PkVectorSearchResult>> results = new
ArrayList<>(queries.length);
+ for (PriorityQueue<PkVectorSearchResult> queryNearest : nearest) {
+ List<PkVectorSearchResult> result = new ArrayList<>(queryNearest);
+ Collections.sort(result, bestFirst);
+ results.add(Collections.unmodifiableList(result));
+ }
+ return Collections.unmodifiableList(results);
+ }
+
+ private static void validateQuery(float[] query, int dimension) {
+ checkArgument(query != null, "Query vector cannot be null.");
+ checkArgument(query.length == dimension, "Query vector dimension does
not match.");
+ for (int i = 0; i < query.length; i++) {
+ checkArgument(
+ Float.isFinite(query[i]),
+ "Query vector element at position %s must be finite.",
+ i);
+ }
}
}
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 ed2e2c4114..65dd95362b 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
@@ -139,15 +139,78 @@ public class PrimaryKeyVectorBucketSearch {
int indexedLimit,
int exactLimit,
Executor executor) {
+ return searchBatchAsync(
+ state,
+ activeFiles,
+ deletionVectors,
+ rowRangesByFile,
+ new float[][] {query},
+ indexedLimit,
+ exactLimit,
+ executor)
+ .thenApply(results -> results.get(0));
+ }
+
+ public List<Result> searchBatch(
+ PkVectorBucketIndexState state,
+ List<DataFileMeta> activeFiles,
+ Map<String, DeletionVector> deletionVectors,
+ float[][] queries,
+ int indexedLimit,
+ int exactLimit)
+ throws IOException {
+ return searchBatch(
+ state,
+ activeFiles,
+ deletionVectors,
+ Collections.emptyMap(),
+ queries,
+ indexedLimit,
+ exactLimit);
+ }
+
+ public List<Result> searchBatch(
+ PkVectorBucketIndexState state,
+ List<DataFileMeta> activeFiles,
+ Map<String, DeletionVector> deletionVectors,
+ Map<String, List<Range>> rowRangesByFile,
+ float[][] queries,
+ int indexedLimit,
+ int exactLimit)
+ throws IOException {
+ return join(
+ searchBatchAsync(
+ state,
+ activeFiles,
+ deletionVectors,
+ rowRangesByFile,
+ queries,
+ indexedLimit,
+ exactLimit,
+ Runnable::run));
+ }
+
+ public CompletableFuture<List<Result>> searchBatchAsync(
+ PkVectorBucketIndexState state,
+ List<DataFileMeta> activeFiles,
+ Map<String, DeletionVector> deletionVectors,
+ Map<String, List<Range>> rowRangesByFile,
+ float[][] queries,
+ int indexedLimit,
+ int exactLimit,
+ Executor executor) {
+ checkArgument(queries != null && queries.length > 0, "Query vectors
cannot be empty.");
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.");
}
Set<String> activeSourceFiles = new HashSet<>(filesByName.keySet());
Set<String> covered = new HashSet<>();
- List<CompletableFuture<List<PkVectorSearchResult>>> indexedFutures =
new ArrayList<>();
+ List<CompletableFuture<List<List<PkVectorSearchResult>>>>
indexedFutures =
+ new ArrayList<>();
for (IndexFileMeta ann : state.annSegments()) {
PrimaryKeyIndexSourceMeta sourceMeta =
PrimaryKeyIndexSourceMeta.fromIndexFile(ann);
for (PrimaryKeyIndexSourceFile source : sourceMeta.sourceFiles()) {
@@ -162,19 +225,34 @@ public class PrimaryKeyVectorBucketSearch {
covered.add(source.fileName());
}
checkArgument(annSearcher != null, "ANN search is not
configured.");
- indexedFutures.add(
- annSearcher.searchAsync(
- ann,
- sourceMeta,
- query,
- indexedLimit,
- deletionVectors,
- activeSourceFiles,
- rowRangesByFile,
- searchOptions));
+ if (queries.length == 1) {
+ indexedFutures.add(
+ annSearcher
+ .searchAsync(
+ ann,
+ sourceMeta,
+ queries[0],
+ indexedLimit,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ searchOptions)
+ .thenApply(Collections::singletonList));
+ } else {
+ indexedFutures.add(
+ annSearcher.searchBatchAsync(
+ ann,
+ sourceMeta,
+ queries,
+ indexedLimit,
+ deletionVectors,
+ activeSourceFiles,
+ rowRangesByFile,
+ searchOptions));
+ }
}
- List<CompletableFuture<List<PkVectorSearchResult>>> exactFutures = new
ArrayList<>();
+ List<CompletableFuture<List<List<PkVectorSearchResult>>>> exactFutures
= new ArrayList<>();
if (searchMode != GlobalIndexSearchMode.FAST) {
for (DataFileMeta file : activeFiles) {
if (covered.contains(file.fileName())) {
@@ -191,7 +269,7 @@ public class PrimaryKeyVectorBucketSearch {
|| (rowRanges != null &&
!contains(rowRanges, position));
exactFutures.add(
CompletableFuture.supplyAsync(
- () -> exactSearch(file, query, exactLimit,
excluded), executor));
+ () -> exactSearch(file, queries, exactLimit,
excluded), executor));
}
}
List<CompletableFuture<?>> futures = new ArrayList<>();
@@ -200,36 +278,63 @@ public class PrimaryKeyVectorBucketSearch {
return CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0]))
.thenApply(
ignored -> {
- PriorityQueue<PkVectorSearchResult> indexedNearest
=
- new PriorityQueue<>(indexedLimit,
BEST_FIRST.reversed());
- for (CompletableFuture<List<PkVectorSearchResult>>
future :
+ List<PriorityQueue<PkVectorSearchResult>>
indexedNearest =
+ queues(queries.length, indexedLimit);
+ for
(CompletableFuture<List<List<PkVectorSearchResult>>> future :
indexedFutures) {
- for (PkVectorSearchResult result :
future.join()) {
- add(indexedNearest, result, indexedLimit);
+ List<List<PkVectorSearchResult>> batchResults
= future.join();
+ checkArgument(
+ batchResults.size() == queries.length,
+ "ANN batch result count does not match
query count.");
+ for (int i = 0; i < queries.length; i++) {
+ for (PkVectorSearchResult result :
batchResults.get(i)) {
+ add(indexedNearest.get(i), result,
indexedLimit);
+ }
}
}
- PriorityQueue<PkVectorSearchResult> exactNearest =
- new PriorityQueue<>(exactLimit,
BEST_FIRST.reversed());
- for (CompletableFuture<List<PkVectorSearchResult>>
future :
+ List<PriorityQueue<PkVectorSearchResult>>
exactNearest =
+ queues(queries.length, exactLimit);
+ for
(CompletableFuture<List<List<PkVectorSearchResult>>> future :
exactFutures) {
- for (PkVectorSearchResult result :
future.join()) {
- add(exactNearest, result, exactLimit);
+ List<List<PkVectorSearchResult>> batchResults
= future.join();
+ checkArgument(
+ batchResults.size() == queries.length,
+ "Exact batch result count does not
match query count.");
+ for (int i = 0; i < queries.length; i++) {
+ for (PkVectorSearchResult result :
batchResults.get(i)) {
+ add(exactNearest.get(i), result,
exactLimit);
+ }
}
}
- return new Result(sorted(indexedNearest),
sorted(exactNearest));
+ List<Result> results = new
ArrayList<>(queries.length);
+ for (int i = 0; i < queries.length; i++) {
+ results.add(
+ new Result(
+ sorted(indexedNearest.get(i)),
+ sorted(exactNearest.get(i))));
+ }
+ return Collections.unmodifiableList(results);
});
}
- private List<PkVectorSearchResult> exactSearch(
- DataFileMeta file, float[] query, int limit, LongPredicate
excluded) {
+ private List<List<PkVectorSearchResult>> exactSearch(
+ DataFileMeta file, float[][] queries, int limit, LongPredicate
excluded) {
try (PkVectorReader reader = vectorReaderFactory.create(file)) {
- return PkVectorExactSearcher.search(
- file.fileName(), reader, query, metric, limit, excluded);
+ return PkVectorExactSearcher.searchBatch(
+ file.fileName(), reader, queries, metric, limit, excluded);
} catch (IOException e) {
throw new CompletionException(e);
}
}
+ private static List<PriorityQueue<PkVectorSearchResult>> queues(int count,
int limit) {
+ List<PriorityQueue<PkVectorSearchResult>> queues = new
ArrayList<>(count);
+ for (int i = 0; i < count; i++) {
+ queues.add(new PriorityQueue<>(limit, BEST_FIRST.reversed()));
+ }
+ return queues;
+ }
+
private static <T> T join(CompletableFuture<T> future) throws IOException {
try {
return future.join();
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
index ab92a1e668..8209f93c4c 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
@@ -125,6 +125,14 @@ public class BatchVectorSearchBuilderImpl implements
BatchVectorSearchBuilder {
@Override
public VectorScan newVectorScan() {
+ if (isPrimaryKeyVectorSearch()) {
+ return new PrimaryKeyVectorScan(
+ table,
+ vectorColumn.id(),
+
table.coreOptions().primaryKeyVectorIndexType(vectorColumn.name()),
+ partitionFilter,
+ filter);
+ }
return new DataEvolutionVectorScan(table, partitionFilter, filter,
vectorColumn, options);
}
@@ -137,7 +145,16 @@ public class BatchVectorSearchBuilderImpl implements
BatchVectorSearchBuilder {
for (float[] vector : vectors) {
checkNotNull(vector, "Search vector element cannot be null");
}
+ if (isPrimaryKeyVectorSearch()) {
+ return new PrimaryKeyBatchVectorRead(
+ table, vectorColumn, vectors, limit, options, filter);
+ }
return new BatchVectorReadImpl(
table, partitionFilter, filter, limit, vectorColumn, vectors,
options);
}
+
+ protected boolean isPrimaryKeyVectorSearch() {
+ return vectorColumn != null
+ &&
table.coreOptions().primaryKeyVectorIndexColumns().contains(vectorColumn.name());
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyBatchVectorRead.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyBatchVectorRead.java
new file mode 100644
index 0000000000..bc63614126
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyBatchVectorRead.java
@@ -0,0 +1,80 @@
+/*
+ * 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.
+ */
+
+package org.apache.paimon.table.source;
+
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.types.DataField;
+
+import javax.annotation.Nullable;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
+
+/** Executes snapshot-consistent primary-key vector search for multiple query
vectors. */
+public class PrimaryKeyBatchVectorRead implements BatchVectorRead,
Serializable {
+
+ private static final long serialVersionUID = 1L;
+
+ private final float[][] queries;
+ private final List<PrimaryKeyVectorRead> queryReads;
+
+ public PrimaryKeyBatchVectorRead(
+ FileStoreTable table,
+ DataField vectorField,
+ float[][] queries,
+ int limit,
+ Map<String, String> searchOptions,
+ @Nullable Predicate filter) {
+ checkArgument(queries != null && queries.length > 0, "Query vectors
cannot be empty.");
+ this.queries = new float[queries.length][];
+ this.queryReads = new ArrayList<>(queries.length);
+ for (int i = 0; i < queries.length; i++) {
+ float[] query = checkNotNull(queries[i], "Query vector cannot be
null.").clone();
+ this.queries[i] = query;
+ this.queryReads.add(
+ new PrimaryKeyVectorRead(
+ table, vectorField, query, limit, searchOptions,
filter));
+ }
+ }
+
+ @Override
+ public List<GlobalIndexResult> readBatch(VectorScan.Plan plan) {
+ PrimaryKeyVectorRead firstRead = queryReads.get(0);
+ PrimaryKeyVectorScan.Plan primaryKeyPlan =
firstRead.primaryKeyPlan(plan);
+ List<PrimaryKeyVectorRead.SearchResult> searchResults =
+
firstRead.searchBuckets(firstRead.bucketSplits(primaryKeyPlan), queries);
+ checkArgument(
+ searchResults.size() == queryReads.size(),
+ "Primary-key vector batch result count does not match query
count.");
+
+ List<GlobalIndexResult> results = new ArrayList<>(queryReads.size());
+ for (int i = 0; i < queryReads.size(); i++) {
+ results.add(queryReads.get(i).createResult(primaryKeyPlan,
searchResults.get(i)));
+ }
+ return Collections.unmodifiableList(results);
+ }
+}
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 ad1a4a2dca..4a33127a29 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
@@ -211,6 +211,40 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
}
}
+ protected List<SearchResult> searchBuckets(
+ List<BucketVectorSearchSplit> splits, float[][] queries) {
+ try {
+ SearchContext context = createSearchContext();
+ List<CompletableFuture<List<SearchResult>>> futures = new
ArrayList<>(splits.size());
+ for (BucketVectorSearchSplit split : splits) {
+ futures.add(searchBatchAsync(split, context, queries));
+ }
+ CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
+ List<List<SearchResult>> resultsByQuery = new
ArrayList<>(queries.length);
+ for (int i = 0; i < queries.length; i++) {
+ resultsByQuery.add(new ArrayList<>());
+ }
+ for (CompletableFuture<List<SearchResult>> future : futures) {
+ List<SearchResult> splitResults = future.join();
+ checkArgument(
+ splitResults.size() == queries.length,
+ "Primary-key vector batch result count does not match
query count.");
+ for (int i = 0; i < queries.length; i++) {
+ resultsByQuery.get(i).add(splitResults.get(i));
+ }
+ }
+ List<SearchResult> results = new ArrayList<>(queries.length);
+ for (List<SearchResult> queryResults : resultsByQuery) {
+ results.add(mergeSearchResults(queryResults));
+ }
+ return Collections.unmodifiableList(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);
}
@@ -241,6 +275,13 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
CompletableFuture<SearchResult> searchAsync(
BucketVectorSearchSplit split, SearchContext context) throws
IOException {
+ return searchBatchAsync(split, context, new float[][] {query})
+ .thenApply(results -> results.get(0));
+ }
+
+ CompletableFuture<List<SearchResult>> searchBatchAsync(
+ BucketVectorSearchSplit split, SearchContext context, float[][]
queries)
+ throws IOException {
DataSplit dataSplit = split.dataSplit();
List<DataFileMeta> activeFiles =
dataSplit.dataFiles().stream()
@@ -274,20 +315,26 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
metric,
table.coreOptions().globalIndexSearchMode());
return bucketSearch
- .searchAsync(
+ .searchBatchAsync(
state,
activeFiles,
deletionVectors,
rowRangesByFile(split),
- query,
+ queries,
indexedLimit,
limit,
context.executor)
.thenApply(
- result ->
- new SearchResult(
- candidates(dataSplit,
result.indexedCandidates()),
- candidates(dataSplit,
result.exactCandidates())));
+ bucketResults -> {
+ List<SearchResult> results = new
ArrayList<>(bucketResults.size());
+ for (PrimaryKeyVectorBucketSearch.Result result :
bucketResults) {
+ results.add(
+ new SearchResult(
+ candidates(dataSplit,
result.indexedCandidates()),
+ candidates(dataSplit,
result.exactCandidates())));
+ }
+ return Collections.unmodifiableList(results);
+ });
}
private Map<String, List<Range>> rowRangesByFile(BucketVectorSearchSplit
split)
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
index 8e89ac6fc8..097ff80195 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
@@ -175,6 +175,61 @@ class PkVectorAnnSegmentFileTest {
org.assertj.core.groups.Tuple.tuple("data-1", 1L));
}
+ @Test
+ void testBatchSearchPreservesQueryOrderAndMapsPhysicalPositions() throws
Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = annFile(fileIO);
+ IndexFileMeta segment =
+ annFile.build(
+ Arrays.asList(
+ new PkVectorAnnSegmentFile.Source(
+ dataFile("data-1", 2),
+ new ArrayReader(new float[][] {{5, 0},
{10, 0}})),
+ new PkVectorAnnSegmentFile.Source(
+ dataFile("data-2", 2),
+ new ArrayReader(new float[][] {{0, 0},
{2, 0}}))),
+ vectorField(),
+ indexOptions(),
+ "l2",
+ "test-vector-ann");
+ PrimaryKeyIndexSourceMeta sourceMeta =
PrimaryKeyIndexSourceMeta.fromIndexFile(segment);
+ BitmapDeletionVector data2Deletes = new BitmapDeletionVector();
+ data2Deletes.delete(0);
+ Map<String, org.apache.paimon.deletionvectors.DeletionVector>
deletionVectors =
+ new HashMap<>();
+ deletionVectors.put("data-2", data2Deletes);
+ Map<String, List<Range>> rowRangesByFile = new HashMap<>();
+ rowRangesByFile.put("data-1", Collections.singletonList(new Range(1,
1)));
+ rowRangesByFile.put("data-2", Collections.singletonList(new Range(1,
1)));
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ List<List<PkVectorSearchResult>> results;
+ try {
+ results =
+ new PkVectorAnnSegmentSearcher(
+ fileIO, annFile, vectorField(),
indexOptions(), "l2", executor)
+ .searchBatch(
+ segment,
+ sourceMeta,
+ new float[][] {{0, 0}, {10, 0}},
+ 1,
+ deletionVectors,
+ new HashSet<>(Arrays.asList("data-1",
"data-2")),
+ rowRangesByFile,
+ Collections.emptyMap());
+ } finally {
+ executor.shutdownNow();
+ }
+
+ assertThat(results).hasSize(2);
+ assertThat(results.get(0))
+ .extracting(PkVectorSearchResult::dataFileName,
PkVectorSearchResult::rowPosition)
+ .containsExactly(org.assertj.core.groups.Tuple.tuple("data-2",
1L));
+ assertThat(results.get(1))
+ .extracting(PkVectorSearchResult::dataFileName,
PkVectorSearchResult::rowPosition)
+ .containsExactly(org.assertj.core.groups.Tuple.tuple("data-1",
1L));
+ }
+
@Test
void testSearchFiltersInactiveSources() throws Exception {
LocalFileIO fileIO = LocalFileIO.create();
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorExactSearcherTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorExactSearcherTest.java
index 8f7e14fe0d..866ca7ceab 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorExactSearcherTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorExactSearcherTest.java
@@ -77,6 +77,32 @@ class PkVectorExactSearcherTest {
org.assertj.core.groups.Tuple.tuple("data-file", 0L,
9F));
}
+ @Test
+ void testBatchSearchPreservesQueryOrderAndScansSourceOnce() throws
Exception {
+ ArrayReader reader = new ArrayReader(new float[][] {{0, 0}, {2, 0},
{5, 0}});
+
+ List<List<PkVectorSearchResult>> results;
+ try (PkVectorReader ignored = reader) {
+ results =
+ PkVectorExactSearcher.searchBatch(
+ "data-file",
+ reader,
+ new float[][] {{0, 0}, {5, 0}},
+ "l2",
+ 1,
+ position -> false);
+ }
+
+ assertThat(results).hasSize(2);
+ assertThat(results.get(0))
+ .extracting(PkVectorSearchResult::rowPosition)
+ .containsExactly(0L);
+ assertThat(results.get(1))
+ .extracting(PkVectorSearchResult::rowPosition)
+ .containsExactly(2L);
+ assertThat(reader.position).isEqualTo(3);
+ }
+
private static float distance(String metric) throws IOException {
try (PkVectorReader reader = new ArrayReader(new float[][] {{1, 0}})) {
return PkVectorExactSearcher.search(
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 a80cb69e0f..70c9f83953 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
@@ -351,6 +351,39 @@ class PrimaryKeyVectorBucketSearchTest {
.containsExactly(org.assertj.core.groups.Tuple.tuple("data-1",
1L, 4F));
}
+ @Test
+ void testBatchExactFallbackPreservesQueryOrderAndOpensFileOnce() throws
Exception {
+ DataFileMeta data = dataFile("data");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ PkVectorDataFileReader dataReader = reader(new float[][] {{0, 0}, {5,
0}});
+ when(readerFactory.create(data)).thenReturn(dataReader);
+
+ List<PrimaryKeyVectorBucketSearch.Result> results =
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ null,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.FULL)
+ .searchBatch(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Collections.emptyList()),
+ Collections.singletonList(data),
+ Collections.emptyMap(),
+ new float[][] {{0, 0}, {5, 0}},
+ 1,
+ 1);
+
+ assertThat(results).hasSize(2);
+ assertThat(results.get(0).exactCandidates())
+ .extracting(PkVectorSearchResult::rowPosition)
+ .containsExactly(0L);
+ assertThat(results.get(1).exactCandidates())
+ .extracting(PkVectorSearchResult::rowPosition)
+ .containsExactly(1L);
+ verify(readerFactory).create(data);
+ }
+
private static PkVectorDataFileReader reader(float[][] vectors) throws
IOException {
PkVectorDataFileReader reader = mock(PkVectorDataFileReader.class);
when(reader.dimension()).thenReturn(2);
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
index d1fb7c2556..06d376be26 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
@@ -153,6 +153,30 @@ class PrimaryKeyVectorSearchTest extends TableTestBase {
assertThat(ids).containsExactly(2, 3);
}
+ @Test
+ void testBatchVectorSearchPreservesQueryOrder() throws Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+ write(
+ table,
+ ioManager,
+ GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[]
{0, 0})),
+ GenericRow.of(2, BinaryVector.fromPrimitiveArray(new float[]
{5, 0})),
+ GenericRow.of(3, BinaryVector.fromPrimitiveArray(new float[]
{10, 0})));
+
+ List<GlobalIndexResult> results =
+ table.newBatchVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVectors(new float[][] {{0, 0}, {10, 0}})
+ .withLimit(1)
+ .withOption("refine_factor", "2")
+ .executeBatchLocal();
+
+ assertThat(results).hasSize(2);
+ assertThat(readIds(table, results.get(0))).containsExactly(1);
+ assertThat(readIds(table, results.get(1))).containsExactly(3);
+ }
+
@Test
void testEmptyVectorSearch() throws Exception {
createTableDefault();
@@ -168,6 +192,21 @@ class PrimaryKeyVectorSearchTest extends TableTestBase {
assertThat(((GlobalIndexSplitResult) result).splits()).isEmpty();
}
+ @Test
+ void testBatchBuilderSelectsPrimaryKeyVectorScan() throws Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+
+ VectorScan scan =
+ table.newBatchVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVectors(new float[][] {{0, 0}})
+ .withLimit(1)
+ .newVectorScan();
+
+ assertThat(scan).isInstanceOf(PrimaryKeyVectorScan.class);
+ }
+
@Test
void testVectorSearchUsesSortedIndexPreFilter() throws Exception {
Schema schema =
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
index b3c23e03d6..74c0e6c47d 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
@@ -19,7 +19,8 @@
package org.apache.paimon.spark.execution
import org.apache.paimon.CoreOptions
-import org.apache.paimon.globalindex.{GlobalIndexResult,
ScoredGlobalIndexResult}
+import org.apache.paimon.data.BinaryRow
+import org.apache.paimon.globalindex.{GlobalIndexResult, IndexedSplit,
ScoredGlobalIndexResult}
import org.apache.paimon.partition.PartitionPredicate
import
org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicatesAndDataPredicates
import org.apache.paimon.predicate.{Predicate, PredicateBuilder}
@@ -31,7 +32,7 @@ import org.apache.paimon.spark.data.SparkInternalRow
import org.apache.paimon.spark.read.VectorSearchResultUtils
import org.apache.paimon.spark.schema.PaimonMetadataColumn
import org.apache.paimon.table.{InnerTable, SpecialFields, Table}
-import org.apache.paimon.table.source.{BatchVectorSearchBuilder,
InnerTableScan, ReadBuilder, VectorScan}
+import org.apache.paimon.table.source.{BatchVectorSearchBuilder, DataSplit,
InnerTableScan, PrimaryKeyScoredResult, PrimaryKeySearchPosition,
PrimaryKeyVectorResult, ReadBuilder, VectorScan}
import org.apache.paimon.types.RowType
import org.apache.paimon.utils.RoaringNavigableMap64
@@ -328,6 +329,9 @@ case class LateralVectorSearchExec(
val readBuilder = innerTable
.newReadBuilder()
.withReadType(rowTypeWithRowId.project(readFieldNamesWithRowId.asJava))
+ val physicalReadBuilder = innerTable
+ .newReadBuilder()
+ .withReadType(readRowType)
val scoreMetadataColumns =
if (vectorSearchOutput.exists(_.name ==
PaimonMetadataColumn.SEARCH_SCORE_COLUMN)) {
Seq(PaimonMetadataColumn.SEARCH_SCORE)
@@ -348,7 +352,7 @@ case class LateralVectorSearchExec(
.withVectorColumn(columnName)
.withLimit(limit)
.withOptions(options.asJava)
- pushSearchFilters(readBuilder, vectorSearchBuilder)
+ pushSearchFilters(Seq(readBuilder, physicalReadBuilder),
vectorSearchBuilder)
val vectorPlan = vectorSearchBuilder.newVectorScan().scan()
val batchSize =
@@ -356,6 +360,7 @@ case class LateralVectorSearchExec(
LateralVectorSearchContext(
readBuilder,
+ physicalReadBuilder,
vectorSearchBuilder,
vectorPlan,
scoreMetadataColumns,
@@ -378,7 +383,7 @@ case class LateralVectorSearchExec(
}
private def pushSearchFilters(
- readBuilder: ReadBuilder,
+ readBuilders: Seq[ReadBuilder],
vectorSearchBuilder: BatchVectorSearchBuilder): Unit = {
val predicates = convertSearchFilters()
if (predicates.nonEmpty) {
@@ -388,12 +393,12 @@ case class LateralVectorSearchExec(
innerTable.partitionKeys())
if (split.getLeft.isPresent) {
val partitionFilter = split.getLeft.get()
- readBuilder.withPartitionFilter(partitionFilter)
+ readBuilders.foreach(_.withPartitionFilter(partitionFilter))
vectorSearchBuilder.withPartitionFilter(partitionFilter)
}
if (!split.getRight.isEmpty) {
val dataFilter = PredicateBuilder.and(split.getRight)
- readBuilder.withFilter(dataFilter)
+ readBuilders.foreach(_.withFilter(dataFilter))
vectorSearchBuilder.withFilter(dataFilter)
}
}
@@ -434,9 +439,14 @@ case class LateralVectorSearchExec(
s"Batch vector search returned ${globalIndexResults.size} results for
${queries.size} " +
"query vectors. The result count must match the query count."
)
+ val primaryKeyResults = primaryKeyVectorResults(globalIndexResults)
if (context.metaColumnsOnly) {
- return searchMetaColumns(queries, globalIndexResults, context)
+ return primaryKeyResults match {
+ case Some(results) => searchPrimaryKeyMetaColumns(queries, results,
context)
+ case None => searchMetaColumns(queries, globalIndexResults, context)
+ }
}
+ primaryKeyResults.foreach(results => return searchPrimaryKeyRows(queries,
results, context))
val rowIdToMatches = createRowIdToMatches(queries, globalIndexResults)
val batchGlobalIndexResult =
createBatchGlobalIndexResult(globalIndexResults)
val scan = context.readBuilder
@@ -472,6 +482,126 @@ case class LateralVectorSearchExec(
}
}
+ private def searchPrimaryKeyMetaColumns(
+ queries: Seq[LateralVectorSearchQuery],
+ results: Seq[PrimaryKeyVectorResult],
+ context: LateralVectorSearchContext): Iterator[(InternalRow,
InternalRow)] = {
+ queries.zip(results).iterator.flatMap {
+ case (query, result) =>
+ result.positions().iterator().asScala.map {
+ position =>
+ val values = vectorSearchOutput.map {
+ attr =>
+ attr.name match {
+ case PaimonMetadataColumn.ROW_ID_COLUMN =>
position.rowPosition()
+ case PaimonMetadataColumn.SEARCH_SCORE_COLUMN =>
position.score()
+ case name =>
+ throw new IllegalArgumentException(
+ s"Unsupported primary-key vector search metadata column:
$name")
+ }
+ }.toArray
+ val projectedRow = context.rightProjection(
+ new JoinedRow(
+ query.outerRow,
+ new GenericInternalRow(values.asInstanceOf[Array[Any]])))
+ (query.outerRow, projectedRow)
+ }
+ }
+ }
+
+ private def primaryKeyVectorResults(
+ globalIndexResults: Seq[GlobalIndexResult]):
Option[Seq[PrimaryKeyVectorResult]] = {
+ val results = globalIndexResults.collect { case result:
PrimaryKeyVectorResult => result }
+ require(
+ results.isEmpty || results.size == globalIndexResults.size,
+ "Batch vector search cannot mix primary-key physical results with global
row-ID results."
+ )
+ if (results.isEmpty) None else Some(results)
+ }
+
+ private def searchPrimaryKeyRows(
+ queries: Seq[LateralVectorSearchQuery],
+ results: Seq[PrimaryKeyVectorResult],
+ context: LateralVectorSearchContext): Iterator[(InternalRow,
InternalRow)] = {
+ val snapshotId = results.head.snapshotId()
+ require(
+ results.forall(_.snapshotId() == snapshotId),
+ "Primary-key batch vector results must belong to the same snapshot."
+ )
+
+ val positionToMatches = scala.collection.mutable
+ .LinkedHashMap[LateralVectorSearchPhysicalPosition,
ArrayBuffer[LateralVectorSearchMatch]]()
+ val uniquePositions = scala.collection.mutable
+ .LinkedHashMap[LateralVectorSearchPhysicalPosition,
PrimaryKeySearchPosition]()
+ queries.zip(results).foreach {
+ case (query, result) =>
+ result.positions().asScala.foreach {
+ position =>
+ val key = LateralVectorSearchPhysicalPosition.from(position)
+ positionToMatches.getOrElseUpdate(key, ArrayBuffer()) +=
+ LateralVectorSearchMatch(query.outerRow, position.score())
+ uniquePositions.getOrElseUpdate(key, position)
+ }
+ }
+
+ val sourceSplits =
+ scala.collection.mutable.LinkedHashMap[LateralVectorSearchPhysicalFile,
DataSplit]()
+ results.foreach {
+ result =>
+ result.splits().asScala.foreach {
+ split =>
+ val dataSplit = split.dataSplit()
+ val key = LateralVectorSearchPhysicalFile.from(dataSplit)
+ sourceSplits.getOrElseUpdate(key, dataSplit)
+ }
+ }
+ val batchResult = new PrimaryKeyScoredResult(
+ snapshotId,
+ sourceSplits.values.toList.asJava,
+ uniquePositions.values.toList.asJava)
+ val scan = context.physicalReadBuilder
+ .newScan()
+ .withGlobalIndexResult(batchResult)
+ .asInstanceOf[InnerTableScan]
+ val read = context.physicalReadBuilder.newRead()
+
+ scan.plan().splits().asScala.iterator.flatMap {
+ split =>
+ val indexedSplit = split.asInstanceOf[IndexedSplit]
+ val dataSplit = indexedSplit.dataSplit()
+ val file = LateralVectorSearchPhysicalFile.from(dataSplit)
+ val reader =
+ PaimonRecordReaderIterator(
+ read.createReader(split),
+ Seq(PaimonMetadataColumn.ROW_ID) ++ context.scoreMetadataColumns,
+ split)
+ val readerState = context.readerTracker.track(reader)
+ new Iterator[Iterator[(InternalRow, InternalRow)]] {
+ override def hasNext: Boolean = {
+ val hasNext = reader.hasNext
+ if (!hasNext) {
+ readerState.closeOnce()
+ }
+ hasNext
+ }
+
+ override def next(): Iterator[(InternalRow, InternalRow)] = {
+ val rightRow = context.sparkRow.replace(reader.next())
+ val position = LateralVectorSearchPhysicalPosition(
+ file.partition,
+ file.bucket,
+ file.dataFileName,
+ rightRow.getLong(context.rowIdOrdinal))
+ positionToMatches.getOrElse(position, Seq.empty).iterator.map {
+ searchMatch =>
+ val projectedRow = projectRightRow(rightRow, searchMatch,
context)
+ (searchMatch.outerRow, projectedRow)
+ }
+ }
+ }.flatMap(identity)
+ }
+ }
+
private def searchMetaColumns(
queries: Seq[LateralVectorSearchQuery],
globalIndexResults: Seq[GlobalIndexResult],
@@ -599,6 +729,7 @@ case class LateralVectorSearchExec(
private case class LateralVectorSearchContext(
readBuilder: ReadBuilder,
+ physicalReadBuilder: ReadBuilder,
vectorSearchBuilder: BatchVectorSearchBuilder,
vectorPlan: VectorScan.Plan,
scoreMetadataColumns: Seq[PaimonMetadataColumn],
@@ -613,4 +744,38 @@ case class LateralVectorSearchExec(
private case class LateralVectorSearchQuery(outerRow: InternalRow,
queryVector: Array[Float])
private case class LateralVectorSearchMatch(outerRow: InternalRow, score:
Float)
+
+ private case class LateralVectorSearchPhysicalFile(
+ partition: BinaryRow,
+ bucket: Int,
+ dataFileName: String)
+
+ private object LateralVectorSearchPhysicalFile {
+ def from(split: DataSplit): LateralVectorSearchPhysicalFile = {
+ require(
+ split.dataFiles().size() == 1,
+ "Primary-key indexed split must contain exactly one data file."
+ )
+ LateralVectorSearchPhysicalFile(
+ split.partition().copy(),
+ split.bucket(),
+ split.dataFiles().get(0).fileName())
+ }
+ }
+
+ private case class LateralVectorSearchPhysicalPosition(
+ partition: BinaryRow,
+ bucket: Int,
+ dataFileName: String,
+ rowPosition: Long)
+
+ private object LateralVectorSearchPhysicalPosition {
+ def from(position: PrimaryKeySearchPosition):
LateralVectorSearchPhysicalPosition = {
+ LateralVectorSearchPhysicalPosition(
+ position.partition().copy(),
+ position.bucket(),
+ position.dataFileName(),
+ position.rowPosition())
+ }
+ }
}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
index b9b71c9045..ebafe40933 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
@@ -398,6 +398,65 @@ class PrimaryKeyVectorSearchTest extends
PaimonSparkTestBase {
}
}
+ test("lateral primary-key vector search reads physical positions") {
+ withTable("T") {
+ createVectorTable(columns = "id INT, embedding ARRAY<FLOAT>,
query_embedding ARRAY<FLOAT>")
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(1.0f, 0.0f), array(0.0f, 0.0f)),
+ | (2, array(5.0f, 0.0f), array(0.5f, 0.0f))
+ |""".stripMargin)
+
+ val rows = spark
+ .sql("""
+ |SELECT q.id AS query_id, r.id AS result_id, r._row_id AS
row_id,
+ | r.__paimon_search_score AS score
+ |FROM T AS q,
+ |LATERAL (
+ | SELECT id, _row_id, __paimon_search_score
+ | FROM vector_search('T', 'embedding', q.query_embedding, 1)
+ |) AS r
+ |ORDER BY query_id
+ |""".stripMargin)
+ .collect()
+
+ assert(rows.length == 2)
+ assert(rows.map(_.getInt(1)).toSeq == Seq(1, 1))
+ assert(rows.map(_.getLong(2)).distinct.length == 1)
+ assert(Math.abs(rows(0).getFloat(3) - 0.5f) < 1e-6)
+ assert(Math.abs(rows(1).getFloat(3) - 0.8f) < 1e-6)
+ }
+ }
+
+ test("lateral primary-key vector search projects physical metadata") {
+ withTable("T") {
+ createVectorTable(columns = "id INT, embedding ARRAY<FLOAT>,
query_embedding ARRAY<FLOAT>")
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(1.0f, 0.0f), array(0.0f, 0.0f)),
+ | (2, array(5.0f, 0.0f), array(0.5f, 0.0f))
+ |""".stripMargin)
+
+ val rows = spark
+ .sql("""
+ |SELECT q.id AS query_id, r._row_id AS row_id,
+ | r.__paimon_search_score AS score
+ |FROM T AS q,
+ |LATERAL (
+ | SELECT _row_id, __paimon_search_score
+ | FROM vector_search('T', 'embedding', q.query_embedding, 1)
+ |) AS r
+ |ORDER BY query_id
+ |""".stripMargin)
+ .collect()
+
+ assert(rows.length == 2)
+ assert(rows.map(_.getLong(1)).distinct.length == 1)
+ assert(Math.abs(rows(0).getFloat(2) - 0.5f) < 1e-6)
+ assert(Math.abs(rows(1).getFloat(2) - 0.8f) < 1e-6)
+ }
+ }
+
private def createVectorTable(
columns: String = "id INT, embedding ARRAY<FLOAT>",
primaryKey: String = "id",
