From 1e547f14722e4b0826a760989c142718f83e2f3d Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 07:52:13 +0800 Subject: [PATCH] [core] Add primary-key vector bucket search --- .../globalindex/VectorSearchMetric.java | 115 ++++++++++ .../pkvector/PkVectorAnnSegmentSearcher.java | 67 ++---- .../index/pkvector/PkVectorExactSearcher.java | 86 ++++++++ .../index/pkvector/PkVectorSearchResult.java | 45 ++++ .../PrimaryKeyVectorBucketSearch.java | 128 ++++++++++++ .../table/source/AbstractVectorRead.java | 53 ++--- .../globalindex/VectorSearchMetricTest.java | 99 +++++++++ .../pkvector/PkVectorAnnSegmentFileTest.java | 6 +- .../pkvector/PkVectorExactSearcherTest.java | 128 ++++++++++++ .../PrimaryKeyVectorBucketSearchTest.java | 197 ++++++++++++++++++ 10 files changed, 828 insertions(+), 96 deletions(-) create mode 100644 paimon-core/src/main/java/org/apache/paimon/globalindex/VectorSearchMetric.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorExactSearcher.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorSearchResult.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/globalindex/VectorSearchMetricTest.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorExactSearcherTest.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java 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 000000000000..2caf81b5de02 --- /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 e3b2fbba7f1a..5722995225ca 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.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.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 @@ /** Searches one ANN payload and maps its segment-local ids back to source row positions. */ public class PkVectorAnnSegmentSearcher { - private static final Comparator BEST_FIRST = - Comparator.comparingDouble((Candidate candidate) -> candidate.distance) - .thenComparing(candidate -> candidate.dataFileName) - .thenComparingLong(candidate -> candidate.rowPosition); + private static final Comparator 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 PkVectorAnnSegmentSearcher( this.annSegmentFile = annSegmentFile; this.vectorField = vectorField; this.indexOptions = indexOptions; - this.metric = normalizeMetric(metric); + this.metric = VectorSearchMetric.normalize(metric); this.executor = executor; } - public List search( + public List search( IndexFileMeta segment, PkVectorSourceMeta sourceMeta, float[] query, @@ -95,7 +95,7 @@ public List search( return search(segment, sourceMeta, query, limit, deletionVectors, searchOptions); } - public List search( + public List search( IndexFileMeta segment, PkVectorSourceMeta sourceMeta, float[] query, @@ -114,7 +114,8 @@ public List search( 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 List search( } long sourceRowCount = totalRowCount(sourceMeta.sourceFiles()); - List candidates = new ArrayList<>(); + List candidates = new ArrayList<>(); ScoredGlobalIndexResult scored = result.get(); for (long ordinal : scored.results()) { checkArgument( @@ -162,10 +163,11 @@ public List search( 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 @@ private static FilePosition filePosition(List sourceFiles, l 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 000000000000..e63fa700b8f1 --- /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 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 bestFirst = + Comparator.comparingDouble(PkVectorSearchResult::distance) + .thenComparingLong(PkVectorSearchResult::rowPosition); + PriorityQueue 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 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 000000000000..18572d25f190 --- /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 000000000000..f7b358cf3b59 --- /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 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 searchOptions; + private final String metric; + + public PrimaryKeyVectorBucketSearch( + PkVectorDataFileReader.Factory vectorReaderFactory, + @Nullable PkVectorAnnSegmentSearcher annSearcher, + Map searchOptions, + String metric) { + this.vectorReaderFactory = vectorReaderFactory; + this.annSearcher = annSearcher; + this.searchOptions = Collections.unmodifiableMap(new HashMap<>(searchOptions)); + this.metric = metric; + } + + public List search( + PkVectorBucketIndexState state, + List activeFiles, + Map deletionVectors, + float[] query, + int limit) + throws IOException { + checkArgument(limit > 0, "Vector search limit must be positive."); + Map filesByName = new HashMap<>(); + for (DataFileMeta file : activeFiles) { + checkArgument(filesByName.put(file.fileName(), file) == null, "Duplicate data file."); + } + PriorityQueue nearest = + new PriorityQueue<>(limit, BEST_FIRST.reversed()); + Set 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 result = new ArrayList<>(nearest); + Collections.sort(result, BEST_FIRST); + return Collections.unmodifiableList(result); + } + + private static void add( + PriorityQueue 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 93341f3f1b6c..dd43a99af833 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.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 @@ private ScoredGlobalIndexResult readRawSearch( 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 @@ protected ScoredGlobalIndexResult scoreRawVectors( 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 @@ protected String rawSearchMetric(@Nullable GlobalIndexer globalIndexer) { if (metric == null) { metric = configuredRawSearchMetric(); } - return metric == null ? "l2" : normalizeMetric(metric); + return metric == null ? "l2" : VectorSearchMetric.normalize(metric); } @Nullable @@ -690,7 +698,7 @@ private String configuredRawSearchMetric(Map options) { for (Map.Entry 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 @@ private String configuredRawSearchMetric(Map options) { @Nullable private static String option(Map 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 @@ private static IndexFileMeta firstVectorIndexFile(List s 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> emptyOptionalResults(int n) { List> 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 000000000000..e13e7c66e257 --- /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 1a30e207f4bd..0f7557c848cd 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 @@ void testBuildsAndSearchesMultiSourceSegment() throws Exception { deletionVectors.put("data-2", data2Deletes); ExecutorService executor = Executors.newSingleThreadExecutor(); - List candidates; + List candidates; try { candidates = new PkVectorAnnSegmentSearcher( @@ -120,9 +120,7 @@ fileIO, annFile, vectorField(), indexOptions(), "l2", executor) } 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 000000000000..8f7e14fe0d84 --- /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 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 000000000000..5e99f96debd3 --- /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 deletionVectors = Collections.emptyMap(); + Map 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 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 deletionVectors = new HashMap<>(); + deletionVectors.put("data-1", data1Deletes); + + List 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); + } +}