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 7e396cff87 Add rerank to primary-key vector search (#8591)
7e396cff87 is described below
commit 7e396cff87033bb61c9cfe55e47daeb5b850ab24
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon Jul 13 14:44:57 2026 +0800
Add rerank to primary-key vector search (#8591)
Add exact reranking to primary-key vector search so approximate ANN
candidates can be reordered
with the original vectors stored in Paimon. The implementation reuses
the Data Evolution refine
configuration semantics and supports both local and distributed Spark
execution.
---
docs/docs/primary-key-table/vector-index.md | 32 +++
.../pkvector/PrimaryKeyVectorBucketSearch.java | 63 +++++-
.../paimon/table/source/AbstractVectorRead.java | 84 +-------
.../paimon/table/source/PrimaryKeyVectorRead.java | 228 ++++++++++++++++++++-
.../table/source/VectorSearchRefineOptions.java | 121 +++++++++++
.../pkvector/PrimaryKeyVectorBucketSearchTest.java | 17 +-
.../table/source/PrimaryKeyVectorSearchTest.java | 94 +++++++--
.../source/VectorSearchRefineOptionsTest.java | 112 ++++++++++
.../spark/read/SparkPrimaryKeyVectorRead.java | 12 +-
.../spark/sql/PrimaryKeyVectorSearchTest.scala | 41 ++++
10 files changed, 680 insertions(+), 124 deletions(-)
diff --git a/docs/docs/primary-key-table/vector-index.md
b/docs/docs/primary-key-table/vector-index.md
index 9604df843e..e6adaa55f2 100644
--- a/docs/docs/primary-key-table/vector-index.md
+++ b/docs/docs/primary-key-table/vector-index.md
@@ -140,6 +140,38 @@ required. Batch writes which wait for compaction can
publish the data and its in
## Search
+### Exact Rerank
+
+Primary-key vector search can retrieve more ANN candidates and rerank them
with the original
+vectors stored in the table. For example, the following table option retrieves
up to four times
+the requested Top-K from the `ivf-flat` index before computing exact distances:
+
+```sql
+'fields.embedding.ivf-flat.refine_factor' = '4'
+```
+
+The option is disabled by default. Its configuration semantics are the same as
for a Data
+Evolution vector index:
+
+- `refine_factor`, `refine-factor`, `rerank_factor`, and `rerank-factor` are
accepted.
+- A query option overrides every table option. Within either set of options,
field and index
+ prefixes take precedence over less specific prefixes. For example,
+ `fields.embedding.ivf-flat.refine_factor` takes precedence over
+ `fields.embedding.ivf.refine_factor`, which takes precedence over
`ivf.refine_factor` and then
+ `refine_factor`. The normalized underscore form of an index type, such as
`ivf_flat`, is also
+ accepted after the configured index name.
+- The factor must be a positive integer. A factor of `1` performs exact
reranking without
+ retrieving additional ANN candidates.
+
+Only candidates returned by ANN 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. Files without
+an active ANN segment are already searched exactly and are kept separate from
approximate
+candidates until the final Top-K merge.
+
+For distributed Spark searches, executors return bounded ANN and exact
candidate streams. The
+driver globally merges the ANN candidates and rereads their original vectors
for exact reranking;
+this does not start a second Spark job.
+
### Spark SQL
Use the `vector_search` table-valued function. Spark exposes the ANN score
through the
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 8c2a77a135..62bb835b8f 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
@@ -73,13 +73,36 @@ public class PrimaryKeyVectorBucketSearch {
float[] query,
int limit)
throws IOException {
- checkArgument(limit > 0, "Vector search limit must be positive.");
+ Result result = search(state, activeFiles, deletionVectors, query,
limit, limit);
+ PriorityQueue<PkVectorSearchResult> nearest =
+ new PriorityQueue<>(limit, BEST_FIRST.reversed());
+ for (PkVectorSearchResult candidate : result.indexedCandidates) {
+ add(nearest, candidate, limit);
+ }
+ for (PkVectorSearchResult candidate : result.exactCandidates) {
+ add(nearest, candidate, limit);
+ }
+ return sorted(nearest);
+ }
+
+ public Result search(
+ PkVectorBucketIndexState state,
+ List<DataFileMeta> activeFiles,
+ Map<String, DeletionVector> deletionVectors,
+ float[] query,
+ int indexedLimit,
+ int exactLimit)
+ throws IOException {
+ 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> nearest =
- new PriorityQueue<>(limit, BEST_FIRST.reversed());
+ 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<>();
for (IndexFileMeta ann : state.annSegments()) {
@@ -101,11 +124,11 @@ public class PrimaryKeyVectorBucketSearch {
ann,
sourceMeta,
query,
- limit,
+ indexedLimit,
deletionVectors,
activeSourceFiles,
searchOptions)) {
- add(nearest, result, limit);
+ add(indexedNearest, result, indexedLimit);
}
}
@@ -119,12 +142,16 @@ public class PrimaryKeyVectorBucketSearch {
try (PkVectorReader reader = vectorReaderFactory.create(file))
{
for (PkVectorSearchResult result :
PkVectorExactSearcher.search(
- file.fileName(), reader, query, metric,
limit, excluded)) {
- add(nearest, result, limit);
+ file.fileName(), reader, query, metric,
exactLimit, excluded)) {
+ add(exactNearest, result, exactLimit);
}
}
}
}
+ return new Result(sorted(indexedNearest), sorted(exactNearest));
+ }
+
+ private static List<PkVectorSearchResult>
sorted(PriorityQueue<PkVectorSearchResult> nearest) {
List<PkVectorSearchResult> result = new ArrayList<>(nearest);
Collections.sort(result, BEST_FIRST);
return Collections.unmodifiableList(result);
@@ -141,4 +168,26 @@ public class PrimaryKeyVectorBucketSearch {
nearest.add(candidate);
}
}
+
+ /** Separately bounded approximate-index and exact-fallback candidates. */
+ public static class Result {
+
+ private final List<PkVectorSearchResult> indexedCandidates;
+ private final List<PkVectorSearchResult> exactCandidates;
+
+ private Result(
+ List<PkVectorSearchResult> indexedCandidates,
+ List<PkVectorSearchResult> exactCandidates) {
+ this.indexedCandidates = indexedCandidates;
+ this.exactCandidates = exactCandidates;
+ }
+
+ public List<PkVectorSearchResult> indexedCandidates() {
+ return indexedCandidates;
+ }
+
+ public List<PkVectorSearchResult> exactCandidates() {
+ return exactCandidates;
+ }
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
index dd43a99af8..2789799bd0 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
@@ -301,17 +301,7 @@ public abstract class AbstractVectorRead implements
Serializable {
}
protected int indexedSearchLimit(String indexType) {
- int refineFactor = configuredRefineFactor(indexType);
- if (refineFactor == 0) {
- return limit;
- }
- if (limit > Integer.MAX_VALUE / refineFactor) {
- throw new IllegalArgumentException(
- String.format(
- "Vector search limit overflow: limit=%d, refine
factor=%d",
- limit, refineFactor));
- }
- return limit * refineFactor;
+ return VectorSearchRefineOptions.searchLimit(limit,
configuredRefineFactor(indexType));
}
protected ScoredGlobalIndexResult maybeRerankIndexedResult(
@@ -721,73 +711,11 @@ public abstract class AbstractVectorRead implements
Serializable {
}
protected int configuredRefineFactor(String indexType) {
- String value = configuredRefineFactor(options, indexType);
- if (value == null) {
- value = configuredRefineFactor(table.options(), indexType);
- }
- if (value == null) {
- return 0;
- }
- try {
- int factor = Integer.parseInt(value);
- if (factor <= 0) {
- throw new IllegalArgumentException(
- "Vector refine factor must be positive, got: " +
value);
- }
- return factor;
- } catch (NumberFormatException e) {
- throw new IllegalArgumentException(
- "Invalid vector refine factor: " + value + ". Must be an
integer.", e);
- }
- }
-
- @Nullable
- private String configuredRefineFactor(Map<String, String> options, String
indexType) {
- List<String> prefixes = new ArrayList<>();
- String fieldPrefix = "fields." + vectorColumn.name() + ".";
- addRefinePrefixes(prefixes, fieldPrefix, indexType);
- addRefinePrefixes(prefixes, "", indexType);
-
- for (String prefix : prefixes) {
- String value = refineFactorOption(options, prefix +
"refine_factor");
- if (value == null) {
- value = refineFactorOption(options, prefix + "refine-factor");
- }
- if (value == null) {
- value = refineFactorOption(options, prefix + "rerank_factor");
- }
- if (value == null) {
- value = refineFactorOption(options, prefix + "rerank-factor");
- }
- if (value != null) {
- return value;
- }
- }
- return null;
- }
-
- private static void addRefinePrefixes(List<String> prefixes, String base,
String indexType) {
- if (indexType != null && !indexType.isEmpty()) {
- prefixes.add(base + indexType + ".");
- String normalizedIndexType = normalizeIndexType(indexType);
- if (!normalizedIndexType.equals(indexType)) {
- prefixes.add(base + normalizedIndexType + ".");
- }
- if (normalizedIndexType.startsWith("ivf")) {
- prefixes.add(base + "ivf.");
- }
- }
- prefixes.add(base);
- }
-
- @Nullable
- private static String refineFactorOption(Map<String, String> options,
String key) {
- String value = options.get(key);
- return value == null ? null : value.trim();
- }
-
- private static String normalizeIndexType(String indexType) {
- return indexType.toLowerCase().replace('-', '_');
+ return VectorSearchRefineOptions.resolve(
+ options,
+ table == null ? Collections.emptyMap() : table.options(),
+ vectorColumn.name(),
+ indexType);
}
private static IndexFileMeta
firstVectorIndexFile(List<IndexVectorSearchSplit> splits) {
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 5f1868f0ec..785a93c5f4 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
@@ -20,10 +20,13 @@ package org.apache.paimon.table.source;
import org.apache.paimon.KeyValueFileStore;
import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.InternalRow;
import org.apache.paimon.deletionvectors.DeletionVector;
import org.apache.paimon.fs.FileIO;
import org.apache.paimon.globalindex.GlobalIndexReadThreadPool;
import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.globalindex.IndexedSplit;
+import org.apache.paimon.globalindex.VectorSearchMetric;
import org.apache.paimon.index.IndexFileHandler;
import org.apache.paimon.index.pkvector.PkVectorAnnSegmentSearcher;
import org.apache.paimon.index.pkvector.PkVectorBucketIndexState;
@@ -34,9 +37,13 @@ import
org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.io.KeyValueFileReaderFactory;
import org.apache.paimon.options.Options;
+import org.apache.paimon.reader.RecordReader;
+import org.apache.paimon.reader.ScoreRecordIterator;
+import org.apache.paimon.reader.ScoreRecordReader;
import org.apache.paimon.table.DelegatedFileStoreTable;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.types.DataField;
+import org.apache.paimon.types.RowType;
import org.apache.paimon.types.VectorType;
import java.io.IOException;
@@ -47,6 +54,7 @@ import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.PriorityQueue;
import java.util.concurrent.ExecutorService;
@@ -87,6 +95,8 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
private final float[] query;
protected final int limit;
private final String metric;
+ private final int refineFactor;
+ private final int indexedLimit;
public PrimaryKeyVectorRead(
FileStoreTable table,
@@ -111,6 +121,10 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
this.limit = limit;
this.metric =
normalize(table.coreOptions().primaryKeyVectorDistanceMetric(vectorField.name()));
+ this.refineFactor =
+ VectorSearchRefineOptions.resolve(
+ this.searchOptions, table.options(),
vectorField.name(), indexType);
+ this.indexedLimit = VectorSearchRefineOptions.searchLimit(limit,
refineFactor);
}
private static KeyValueFileStore keyValueStore(FileStoreTable table) {
@@ -152,25 +166,48 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
return splits;
}
- protected List<Candidate> searchBuckets(List<BucketVectorSearchSplit>
splits) {
+ protected SearchResult searchBuckets(List<BucketVectorSearchSplit> splits)
{
try {
SearchContext context = new SearchContext(table);
- List<Candidate> candidates = new ArrayList<>();
+ List<Candidate> indexedCandidates = new ArrayList<>();
+ List<Candidate> exactCandidates = new ArrayList<>();
for (BucketVectorSearchSplit split : splits) {
- candidates.addAll(search(split, context));
+ SearchResult result = search(split, context);
+ indexedCandidates.addAll(result.indexedCandidates());
+ exactCandidates.addAll(result.exactCandidates());
}
- return topK(candidates, limit);
+ return new SearchResult(
+ topK(indexedCandidates, indexedLimit),
topK(exactCandidates, limit));
} catch (IOException e) {
throw new RuntimeException("Failed to search primary-key vector
index.", e);
}
}
protected GlobalIndexResult createResult(
- PrimaryKeyVectorScan.Plan plan, List<Candidate> candidates) {
+ PrimaryKeyVectorScan.Plan plan, SearchResult searchResult) {
+ List<Candidate> indexedCandidates =
topK(searchResult.indexedCandidates(), indexedLimit);
+ if (refineFactor > 0 && !indexedCandidates.isEmpty()) {
+ indexedCandidates = rerank(plan, indexedCandidates);
+ } else {
+ indexedCandidates = topK(indexedCandidates, limit);
+ }
+ List<Candidate> candidates = new ArrayList<>(indexedCandidates);
+ candidates.addAll(topK(searchResult.exactCandidates(), limit));
return new PrimaryKeyVectorResult(plan, topK(candidates, limit),
metric);
}
- private List<Candidate> search(BucketVectorSearchSplit split,
SearchContext context)
+ protected SearchResult mergeSearchResults(List<SearchResult> results) {
+ List<Candidate> indexedCandidates = new ArrayList<>();
+ List<Candidate> exactCandidates = new ArrayList<>();
+ for (SearchResult result : results) {
+ indexedCandidates.addAll(result.indexedCandidates());
+ exactCandidates.addAll(result.exactCandidates());
+ }
+ return new SearchResult(
+ topK(indexedCandidates, indexedLimit), topK(exactCandidates,
limit));
+ }
+
+ private SearchResult search(BucketVectorSearchSplit split, SearchContext
context)
throws IOException {
DataSplit dataSplit = split.dataSplit();
List<DataFileMeta> activeFiles =
@@ -204,13 +241,22 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
searchOptions,
metric,
table.coreOptions().globalIndexSearchMode());
- List<Candidate> candidates = new ArrayList<>();
- for (PkVectorSearchResult result :
- bucketSearch.search(state, activeFiles, deletionVectors,
query, limit)) {
+ PrimaryKeyVectorBucketSearch.Result result =
+ bucketSearch.search(
+ state, activeFiles, deletionVectors, query,
indexedLimit, limit);
+ return new SearchResult(
+ candidates(dataSplit, result.indexedCandidates()),
+ candidates(dataSplit, result.exactCandidates()));
+ }
+
+ private static List<Candidate> candidates(
+ DataSplit split, List<PkVectorSearchResult> searchResults) {
+ List<Candidate> candidates = new ArrayList<>(searchResults.size());
+ for (PkVectorSearchResult result : searchResults) {
candidates.add(
new Candidate(
- dataSplit.partition(),
- dataSplit.bucket(),
+ split.partition(),
+ split.bucket(),
result.dataFileName(),
result.rowPosition(),
result.distance()));
@@ -218,6 +264,94 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
return candidates;
}
+ private List<Candidate> rerank(
+ PrimaryKeyVectorScan.Plan plan, List<Candidate> indexedCandidates)
{
+ Map<PhysicalPosition, Candidate> candidatesByPosition = new
HashMap<>();
+ for (Candidate candidate : indexedCandidates) {
+ PhysicalPosition position = new PhysicalPosition(candidate);
+ checkArgument(
+ candidatesByPosition.put(position, candidate) == null,
+ "Duplicate primary-key vector candidate %s.",
+ position);
+ }
+
+ List<IndexedSplit> splits =
+ new PrimaryKeyVectorResult(plan, indexedCandidates,
metric).splits();
+ TableRead read =
table.newReadBuilder().withReadType(RowType.of(vectorField)).newRead();
+ List<Candidate> reranked = new ArrayList<>(indexedCandidates.size());
+ try {
+ for (IndexedSplit split : splits) {
+ rerankSplit(read, split, candidatesByPosition, reranked);
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to rerank primary-key vector
candidates.", e);
+ }
+ checkArgument(
+ candidatesByPosition.isEmpty(),
+ "Failed to read %s primary-key vector candidates for
reranking.",
+ candidatesByPosition.size());
+ return topK(reranked, limit);
+ }
+
+ @SuppressWarnings("unchecked")
+ private void rerankSplit(
+ TableRead read,
+ IndexedSplit split,
+ Map<PhysicalPosition, Candidate> candidatesByPosition,
+ List<Candidate> reranked)
+ throws IOException {
+ DataSplit dataSplit = split.dataSplit();
+ checkArgument(
+ dataSplit.dataFiles().size() == 1,
+ "Primary-key vector rerank split must contain exactly one data
file.");
+ String dataFileName = dataSplit.dataFiles().get(0).fileName();
+ try (RecordReader<InternalRow> reader = read.createReader(split)) {
+ checkArgument(
+ reader instanceof ScoreRecordReader,
+ "Primary-key vector rerank requires a score record
reader.");
+ ScoreRecordReader<InternalRow> scoreReader =
(ScoreRecordReader<InternalRow>) reader;
+ ScoreRecordIterator<InternalRow> batch;
+ while ((batch = scoreReader.readBatch()) != null) {
+ try {
+ InternalRow row;
+ while ((row = batch.next()) != null) {
+ PhysicalPosition position =
+ new PhysicalPosition(
+ dataSplit.partition(),
+ dataSplit.bucket(),
+ dataFileName,
+ batch.returnedRowId());
+ Candidate candidate =
candidatesByPosition.remove(position);
+ checkArgument(
+ candidate != null,
+ "Primary-key vector rerank read unexpected
position %s.",
+ position);
+ checkArgument(
+ !row.isNullAt(0),
+ "Primary-key vector candidate %s contains a
null vector.",
+ position);
+ float[] vector = row.getVector(0).toFloatArray();
+ checkArgument(
+ vector.length == query.length,
+ "Primary-key vector candidate %s has dimension
%s instead of %s.",
+ position,
+ vector.length,
+ query.length);
+ reranked.add(
+ new Candidate(
+ candidate.partition(),
+ candidate.bucket(),
+ candidate.dataFileName(),
+ candidate.rowPosition(),
+
VectorSearchMetric.computeDistance(query, vector, metric)));
+ }
+ } finally {
+ batch.releaseBatch();
+ }
+ }
+ }
+ }
+
private Map<String, DeletionVector> deletionVectors(DataSplit split,
FileIO fileIO)
throws IOException {
DeletionVector.Factory factory =
@@ -305,6 +439,78 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
}
}
+ /** Separately bounded approximate-index and exact-fallback candidates. */
+ public static class SearchResult implements Serializable {
+
+ private static final long serialVersionUID = 1L;
+
+ private final List<Candidate> indexedCandidates;
+ private final List<Candidate> exactCandidates;
+
+ public SearchResult(List<Candidate> indexedCandidates, List<Candidate>
exactCandidates) {
+ this.indexedCandidates =
+ Collections.unmodifiableList(new
ArrayList<>(indexedCandidates));
+ this.exactCandidates = Collections.unmodifiableList(new
ArrayList<>(exactCandidates));
+ }
+
+ public List<Candidate> indexedCandidates() {
+ return indexedCandidates;
+ }
+
+ public List<Candidate> exactCandidates() {
+ return exactCandidates;
+ }
+ }
+
+ private static class PhysicalPosition {
+
+ private final BinaryRow partition;
+ private final int bucket;
+ private final String dataFileName;
+ private final long rowPosition;
+
+ private PhysicalPosition(Candidate candidate) {
+ this(
+ candidate.partition(),
+ candidate.bucket(),
+ candidate.dataFileName(),
+ candidate.rowPosition());
+ }
+
+ private PhysicalPosition(
+ BinaryRow partition, int bucket, String dataFileName, long
rowPosition) {
+ this.partition = partition.copy();
+ this.bucket = bucket;
+ this.dataFileName = dataFileName;
+ this.rowPosition = rowPosition;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ PhysicalPosition that = (PhysicalPosition) o;
+ return bucket == that.bucket
+ && rowPosition == that.rowPosition
+ && partition.equals(that.partition)
+ && dataFileName.equals(that.dataFileName);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(partition, bucket, dataFileName, rowPosition);
+ }
+
+ @Override
+ public String toString() {
+ return dataFileName + '@' + rowPosition + " in bucket " + bucket;
+ }
+ }
+
private static class SearchContext {
private final FileIO fileIO;
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchRefineOptions.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchRefineOptions.java
new file mode 100644
index 0000000000..a27e100c58
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchRefineOptions.java
@@ -0,0 +1,121 @@
+/*
+ * 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 javax.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+
+/** Shared refine/rerank option resolution for vector search reads. */
+final class VectorSearchRefineOptions {
+
+ private VectorSearchRefineOptions() {}
+
+ static int resolve(
+ Map<String, String> queryOptions,
+ Map<String, String> tableOptions,
+ String vectorColumn,
+ @Nullable String indexType) {
+ String value = resolve(queryOptions, vectorColumn, indexType);
+ if (value == null) {
+ value = resolve(tableOptions, vectorColumn, indexType);
+ }
+ if (value == null) {
+ return 0;
+ }
+ try {
+ int factor = Integer.parseInt(value);
+ if (factor <= 0) {
+ throw new IllegalArgumentException(
+ "Vector refine factor must be positive, got: " +
value);
+ }
+ return factor;
+ } catch (NumberFormatException e) {
+ throw new IllegalArgumentException(
+ "Invalid vector refine factor: " + value + ". Must be an
integer.", e);
+ }
+ }
+
+ static int searchLimit(int limit, int refineFactor) {
+ if (refineFactor == 0) {
+ return limit;
+ }
+ if (limit > Integer.MAX_VALUE / refineFactor) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Vector search limit overflow: limit=%d, refine
factor=%d",
+ limit, refineFactor));
+ }
+ return limit * refineFactor;
+ }
+
+ @Nullable
+ private static String resolve(
+ Map<String, String> options, String vectorColumn, @Nullable String
indexType) {
+ List<String> prefixes = new ArrayList<>();
+ String fieldPrefix = "fields." + vectorColumn + ".";
+ addPrefixes(prefixes, fieldPrefix, indexType);
+ addPrefixes(prefixes, "", indexType);
+
+ for (String prefix : prefixes) {
+ String value = option(options, prefix + "refine_factor");
+ if (value == null) {
+ value = option(options, prefix + "refine-factor");
+ }
+ if (value == null) {
+ value = option(options, prefix + "rerank_factor");
+ }
+ if (value == null) {
+ value = option(options, prefix + "rerank-factor");
+ }
+ if (value != null) {
+ return value;
+ }
+ }
+ return null;
+ }
+
+ private static void addPrefixes(
+ List<String> prefixes, String base, @Nullable String indexType) {
+ if (indexType != null && !indexType.isEmpty()) {
+ prefixes.add(base + indexType + ".");
+ String normalizedIndexType = normalizeIndexType(indexType);
+ if (!normalizedIndexType.equals(indexType)) {
+ prefixes.add(base + normalizedIndexType + ".");
+ }
+ if (normalizedIndexType.startsWith("ivf")) {
+ prefixes.add(base + "ivf.");
+ }
+ }
+ prefixes.add(base);
+ }
+
+ @Nullable
+ private static String option(Map<String, String> options, String key) {
+ String value = options.get(key);
+ return value == null ? null : value.trim();
+ }
+
+ private static String normalizeIndexType(String indexType) {
+ return indexType.toLowerCase(Locale.ROOT).replace('-', '_');
+ }
+}
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 fbc40071f0..fc52293be7 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
@@ -95,7 +95,7 @@ class PrimaryKeyVectorBucketSearchTest {
}
@Test
- void testMergesAnnAndExactFallbackWithoutRescanningCoveredFiles() throws
Exception {
+ void testSeparatesAnnAndExactFallbackWithoutRescanningCoveredFiles()
throws Exception {
DataFileMeta data1 = dataFile("data-1");
DataFileMeta data2 = dataFile("data-2");
IndexFileMeta ann = segment("ann", data1);
@@ -120,7 +120,7 @@ class PrimaryKeyVectorBucketSearchTest {
org.mockito.ArgumentMatchers.eq(searchOptions)))
.thenReturn(Collections.singletonList(new
PkVectorSearchResult("data-1", 1, 0.5F)));
- List<PkVectorSearchResult> results =
+ PrimaryKeyVectorBucketSearch.Result results =
new PrimaryKeyVectorBucketSearch(
readerFactory,
annSearcher,
@@ -132,16 +132,23 @@ class PrimaryKeyVectorBucketSearchTest {
Arrays.asList(data1, data2),
deletionVectors,
new float[] {0, 0},
+ 2,
2);
- assertThat(results)
+ assertThat(results.indexedCandidates())
+ .extracting(
+ PkVectorSearchResult::dataFileName,
+ PkVectorSearchResult::rowPosition,
+ PkVectorSearchResult::distance)
+ .containsExactly(org.assertj.core.groups.Tuple.tuple("data-1",
1L, 0.5F));
+ assertThat(results.exactCandidates())
.extracting(
PkVectorSearchResult::dataFileName,
PkVectorSearchResult::rowPosition,
PkVectorSearchResult::distance)
.containsExactly(
- org.assertj.core.groups.Tuple.tuple("data-1", 1L,
0.5F),
- org.assertj.core.groups.Tuple.tuple("data-2", 0L, 1F));
+ org.assertj.core.groups.Tuple.tuple("data-2", 0L, 1F),
+ org.assertj.core.groups.Tuple.tuple("data-2", 1L, 9F));
verify(readerFactory, never()).create(data1);
}
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 5c6558fdfe..dd34a26638 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
@@ -50,23 +50,72 @@ class PrimaryKeyVectorSearchTest extends TableTestBase {
}
private Schema vectorSchema(String mergeEngine, boolean
deletionVectorsEnabled) {
- return Schema.newBuilder()
- .column("id", DataTypes.INT())
- .column("embedding", DataTypes.VECTOR(2, DataTypes.FLOAT()))
- .primaryKey("id")
- .option(CoreOptions.BUCKET.key(), "1")
- .option(CoreOptions.MERGE_ENGINE.key(), mergeEngine)
- .option(
- CoreOptions.DELETION_VECTORS_ENABLED.key(),
- Boolean.toString(deletionVectorsEnabled))
- .option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding")
- .option(
- "fields.embedding.pk-vector.index.type",
- TestVectorGlobalIndexerFactory.IDENTIFIER)
- .option("fields.embedding.pk-vector.distance.metric", "l2")
- .option("test.vector.dimension", "2")
- .option("test.vector.metric", "l2")
- .build();
+ return vectorSchema(mergeEngine, deletionVectorsEnabled, false);
+ }
+
+ private Schema vectorSchema(
+ String mergeEngine, boolean deletionVectorsEnabled, boolean
reverseScore) {
+ Schema.Builder builder =
+ Schema.newBuilder()
+ .column("id", DataTypes.INT())
+ .column("embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()))
+ .primaryKey("id")
+ .option(CoreOptions.BUCKET.key(), "1")
+ .option(CoreOptions.MERGE_ENGINE.key(), mergeEngine)
+ .option(
+ CoreOptions.DELETION_VECTORS_ENABLED.key(),
+ Boolean.toString(deletionVectorsEnabled))
+ .option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(),
"embedding")
+ .option(
+ "fields.embedding.pk-vector.index.type",
+ TestVectorGlobalIndexerFactory.IDENTIFIER)
+ .option("fields.embedding.pk-vector.distance.metric",
"l2")
+ .option("test.vector.dimension", "2")
+ .option("test.vector.metric", "l2");
+ if (reverseScore) {
+ builder.option("test.vector.reverse-score", "true");
+ }
+ return builder.build();
+ }
+
+ @Test
+ void testRefineFactorReranksAnnCandidates() throws Exception {
+ catalog.createTable(identifier(), vectorSchema("deduplicate", true,
true), false);
+ FileStoreTable table = getTableDefault();
+ write(
+ table,
+ ioManager,
+ GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[]
{1, 0})),
+ GenericRow.of(2, BinaryVector.fromPrimitiveArray(new float[]
{2, 0})),
+ GenericRow.of(3, BinaryVector.fromPrimitiveArray(new float[]
{3, 0})));
+
+ GlobalIndexResult approximate =
+ table.newVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVector(new float[] {0, 0})
+ .withLimit(1)
+ .executeLocal();
+ assertThat(readIds(table, approximate)).containsExactly(3);
+
+ GlobalIndexResult refined =
+ table.newVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVector(new float[] {0, 0})
+ .withLimit(1)
+ .withOption("refine_factor", "3")
+ .executeLocal();
+ assertThat(readIds(table, refined)).containsExactly(1);
+
+ GlobalIndexResult factorOne =
+ table.newVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVector(new float[] {0, 0})
+ .withLimit(1)
+ .withOption("refine_factor", "1")
+ .executeLocal();
+ assertThat(readIds(table, factorOne)).containsExactly(3);
+ assertThat(((PrimaryKeyVectorResult)
factorOne).splits().get(0).scores())
+ .containsExactly(0.1F);
}
@Test
@@ -163,4 +212,15 @@ class PrimaryKeyVectorSearchTest extends TableTestBase {
assertThat(ids).containsExactly(1);
}
+
+ private static List<Integer> readIds(FileStoreTable table,
GlobalIndexResult result)
+ throws Exception {
+ ReadBuilder readBuilder = table.newReadBuilder();
+ TableScan.Plan plan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+ List<Integer> ids = new ArrayList<>();
+ try (RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)) {
+ reader.forEachRemaining(row -> ids.add(row.getInt(0)));
+ }
+ return ids;
+ }
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchRefineOptionsTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchRefineOptionsTest.java
new file mode 100644
index 0000000000..6c464160ce
--- /dev/null
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchRefineOptionsTest.java
@@ -0,0 +1,112 @@
+/*
+ * 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.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static
org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+
+/** Tests for shared vector refine option resolution. */
+class VectorSearchRefineOptionsTest {
+
+ @Test
+ void testQueryOptionsTakePrecedenceOverMoreSpecificTableOptions() {
+ Map<String, String> query = Collections.singletonMap("rerank-factor",
"2");
+ Map<String, String> table =
+
Collections.singletonMap("fields.embedding.ivf-pq.refine_factor", "7");
+
+ assertThat(VectorSearchRefineOptions.resolve(query, table,
"embedding", "ivf-pq"))
+ .isEqualTo(2);
+ }
+
+ @Test
+ void testPrefixAndAliasPrecedence() {
+ Map<String, String> options = new HashMap<>();
+ options.put("refine_factor", "1");
+ options.put("ivf.refine_factor", "2");
+ options.put("ivf_pq.refine_factor", "3");
+ options.put("ivf-pq.refine_factor", "4");
+ options.put("fields.embedding.refine_factor", "5");
+ options.put("fields.embedding.ivf.refine_factor", "6");
+ options.put("fields.embedding.ivf_pq.refine_factor", "7");
+ options.put("fields.embedding.ivf-pq.rerank-factor", "8");
+
+ assertThat(
+ VectorSearchRefineOptions.resolve(
+ Collections.emptyMap(), options, "embedding",
"ivf-pq"))
+ .isEqualTo(8);
+ }
+
+ @Test
+ void testAliasPrecedenceWithinPrefix() {
+ Map<String, String> options = new HashMap<>();
+ options.put("refine_factor", "2");
+ options.put("refine-factor", "3");
+ options.put("rerank_factor", "4");
+ options.put("rerank-factor", "5");
+
+ assertThat(
+ VectorSearchRefineOptions.resolve(
+ options, Collections.emptyMap(), "embedding",
"ivf-pq"))
+ .isEqualTo(2);
+ }
+
+ @Test
+ void testDisabledAndCheckedSearchLimit() {
+ assertThat(
+ VectorSearchRefineOptions.resolve(
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ "embedding",
+ "ivf-pq"))
+ .isZero();
+ assertThat(VectorSearchRefineOptions.searchLimit(10, 0)).isEqualTo(10);
+ assertThat(VectorSearchRefineOptions.searchLimit(10, 1)).isEqualTo(10);
+ assertThat(VectorSearchRefineOptions.searchLimit(10, 3)).isEqualTo(30);
+ }
+
+ @Test
+ void testRejectsInvalidFactorsAndOverflow() {
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> resolveQuery("0"))
+ .withMessageContaining("refine factor must be positive");
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> resolveQuery("-1"))
+ .withMessageContaining("refine factor must be positive");
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> resolveQuery("abc"))
+ .withMessageContaining("Invalid vector refine factor");
+ assertThatIllegalArgumentException()
+ .isThrownBy(() ->
VectorSearchRefineOptions.searchLimit(Integer.MAX_VALUE, 2))
+ .withMessageContaining("limit overflow");
+ }
+
+ private static int resolveQuery(String value) {
+ return VectorSearchRefineOptions.resolve(
+ Collections.singletonMap("refine_factor", value),
+ Collections.emptyMap(),
+ "embedding",
+ "ivf-pq");
+ }
+}
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
index 053abc9677..5fdd42895f 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
@@ -81,11 +81,11 @@ public class SparkPrimaryKeyVectorRead extends
PrimaryKeyVectorRead {
}
};
List<byte[]> groupResults = mapInSpark(groups, task, groups.size());
- List<Candidate> candidates = new ArrayList<>();
+ List<SearchResult> searchResults = new
ArrayList<>(groupResults.size());
for (byte[] groupResult : groupResults) {
- candidates.addAll(deserializeCandidates(groupResult));
+ searchResults.add(deserializeSearchResult(groupResult));
}
- return createResult(primaryKeyPlan, candidates);
+ return createResult(primaryKeyPlan, mergeSearchResults(searchResults));
}
protected int sparkParallelism() {
@@ -121,13 +121,13 @@ public class SparkPrimaryKeyVectorRead extends
PrimaryKeyVectorRead {
}
}
- @SuppressWarnings("unchecked")
- private List<Candidate> deserializeCandidates(byte[] bytes) {
+ private SearchResult deserializeSearchResult(byte[] bytes) {
try {
return InstantiationUtil.deserializeObject(
bytes, Thread.currentThread().getContextClassLoader());
} catch (IOException | ClassNotFoundException e) {
- throw new RuntimeException("Failed to deserialize primary-key
vector candidates.", e);
+ throw new RuntimeException(
+ "Failed to deserialize primary-key vector search result.",
e);
}
}
}
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 84e896ae50..155f314eba 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
@@ -152,6 +152,47 @@ class PrimaryKeyVectorSearchTest extends
PaimonSparkTestBase {
}
}
+ test("distributed primary-key vector search reranks candidates on driver") {
+ withTable("T") {
+ createVectorTable(
+ bucket = 4,
+ extraOptions = Seq(
+ "global-index.thread-num" -> "2",
+ "test.vector.reverse-score" -> "true",
+ "fields.embedding.test-vector-ann.refine_factor" -> "16")
+ )
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(1.0f, 0.0f)),
+ | (2, array(2.0f, 0.0f)),
+ | (3, array(3.0f, 0.0f)),
+ | (4, array(4.0f, 0.0f)),
+ | (5, array(5.0f, 0.0f)),
+ | (6, array(6.0f, 0.0f)),
+ | (7, array(7.0f, 0.0f)),
+ | (8, array(8.0f, 0.0f))
+ |""".stripMargin)
+
+ val jobGroup = s"primary-key-vector-rerank-${System.nanoTime()}"
+ spark.sparkContext.setJobGroup(jobGroup, jobGroup)
+ try {
+ withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" ->
"true") {
+ val ids = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 1)
+ |""".stripMargin)
+ .collect()
+ .map(_.getInt(0))
+ assert(ids.toSeq == Seq(1))
+ }
+ } finally {
+ spark.sparkContext.clearJobGroup()
+ }
+
assert(spark.sparkContext.statusTracker.getJobIdsForGroup(jobGroup).nonEmpty)
+ }
+ }
+
test("primary-key vector search prunes partitions before top k") {
withTable("T") {
createVectorTable(
