[core] Add primary-key vector bucket search

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);
+    }
+}

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<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 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 List<Candidate> search(
         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 List<Candidate> 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<Candidate> search(
             }
 
             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 List<Candidate> 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<PkVectorSourceFile> 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;

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);
+    }
+}

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;
+    }
+}