Add rerank to primary-key vector search

docs/docs/primary-key-table/vector-index.md

@@ -140,6 +140,38 @@ required. Batch writes which wait for compaction can publish the data and its in
 
 ## Search
 
+### Exact Rerank
+
+Primary-key vector search can retrieve more ANN candidates and rerank them with the original
+vectors stored in the table. For example, the following table option retrieves up to four times
+the requested Top-K from the `ivf-flat` index before computing exact distances:
+
+```sql
+'fields.embedding.ivf-flat.refine_factor' = '4'
+```
+
+The option is disabled by default. Its configuration semantics are the same as for a Data
+Evolution vector index:
+
+- `refine_factor`, `refine-factor`, `rerank_factor`, and `rerank-factor` are accepted.
+- A query option overrides every table option. Within either set of options, field and index
+  prefixes take precedence over less specific prefixes. For example,
+  `fields.embedding.ivf-flat.refine_factor` takes precedence over
+  `fields.embedding.ivf.refine_factor`, which takes precedence over `ivf.refine_factor` and then
+  `refine_factor`. The normalized underscore form of an index type, such as `ivf_flat`, is also
+  accepted after the configured index name.
+- The factor must be a positive integer. A factor of `1` performs exact reranking without
+  retrieving additional ANN candidates.
+
+Only candidates returned by ANN can win the rerank, so a larger factor can improve recall but does
+not guarantee the exact global Top-K. It also increases ANN work and data-file I/O. Files without
+an active ANN segment are already searched exactly and are kept separate from approximate
+candidates until the final Top-K merge.
+
+For distributed Spark searches, executors return bounded ANN and exact candidate streams. The
+driver globally merges the ANN candidates and rereads their original vectors for exact reranking;
+this does not start a second Spark job.
+
 ### Spark SQL
 
 Use the `vector_search` table-valued function. Spark exposes the ANN score through the

paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java

@@ -73,13 +73,36 @@ public List<PkVectorSearchResult> search(
             float[] query,
             int limit)
             throws IOException {
-        checkArgument(limit > 0, "Vector search limit must be positive.");
+        Result result = search(state, activeFiles, deletionVectors, query, limit, limit);
+        PriorityQueue<PkVectorSearchResult> nearest =
+                new PriorityQueue<>(limit, BEST_FIRST.reversed());
+        for (PkVectorSearchResult candidate : result.indexedCandidates) {
+            add(nearest, candidate, limit);
+        }
+        for (PkVectorSearchResult candidate : result.exactCandidates) {
+            add(nearest, candidate, limit);
+        }
+        return sorted(nearest);
+    }
+
+    public Result search(
+            PkVectorBucketIndexState state,
+            List<DataFileMeta> activeFiles,
+            Map<String, DeletionVector> deletionVectors,
+            float[] query,
+            int indexedLimit,
+            int exactLimit)
+            throws IOException {
+        checkArgument(indexedLimit > 0, "Vector indexed search limit must be positive.");
+        checkArgument(exactLimit > 0, "Vector exact search limit must be positive.");
         Map<String, DataFileMeta> filesByName = new HashMap<>();
         for (DataFileMeta file : activeFiles) {
             checkArgument(filesByName.put(file.fileName(), file) == null, "Duplicate data file.");
         }
-        PriorityQueue<PkVectorSearchResult> nearest =
-                new PriorityQueue<>(limit, BEST_FIRST.reversed());
+        PriorityQueue<PkVectorSearchResult> indexedNearest =
+                new PriorityQueue<>(indexedLimit, BEST_FIRST.reversed());
+        PriorityQueue<PkVectorSearchResult> exactNearest =
+                new PriorityQueue<>(exactLimit, BEST_FIRST.reversed());
         Set<String> activeSourceFiles = new HashSet<>(filesByName.keySet());
         Set<String> covered = new HashSet<>();
         for (IndexFileMeta ann : state.annSegments()) {
@@ -101,11 +124,11 @@ public List<PkVectorSearchResult> search(
                             ann,
                             sourceMeta,
                             query,
-                            limit,
+                            indexedLimit,
                             deletionVectors,
                             activeSourceFiles,
                             searchOptions)) {
-                add(nearest, result, limit);
+                add(indexedNearest, result, indexedLimit);
             }
         }
 
@@ -119,12 +142,16 @@ public List<PkVectorSearchResult> search(
                 try (PkVectorReader reader = vectorReaderFactory.create(file)) {
                     for (PkVectorSearchResult result :
                             PkVectorExactSearcher.search(
-                                    file.fileName(), reader, query, metric, limit, excluded)) {
-                        add(nearest, result, limit);
+                                    file.fileName(), reader, query, metric, exactLimit, excluded)) {
+                        add(exactNearest, result, exactLimit);
                     }
                 }
             }
         }
+        return new Result(sorted(indexedNearest), sorted(exactNearest));
+    }
+
+    private static List<PkVectorSearchResult> sorted(PriorityQueue<PkVectorSearchResult> nearest) {
         List<PkVectorSearchResult> result = new ArrayList<>(nearest);
         Collections.sort(result, BEST_FIRST);
         return Collections.unmodifiableList(result);
@@ -141,4 +168,26 @@ private static void add(
             nearest.add(candidate);
         }
     }
+
+    /** Separately bounded approximate-index and exact-fallback candidates. */
+    public static class Result {
+
+        private final List<PkVectorSearchResult> indexedCandidates;
+        private final List<PkVectorSearchResult> exactCandidates;
+
+        private Result(
+                List<PkVectorSearchResult> indexedCandidates,
+                List<PkVectorSearchResult> exactCandidates) {
+            this.indexedCandidates = indexedCandidates;
+            this.exactCandidates = exactCandidates;
+        }
+
+        public List<PkVectorSearchResult> indexedCandidates() {
+            return indexedCandidates;
+        }
+
+        public List<PkVectorSearchResult> exactCandidates() {
+            return exactCandidates;
+        }
+    }
 }

paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java

@@ -301,17 +301,7 @@ protected ScoredGlobalIndexResult withRawSearch(
     }
 
     protected int indexedSearchLimit(String indexType) {
-        int refineFactor = configuredRefineFactor(indexType);
-        if (refineFactor == 0) {
-            return limit;
-        }
-        if (limit > Integer.MAX_VALUE / refineFactor) {
-            throw new IllegalArgumentException(
-                    String.format(
-                            "Vector search limit overflow: limit=%d, refine factor=%d",
-                            limit, refineFactor));
-        }
-        return limit * refineFactor;
+        return VectorSearchRefineOptions.searchLimit(limit, configuredRefineFactor(indexType));
     }
 
     protected ScoredGlobalIndexResult maybeRerankIndexedResult(
@@ -721,73 +711,11 @@ private static boolean isRawSearchMetric(String metric) {
     }
 
     protected int configuredRefineFactor(String indexType) {
-        String value = configuredRefineFactor(options, indexType);
-        if (value == null) {
-            value = configuredRefineFactor(table.options(), indexType);
-        }
-        if (value == null) {
-            return 0;
-        }
-        try {
-            int factor = Integer.parseInt(value);
-            if (factor <= 0) {
-                throw new IllegalArgumentException(
-                        "Vector refine factor must be positive, got: " + value);
-            }
-            return factor;
-        } catch (NumberFormatException e) {
-            throw new IllegalArgumentException(
-                    "Invalid vector refine factor: " + value + ". Must be an integer.", e);
-        }
-    }
-
-    @Nullable
-    private String configuredRefineFactor(Map<String, String> options, String indexType) {
-        List<String> prefixes = new ArrayList<>();
-        String fieldPrefix = "fields." + vectorColumn.name() + ".";
-        addRefinePrefixes(prefixes, fieldPrefix, indexType);
-        addRefinePrefixes(prefixes, "", indexType);
-
-        for (String prefix : prefixes) {
-            String value = refineFactorOption(options, prefix + "refine_factor");
-            if (value == null) {
-                value = refineFactorOption(options, prefix + "refine-factor");
-            }
-            if (value == null) {
-                value = refineFactorOption(options, prefix + "rerank_factor");
-            }
-            if (value == null) {
-                value = refineFactorOption(options, prefix + "rerank-factor");
-            }
-            if (value != null) {
-                return value;
-            }
-        }
-        return null;
-    }
-
-    private static void addRefinePrefixes(List<String> prefixes, String base, String indexType) {
-        if (indexType != null && !indexType.isEmpty()) {
-            prefixes.add(base + indexType + ".");
-            String normalizedIndexType = normalizeIndexType(indexType);
-            if (!normalizedIndexType.equals(indexType)) {
-                prefixes.add(base + normalizedIndexType + ".");
-            }
-            if (normalizedIndexType.startsWith("ivf")) {
-                prefixes.add(base + "ivf.");
-            }
-        }
-        prefixes.add(base);
-    }
-
-    @Nullable
-    private static String refineFactorOption(Map<String, String> options, String key) {
-        String value = options.get(key);
-        return value == null ? null : value.trim();
-    }
-
-    private static String normalizeIndexType(String indexType) {
-        return indexType.toLowerCase().replace('-', '_');
+        return VectorSearchRefineOptions.resolve(
+                options,
+                table == null ? Collections.emptyMap() : table.options(),
+                vectorColumn.name(),
+                indexType);
     }
 
     private static IndexFileMeta firstVectorIndexFile(List<IndexVectorSearchSplit> splits) {