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 cd0f9251bc [core] Add primary-key vector bucket search (#8569)
cd0f9251bc is described below
commit cd0f9251bc55f3f25cb463a17849bd17a2d614bb
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Jul 12 14:21:56 2026 +0800
[core] Add primary-key vector bucket search (#8569)
Adds the bucket-local search kernel for primary-key vector indexes. A
bucket can now merge ANN results with exact fallback results from active
data files that are not covered by ANN yet, while sharing the existing
vector metric semantics and returning physical source positions for the
later read path.
---
.../paimon/globalindex/VectorSearchMetric.java | 115 ++++++++++++
.../index/pkvector/PkVectorAnnSegmentSearcher.java | 67 ++-----
.../index/pkvector/PkVectorExactSearcher.java | 86 +++++++++
.../index/pkvector/PkVectorSearchResult.java | 45 +++++
.../pkvector/PrimaryKeyVectorBucketSearch.java | 128 +++++++++++++
.../paimon/table/source/AbstractVectorRead.java | 53 ++----
.../paimon/globalindex/VectorSearchMetricTest.java | 99 +++++++++++
.../index/pkvector/PkVectorAnnSegmentFileTest.java | 6 +-
.../index/pkvector/PkVectorExactSearcherTest.java | 128 +++++++++++++
.../pkvector/PrimaryKeyVectorBucketSearchTest.java | 197 +++++++++++++++++++++
10 files changed, 828 insertions(+), 96 deletions(-)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/globalindex/VectorSearchMetric.java
b/paimon-core/src/main/java/org/apache/paimon/globalindex/VectorSearchMetric.java
new file mode 100644
index 0000000000..2caf81b5de
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/globalindex/VectorSearchMetric.java
@@ -0,0 +1,115 @@
+/*
+ * 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.globalindex;
+
+import java.util.Locale;
+
+/** Common numerical semantics for vector search metrics. */
+public final class VectorSearchMetric {
+
+ private VectorSearchMetric() {}
+
+ public static String normalize(String metric) {
+ return metric.toLowerCase(Locale.ROOT).replace('-', '_');
+ }
+
+ public static boolean isSupported(String metric) {
+ if (metric == null) {
+ return false;
+ }
+ String normalized = normalize(metric);
+ return "l2".equals(normalized)
+ || "cosine".equals(normalized)
+ || "inner_product".equals(normalized);
+ }
+
+ /** Returns a higher-is-better score for exact vector search. */
+ public static float computeScore(float[] query, float[] stored, String
metric) {
+ if ("l2".equals(metric)) {
+ return 1.0f / (1.0f + squaredL2(query, stored));
+ } else if ("cosine".equals(metric)) {
+ return cosineSimilarity(query, stored);
+ } else if ("inner_product".equals(metric)) {
+ return innerProduct(query, stored);
+ }
+ throw unknownMetric(metric);
+ }
+
+ /** Returns a lower-is-better distance for exact vector search. */
+ public static float computeDistance(float[] query, float[] stored, String
metric) {
+ if ("l2".equals(metric)) {
+ return squaredL2(query, stored);
+ } else if ("cosine".equals(metric)) {
+ return cosineDistance(cosineSimilarity(query, stored));
+ } else if ("inner_product".equals(metric)) {
+ return -innerProduct(query, stored);
+ }
+ throw unknownMetric(metric);
+ }
+
+ /** Converts a standardized higher-is-better index score to a
lower-is-better distance. */
+ public static float scoreToDistance(float score, String metric) {
+ if ("l2".equals(metric)) {
+ return 1.0f / score - 1.0f;
+ } else if ("cosine".equals(metric)) {
+ return cosineDistance(score);
+ } else if ("inner_product".equals(metric)) {
+ return -score;
+ }
+ throw unknownMetric(metric);
+ }
+
+ private static float squaredL2(float[] query, float[] stored) {
+ float squared = 0;
+ for (int i = 0; i < query.length; i++) {
+ float delta = query[i] - stored[i];
+ squared += delta * delta;
+ }
+ return squared;
+ }
+
+ private static float cosineSimilarity(float[] query, float[] stored) {
+ float dot = 0;
+ float queryNorm = 0;
+ float storedNorm = 0;
+ for (int i = 0; i < query.length; i++) {
+ dot += query[i] * stored[i];
+ queryNorm += query[i] * query[i];
+ storedNorm += stored[i] * stored[i];
+ }
+ float denominator = (float) (Math.sqrt(queryNorm) *
Math.sqrt(storedNorm));
+ return denominator == 0 ? 0 : dot / denominator;
+ }
+
+ private static float innerProduct(float[] query, float[] stored) {
+ float dot = 0;
+ for (int i = 0; i < query.length; i++) {
+ dot += query[i] * stored[i];
+ }
+ return dot;
+ }
+
+ private static float cosineDistance(float similarity) {
+ return 1.0f - Math.max(-1.0f, Math.min(1.0f, similarity));
+ }
+
+ private static IllegalArgumentException unknownMetric(String metric) {
+ return new IllegalArgumentException("Unknown vector search metric: " +
metric);
+ }
+}
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 e3b2fbba7f..5722995225 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
@@ -25,6 +25,7 @@ import org.apache.paimon.globalindex.GlobalIndexReader;
import org.apache.paimon.globalindex.GlobalIndexer;
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
import org.apache.paimon.globalindex.VectorGlobalIndexer;
+import org.apache.paimon.globalindex.VectorSearchMetric;
import org.apache.paimon.index.GlobalIndexMeta;
import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.options.Options;
@@ -41,7 +42,6 @@ import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
-import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ExecutorService;
@@ -51,10 +51,10 @@ import static
org.apache.paimon.utils.Preconditions.checkArgument;
/** Searches one ANN payload and maps its segment-local ids back to source row
positions. */
public class PkVectorAnnSegmentSearcher {
- private static final Comparator<Candidate> BEST_FIRST =
- Comparator.comparingDouble((Candidate candidate) ->
candidate.distance)
- .thenComparing(candidate -> candidate.dataFileName)
- .thenComparingLong(candidate -> candidate.rowPosition);
+ private static final Comparator<PkVectorSearchResult> BEST_FIRST =
+ Comparator.comparingDouble(PkVectorSearchResult::distance)
+ .thenComparing(PkVectorSearchResult::dataFileName)
+ .thenComparingLong(PkVectorSearchResult::rowPosition);
private final FileIO fileIO;
private final PkVectorAnnSegmentFile annSegmentFile;
@@ -74,11 +74,11 @@ public class PkVectorAnnSegmentSearcher {
this.annSegmentFile = annSegmentFile;
this.vectorField = vectorField;
this.indexOptions = indexOptions;
- this.metric = normalizeMetric(metric);
+ this.metric = VectorSearchMetric.normalize(metric);
this.executor = executor;
}
- public List<Candidate> search(
+ public List<PkVectorSearchResult> search(
IndexFileMeta segment,
PkVectorSourceMeta sourceMeta,
float[] query,
@@ -95,7 +95,7 @@ public class PkVectorAnnSegmentSearcher {
return search(segment, sourceMeta, query, limit, deletionVectors,
searchOptions);
}
- public List<Candidate> search(
+ public List<PkVectorSearchResult> search(
IndexFileMeta segment,
PkVectorSourceMeta sourceMeta,
float[] query,
@@ -114,7 +114,8 @@ public class PkVectorAnnSegmentSearcher {
indexer instanceof VectorGlobalIndexer,
"Index algorithm %s does not implement VectorGlobalIndexer.",
segment.indexType());
- String readerMetric = normalizeMetric(((VectorGlobalIndexer)
indexer).metric());
+ String readerMetric =
+ VectorSearchMetric.normalize(((VectorGlobalIndexer)
indexer).metric());
checkArgument(
metric.equals(readerMetric),
"ANN segment metric %s does not match index reader metric %s.",
@@ -144,7 +145,7 @@ public class PkVectorAnnSegmentSearcher {
}
long sourceRowCount = totalRowCount(sourceMeta.sourceFiles());
- List<Candidate> candidates = new ArrayList<>();
+ List<PkVectorSearchResult> candidates = new ArrayList<>();
ScoredGlobalIndexResult scored = result.get();
for (long ordinal : scored.results()) {
checkArgument(
@@ -162,10 +163,11 @@ public class PkVectorAnnSegmentSearcher {
segment.fileName(),
filePosition.rowPosition);
candidates.add(
- new Candidate(
+ new PkVectorSearchResult(
filePosition.dataFileName,
filePosition.rowPosition,
-
scoreToDistance(scored.scoreGetter().score(ordinal), metric)));
+ VectorSearchMetric.scoreToDistance(
+ scored.scoreGetter().score(ordinal),
metric)));
}
Collections.sort(candidates, BEST_FIRST);
return Collections.unmodifiableList(candidates);
@@ -218,47 +220,6 @@ public class PkVectorAnnSegmentSearcher {
throw new IllegalArgumentException("ANN ordinal is outside source
files: " + ordinal);
}
- private static float scoreToDistance(float score, String metric) {
- if ("l2".equals(metric)) {
- return 1F / score - 1F;
- } else if ("cosine".equals(metric)) {
- return 1F - score;
- } else if ("inner_product".equals(metric)) {
- return -score;
- }
- throw new IllegalArgumentException("Unsupported ANN vector metric: " +
metric);
- }
-
- private static String normalizeMetric(String metric) {
- return metric.toLowerCase(Locale.ROOT).replace('-', '_');
- }
-
- /** One ANN candidate addressed by source-file row position. */
- public static class Candidate {
-
- private final long rowPosition;
- private final float distance;
- private final String dataFileName;
-
- private Candidate(String dataFileName, long rowPosition, float
distance) {
- this.dataFileName = dataFileName;
- this.rowPosition = rowPosition;
- this.distance = distance;
- }
-
- public String dataFileName() {
- return dataFileName;
- }
-
- public long rowPosition() {
- return rowPosition;
- }
-
- public float distance() {
- return distance;
- }
- }
-
private static class FilePosition {
private final String dataFileName;
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
new file mode 100644
index 0000000000..e63fa700b8
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorExactSearcher.java
@@ -0,0 +1,86 @@
+/*
+ * 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.index.pkvector;
+
+import org.apache.paimon.globalindex.VectorSearchMetric;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.PriorityQueue;
+import java.util.function.LongPredicate;
+
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+
+/** Exact Top-K over one sequential physical-row vector source. */
+public final class PkVectorExactSearcher {
+
+ private PkVectorExactSearcher() {}
+
+ public static List<PkVectorSearchResult> search(
+ String dataFileName,
+ PkVectorReader reader,
+ float[] query,
+ String metric,
+ int limit,
+ LongPredicate excludedPosition)
+ throws IOException {
+ checkArgument(query.length == reader.dimension(), "Query vector
dimension does not match.");
+ 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());
+ 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);
+ }
+ }
+ List<PkVectorSearchResult> result = new ArrayList<>(nearest);
+ Collections.sort(result, bestFirst);
+ return Collections.unmodifiableList(result);
+ }
+}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorSearchResult.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorSearchResult.java
new file mode 100644
index 0000000000..18572d25f1
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorSearchResult.java
@@ -0,0 +1,45 @@
+/*
+ * 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.index.pkvector;
+
+/** One vector search result addressed by source data-file row position. */
+public final class PkVectorSearchResult {
+
+ private final String dataFileName;
+ private final long rowPosition;
+ private final float distance;
+
+ public PkVectorSearchResult(String dataFileName, long rowPosition, float
distance) {
+ this.dataFileName = dataFileName;
+ this.rowPosition = rowPosition;
+ this.distance = distance;
+ }
+
+ public String dataFileName() {
+ return dataFileName;
+ }
+
+ public long rowPosition() {
+ return rowPosition;
+ }
+
+ public float distance() {
+ return distance;
+ }
+}
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
new file mode 100644
index 0000000000..f7b358cf3b
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
@@ -0,0 +1,128 @@
+/*
+ * 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.index.pkvector;
+
+import org.apache.paimon.deletionvectors.DeletionVector;
+import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.io.DataFileMeta;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.function.LongPredicate;
+
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+
+/** ANN plus exact data-file fallback search for one snapshot bucket. */
+public class PrimaryKeyVectorBucketSearch {
+
+ private static final Comparator<PkVectorSearchResult> BEST_FIRST =
+ Comparator.comparingDouble(PkVectorSearchResult::distance)
+ .thenComparing(PkVectorSearchResult::dataFileName)
+ .thenComparingLong(PkVectorSearchResult::rowPosition);
+
+ private final PkVectorDataFileReader.Factory vectorReaderFactory;
+ @Nullable private final PkVectorAnnSegmentSearcher annSearcher;
+ private final Map<String, String> searchOptions;
+ private final String metric;
+
+ public PrimaryKeyVectorBucketSearch(
+ PkVectorDataFileReader.Factory vectorReaderFactory,
+ @Nullable PkVectorAnnSegmentSearcher annSearcher,
+ Map<String, String> searchOptions,
+ String metric) {
+ this.vectorReaderFactory = vectorReaderFactory;
+ this.annSearcher = annSearcher;
+ this.searchOptions = Collections.unmodifiableMap(new
HashMap<>(searchOptions));
+ this.metric = metric;
+ }
+
+ public List<PkVectorSearchResult> search(
+ PkVectorBucketIndexState state,
+ List<DataFileMeta> activeFiles,
+ Map<String, DeletionVector> deletionVectors,
+ float[] query,
+ int limit)
+ throws IOException {
+ checkArgument(limit > 0, "Vector 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());
+ Set<String> covered = new HashSet<>();
+ for (IndexFileMeta ann : state.annSegments()) {
+ PkVectorSourceMeta sourceMeta =
PkVectorSourceMeta.fromIndexFile(ann);
+ for (PkVectorSourceFile source : sourceMeta.sourceFiles()) {
+ DataFileMeta file = filesByName.get(source.fileName());
+ checkArgument(
+ file != null && file.rowCount() == source.rowCount(),
+ "ANN source %s does not match the active data file.",
+ source.fileName());
+ covered.add(source.fileName());
+ }
+ checkArgument(annSearcher != null, "ANN search is not
configured.");
+ for (PkVectorSearchResult result :
+ annSearcher.search(
+ ann, sourceMeta, query, limit, deletionVectors,
searchOptions)) {
+ add(nearest, result, limit);
+ }
+ }
+
+ for (DataFileMeta file : activeFiles) {
+ if (covered.contains(file.fileName())) {
+ continue;
+ }
+ DeletionVector dv = deletionVectors.get(file.fileName());
+ LongPredicate excluded = dv == null ? position -> false :
dv::isDeleted;
+ try (PkVectorReader reader = vectorReaderFactory.create(file)) {
+ for (PkVectorSearchResult result :
+ PkVectorExactSearcher.search(
+ file.fileName(), reader, query, metric, limit,
excluded)) {
+ add(nearest, result, limit);
+ }
+ }
+ }
+ List<PkVectorSearchResult> result = new ArrayList<>(nearest);
+ Collections.sort(result, BEST_FIRST);
+ return Collections.unmodifiableList(result);
+ }
+
+ private static void add(
+ PriorityQueue<PkVectorSearchResult> nearest,
+ PkVectorSearchResult candidate,
+ int limit) {
+ if (nearest.size() < limit) {
+ nearest.add(candidate);
+ } else if (BEST_FIRST.compare(candidate, nearest.peek()) < 0) {
+ nearest.poll();
+ nearest.add(candidate);
+ }
+ }
+}
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 93341f3f1b..dd43a99af8 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
@@ -31,6 +31,7 @@ import
org.apache.paimon.globalindex.GlobalIndexerFactoryUtils;
import org.apache.paimon.globalindex.OffsetGlobalIndexReader;
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
import org.apache.paimon.globalindex.VectorGlobalIndexer;
+import org.apache.paimon.globalindex.VectorSearchMetric;
import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
import org.apache.paimon.index.GlobalIndexMeta;
import org.apache.paimon.index.IndexFileMeta;
@@ -447,7 +448,10 @@ public abstract class AbstractVectorRead implements
Serializable {
float[] stored = getVector(row, vectorIndex);
checkVectorDimension(queryVector, stored);
offerScore(
- topKHeap, limit, rowId,
computeScore(queryVector, stored, metric));
+ topKHeap,
+ limit,
+ rowId,
+ VectorSearchMetric.computeScore(queryVector,
stored, metric));
});
} catch (IOException e) {
throw new RuntimeException("Failed to read raw vectors for vector
search.", e);
@@ -523,7 +527,11 @@ public abstract class AbstractVectorRead implements
Serializable {
continue;
}
checkVectorDimension(queryVector, stored);
- offerScore(topKHeap, topK, rowId, computeScore(queryVector,
stored, metric));
+ offerScore(
+ topKHeap,
+ topK,
+ rowId,
+ VectorSearchMetric.computeScore(queryVector, stored,
metric));
}
return scoredResult(topKHeap);
}
@@ -655,7 +663,7 @@ public abstract class AbstractVectorRead implements
Serializable {
if (metric == null) {
metric = configuredRawSearchMetric();
}
- return metric == null ? "l2" : normalizeMetric(metric);
+ return metric == null ? "l2" : VectorSearchMetric.normalize(metric);
}
@Nullable
@@ -690,7 +698,7 @@ public abstract class AbstractVectorRead implements
Serializable {
for (Map.Entry<String, String> entry : options.entrySet()) {
String key = entry.getKey();
if (key.endsWith(".distance.metric") || key.endsWith(".metric")) {
- String value = normalizeMetric(entry.getValue());
+ String value = VectorSearchMetric.normalize(entry.getValue());
if (isRawSearchMetric(value)) {
if (metric != null && !metric.equals(value)) {
return null;
@@ -705,15 +713,11 @@ public abstract class AbstractVectorRead implements
Serializable {
@Nullable
private static String option(Map<String, String> options, String key) {
String value = options.get(key);
- return value == null ? null : normalizeMetric(value);
+ return value == null ? null : VectorSearchMetric.normalize(value);
}
private static boolean isRawSearchMetric(String metric) {
- return "l2".equals(metric) || "cosine".equals(metric) ||
"inner_product".equals(metric);
- }
-
- private static String normalizeMetric(String metric) {
- return metric.toLowerCase().replace('-', '_');
+ return VectorSearchMetric.isSupported(metric);
}
protected int configuredRefineFactor(String indexType) {
@@ -795,35 +799,6 @@ public abstract class AbstractVectorRead implements
Serializable {
throw new IllegalArgumentException("No vector index files found.");
}
- private static float computeScore(float[] query, float[] stored, String
metric) {
- if ("l2".equals(metric)) {
- float sumSq = 0;
- for (int i = 0; i < query.length; i++) {
- float diff = query[i] - stored[i];
- sumSq += diff * diff;
- }
- return 1.0f / (1.0f + sumSq);
- } else if ("cosine".equals(metric)) {
- float dot = 0;
- float normA = 0;
- float normB = 0;
- for (int i = 0; i < query.length; i++) {
- dot += query[i] * stored[i];
- normA += query[i] * query[i];
- normB += stored[i] * stored[i];
- }
- float denominator = (float) (Math.sqrt(normA) * Math.sqrt(normB));
- return denominator == 0 ? 0 : dot / denominator;
- } else if ("inner_product".equals(metric)) {
- float dot = 0;
- for (int i = 0; i < query.length; i++) {
- dot += query[i] * stored[i];
- }
- return dot;
- }
- throw new IllegalArgumentException("Unknown vector search metric: " +
metric);
- }
-
private static List<Optional<ScoredGlobalIndexResult>>
emptyOptionalResults(int n) {
List<Optional<ScoredGlobalIndexResult>> results = new ArrayList<>(n);
for (int i = 0; i < n; i++) {
diff --git
a/paimon-core/src/test/java/org/apache/paimon/globalindex/VectorSearchMetricTest.java
b/paimon-core/src/test/java/org/apache/paimon/globalindex/VectorSearchMetricTest.java
new file mode 100644
index 0000000000..e13e7c66e2
--- /dev/null
+++
b/paimon-core/src/test/java/org/apache/paimon/globalindex/VectorSearchMetricTest.java
@@ -0,0 +1,99 @@
+/*
+ * 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.globalindex;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static
org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+
+/** Tests for {@link VectorSearchMetric}. */
+class VectorSearchMetricTest {
+
+ @Test
+ void testNormalizeAndSupportedMetrics() {
+
assertThat(VectorSearchMetric.normalize("INNER-PRODUCT")).isEqualTo("inner_product");
+ assertThat(VectorSearchMetric.isSupported("L2")).isTrue();
+ assertThat(VectorSearchMetric.isSupported("Cosine")).isTrue();
+ assertThat(VectorSearchMetric.isSupported("INNER-PRODUCT")).isTrue();
+ assertThat(VectorSearchMetric.isSupported("manhattan")).isFalse();
+ assertThat(VectorSearchMetric.isSupported(null)).isFalse();
+ }
+
+ @Test
+ void testScoreAndDistanceForSupportedMetrics() {
+ float[] query = {2, 0};
+ float[] stored = {1, 0};
+
+ assertThat(VectorSearchMetric.computeScore(query, stored,
"l2")).isEqualTo(0.5F);
+ assertThat(VectorSearchMetric.computeDistance(query, stored,
"l2")).isEqualTo(1F);
+ assertThat(VectorSearchMetric.computeScore(query, stored,
"cosine")).isEqualTo(1F);
+ assertThat(VectorSearchMetric.computeDistance(query, stored,
"cosine")).isEqualTo(0F);
+ assertThat(VectorSearchMetric.computeScore(query, stored,
"inner_product")).isEqualTo(2F);
+ assertThat(VectorSearchMetric.computeDistance(query, stored,
"inner_product"))
+ .isEqualTo(-2F);
+ }
+
+ @Test
+ void testZeroVectorCosine() {
+ float[] zero = {0, 0};
+ float[] stored = {1, 2};
+
+ assertThat(VectorSearchMetric.computeScore(zero, stored,
"cosine")).isZero();
+ assertThat(VectorSearchMetric.computeDistance(zero, stored,
"cosine")).isEqualTo(1F);
+ }
+
+ @Test
+ void testCosineScoreIsUnclampedButDistanceIsBounded() {
+ float[] query = {0.34558418F, 0.82161814F};
+ float[] stored = {0.3455843F, 0.82161707F};
+
+ assertThat(VectorSearchMetric.computeScore(query, stored,
"cosine")).isEqualTo(1.0000001F);
+ assertThat(VectorSearchMetric.computeDistance(query, stored,
"cosine")).isZero();
+ }
+
+ @Test
+ void testRepresentativeAndLargeL2UseFloatAccumulation() {
+ assertThat(
+ VectorSearchMetric.computeDistance(
+ new float[] {0.1F, 0.2F, 0.3F},
+ new float[] {0.4F, 0.5F, 0.6F},
+ "l2"))
+ .isEqualTo(0.27F);
+
+ float[] query = {100_000F, 100_000F, 1F};
+ float[] stored = {-100_000F, -100_000F, 0F};
+ float distance = VectorSearchMetric.computeDistance(query, stored,
"l2");
+ assertThat(distance).isEqualTo(80_000_000_000F);
+ assertThat(VectorSearchMetric.computeScore(query, stored, "l2"))
+ .isEqualTo(1F / (1F + distance));
+ }
+
+ @Test
+ void testScoreToDistance() {
+ assertThat(VectorSearchMetric.scoreToDistance(0.25F,
"l2")).isEqualTo(3F);
+ assertThat(VectorSearchMetric.scoreToDistance(0.25F,
"cosine")).isEqualTo(0.75F);
+ assertThat(VectorSearchMetric.scoreToDistance(1.0000001F,
"cosine")).isZero();
+ assertThat(VectorSearchMetric.scoreToDistance(-1.0000001F,
"cosine")).isEqualTo(2F);
+ assertThat(VectorSearchMetric.scoreToDistance(0.25F,
"inner_product")).isEqualTo(-0.25F);
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> VectorSearchMetric.scoreToDistance(1F,
"manhattan"))
+ .withMessageContaining("Unknown vector search metric");
+ }
+}
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 1a30e207f4..0f7557c848 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
@@ -103,7 +103,7 @@ class PkVectorAnnSegmentFileTest {
deletionVectors.put("data-2", data2Deletes);
ExecutorService executor = Executors.newSingleThreadExecutor();
- List<PkVectorAnnSegmentSearcher.Candidate> candidates;
+ List<PkVectorSearchResult> candidates;
try {
candidates =
new PkVectorAnnSegmentSearcher(
@@ -120,9 +120,7 @@ class PkVectorAnnSegmentFileTest {
}
assertThat(candidates)
- .extracting(
- PkVectorAnnSegmentSearcher.Candidate::dataFileName,
- PkVectorAnnSegmentSearcher.Candidate::rowPosition)
+ .extracting(PkVectorSearchResult::dataFileName,
PkVectorSearchResult::rowPosition)
.containsExactly(
org.assertj.core.groups.Tuple.tuple("data-2", 1L),
org.assertj.core.groups.Tuple.tuple("data-1", 0L),
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
new file mode 100644
index 0000000000..8f7e14fe0d
--- /dev/null
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorExactSearcherTest.java
@@ -0,0 +1,128 @@
+/*
+ * 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.index.pkvector;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static
org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+
+/** Tests for {@link PkVectorExactSearcher}. */
+class PkVectorExactSearcherTest {
+
+ @Test
+ void testDistancesForSupportedMetrics() throws Exception {
+ assertThat(distance("l2")).isEqualTo(1F);
+ assertThat(distance("cosine")).isEqualTo(0F);
+ assertThat(distance("inner_product")).isEqualTo(-2F);
+ }
+
+ @Test
+ void testRejectsInvalidSearchInput() {
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> search(new float[] {1}, "l2", 1))
+ .withMessageContaining("dimension");
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> search(new float[] {1, 0}, "l2", 0))
+ .withMessageContaining("positive");
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> search(new float[] {1, 0}, "manhattan", 1))
+ .withMessageContaining("Unsupported");
+ assertThatIllegalArgumentException()
+ .isThrownBy(() -> search(new float[] {Float.NaN, 0}, "l2", 1))
+ .withMessageContaining("finite");
+ }
+
+ @Test
+ void testPreservesNullAndDeletedPhysicalPositions() throws Exception {
+ float[][] vectors = {{3, 0}, null, {1, 0}, {2, 0}};
+ List<PkVectorSearchResult> results;
+ try (PkVectorReader reader = new ArrayReader(vectors)) {
+ results =
+ PkVectorExactSearcher.search(
+ "data-file",
+ reader,
+ new float[] {0, 0},
+ "l2",
+ 2,
+ position -> position == 2);
+ }
+
+ assertThat(results)
+ .extracting(
+ PkVectorSearchResult::dataFileName,
+ PkVectorSearchResult::rowPosition,
+ PkVectorSearchResult::distance)
+ .containsExactly(
+ org.assertj.core.groups.Tuple.tuple("data-file", 3L,
4F),
+ org.assertj.core.groups.Tuple.tuple("data-file", 0L,
9F));
+ }
+
+ private static float distance(String metric) throws IOException {
+ try (PkVectorReader reader = new ArrayReader(new float[][] {{1, 0}})) {
+ return PkVectorExactSearcher.search(
+ "data-file", reader, new float[] {2, 0}, metric,
1, position -> false)
+ .get(0)
+ .distance();
+ }
+ }
+
+ private static void search(float[] query, String metric, int limit) throws
IOException {
+ try (PkVectorReader reader = new ArrayReader(new float[][] {{1, 0}})) {
+ PkVectorExactSearcher.search(
+ "data-file", reader, query, metric, limit, position ->
false);
+ }
+ }
+
+ private static class ArrayReader implements PkVectorReader {
+
+ private final float[][] vectors;
+ private int position;
+
+ private ArrayReader(float[][] vectors) {
+ this.vectors = vectors;
+ }
+
+ @Override
+ public int dimension() {
+ return 2;
+ }
+
+ @Override
+ public long rowCount() {
+ return vectors.length;
+ }
+
+ @Override
+ public boolean readNextVector(float[] reuse) {
+ float[] vector = vectors[position++];
+ if (vector == null) {
+ return false;
+ }
+ System.arraycopy(vector, 0, reuse, 0, reuse.length);
+ return true;
+ }
+
+ @Override
+ public void close() throws IOException {}
+ }
+}
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
new file mode 100644
index 0000000000..5e99f96deb
--- /dev/null
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
@@ -0,0 +1,197 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.index.pkvector;
+
+import org.apache.paimon.deletionvectors.BitmapDeletionVector;
+import org.apache.paimon.deletionvectors.DeletionVector;
+import org.apache.paimon.index.GlobalIndexMeta;
+import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.io.DataFileMeta;
+import org.apache.paimon.manifest.FileSource;
+import org.apache.paimon.stats.SimpleStats;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static
org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Tests for {@link PrimaryKeyVectorBucketSearch}. */
+class PrimaryKeyVectorBucketSearchTest {
+
+ @Test
+ void testRejectsNonPositiveLimitForEmptyBucket() {
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ PrimaryKeyVectorBucketSearch search =
+ new PrimaryKeyVectorBucketSearch(readerFactory, null,
Collections.emptyMap(), "l2");
+
+ assertThatIllegalArgumentException()
+ .isThrownBy(
+ () ->
+ search.search(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Collections.emptyList()),
+ Collections.emptyList(),
+ Collections.emptyMap(),
+ new float[] {0, 0},
+ 0))
+ .withMessageContaining("positive");
+ }
+
+ @Test
+ void testMergesAnnAndExactFallbackWithoutRescanningCoveredFiles() throws
Exception {
+ DataFileMeta data1 = dataFile("data-1");
+ DataFileMeta data2 = dataFile("data-2");
+ IndexFileMeta ann = segment("ann", data1);
+ PkVectorBucketIndexState state =
+ new PkVectorBucketIndexState(7, "test-vector-ann",
Collections.singletonList(ann));
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ PkVectorDataFileReader reader1 = reader(new float[][] {{5, 0}, {6,
0}});
+ PkVectorDataFileReader reader2 = reader(new float[][] {{1, 0}, {3,
0}});
+ when(readerFactory.create(data1)).thenReturn(reader1);
+ when(readerFactory.create(data2)).thenReturn(reader2);
+ PkVectorAnnSegmentSearcher annSearcher =
mock(PkVectorAnnSegmentSearcher.class);
+ Map<String, DeletionVector> deletionVectors = Collections.emptyMap();
+ Map<String, String> searchOptions =
Collections.singletonMap("nprobes", "8");
+ when(annSearcher.search(
+ org.mockito.ArgumentMatchers.eq(ann),
+
org.mockito.ArgumentMatchers.any(PkVectorSourceMeta.class),
+ org.mockito.ArgumentMatchers.any(float[].class),
+ org.mockito.ArgumentMatchers.eq(2),
+ org.mockito.ArgumentMatchers.eq(deletionVectors),
+ org.mockito.ArgumentMatchers.eq(searchOptions)))
+ .thenReturn(Collections.singletonList(new
PkVectorSearchResult("data-1", 1, 0.5F)));
+
+ List<PkVectorSearchResult> results =
+ new PrimaryKeyVectorBucketSearch(readerFactory, annSearcher,
searchOptions, "l2")
+ .search(
+ state,
+ Arrays.asList(data1, data2),
+ deletionVectors,
+ new float[] {0, 0},
+ 2);
+
+ assertThat(results)
+ .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));
+ verify(readerFactory, never()).create(data1);
+ }
+
+ @Test
+ void testExactFallbackMergesFilesAndAppliesDeletionVectors() throws
Exception {
+ DataFileMeta data1 = dataFile("data-1");
+ DataFileMeta data2 = dataFile("data-2");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ PkVectorDataFileReader reader1 = reader(new float[][] {{0, 0}, {2,
0}});
+ PkVectorDataFileReader reader2 = reader(new float[][] {{1, 0}, null});
+ when(readerFactory.create(data1)).thenReturn(reader1);
+ when(readerFactory.create(data2)).thenReturn(reader2);
+ BitmapDeletionVector data1Deletes = new BitmapDeletionVector();
+ data1Deletes.delete(0);
+ Map<String, DeletionVector> deletionVectors = new HashMap<>();
+ deletionVectors.put("data-1", data1Deletes);
+
+ List<PkVectorSearchResult> results =
+ new PrimaryKeyVectorBucketSearch(readerFactory, null,
Collections.emptyMap(), "l2")
+ .search(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Collections.emptyList()),
+ Arrays.asList(data1, data2),
+ deletionVectors,
+ new float[] {0, 0},
+ 2);
+
+ assertThat(results)
+ .extracting(
+ PkVectorSearchResult::dataFileName,
+ PkVectorSearchResult::rowPosition,
+ PkVectorSearchResult::distance)
+ .containsExactly(
+ org.assertj.core.groups.Tuple.tuple("data-2", 0L, 1F),
+ org.assertj.core.groups.Tuple.tuple("data-1", 1L, 4F));
+ }
+
+ private static PkVectorDataFileReader reader(float[][] vectors) throws
IOException {
+ PkVectorDataFileReader reader = mock(PkVectorDataFileReader.class);
+ when(reader.dimension()).thenReturn(2);
+ when(reader.rowCount()).thenReturn((long) vectors.length);
+ final int[] position = {0};
+
when(reader.readNextVector(org.mockito.ArgumentMatchers.any(float[].class)))
+ .thenAnswer(
+ invocation -> {
+ float[] vector = vectors[position[0]++];
+ if (vector == null) {
+ return false;
+ }
+ System.arraycopy(
+ vector, 0, invocation.getArgument(0), 0,
vector.length);
+ return true;
+ });
+ return reader;
+ }
+
+ private static DataFileMeta dataFile(String fileName) {
+ return DataFileMeta.forAppend(
+ fileName,
+ 100,
+ 2,
+ SimpleStats.EMPTY_STATS,
+ 0,
+ 0,
+ 1,
+ Collections.emptyList(),
+ null,
+ FileSource.COMPACT,
+ null,
+ null,
+ null,
+ null);
+ }
+
+ private static IndexFileMeta segment(String fileName, DataFileMeta source)
{
+ byte[] sourceMeta =
+ new PkVectorSourceMeta(
+ Collections.singletonList(
+ new PkVectorSourceFile(
+ source.fileName(),
source.rowCount())))
+ .serialize();
+ return new IndexFileMeta(
+ "test-vector-ann",
+ fileName,
+ 100,
+ source.rowCount(),
+ new GlobalIndexMeta(0, source.rowCount() - 1, 7, null, new
byte[] {1}, sourceMeta),
+ null);
+ }
+}
