SOLR-18424: add oversampling + raw dense vector reranking

changelog/unreleased/SOLR-18424.yml

@@ -0,0 +1,11 @@
+title: >
+  Added `rerankOversample` parameter to the `{!knn}` query parser. The approximate vector search
+  collects `topK * rerankOversample` candidates, which are then rescored against the raw
+  full-precision vectors and reranked down to `topK` results. Recommended for quantized dense
+  vector fields, particularly binary quantized ones, where it recovers recall lost to quantization.
+type: added
+authors:
+  - name: Kevin Liang
+links:
+  - name: SOLR-18424
+    url: https://issues.apache.org/jira/browse/SOLR-18424

solr/core/src/java/org/apache/solr/schema/DenseVectorField.java

@@ -38,6 +38,7 @@
 import org.apache.lucene.queries.function.valuesource.ByteKnnVectorFieldSource;
 import org.apache.lucene.queries.function.valuesource.FloatKnnVectorFieldSource;
 import org.apache.lucene.search.FieldExistsQuery;
+import org.apache.lucene.search.FullPrecisionFloatVectorSimilarityValuesSource;
 import org.apache.lucene.search.PatienceKnnVectorQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.SeededKnnVectorQuery;
@@ -51,6 +52,7 @@
 import org.apache.solr.search.vector.KnnQParser.EarlyTerminationParams;
 import org.apache.solr.search.vector.SolrKnnByteVectorQuery;
 import org.apache.solr.search.vector.SolrKnnFloatVectorQuery;
+import org.apache.solr.search.vector.SolrRescoreTopNQuery;
 import org.apache.solr.uninverting.UninvertingReader;
 import org.apache.solr.util.vector.ByteDenseVectorParser;
 import org.apache.solr.util.vector.DenseVectorParser;
@@ -490,6 +492,24 @@ public ValueSource getValueSource(SchemaField field, QParser parser) {
         SolrException.ErrorCode.BAD_REQUEST, "Vector encoding not supported for function queries.");
   }
 
+  /**
+   * Re-ranking compares candidates against the raw full precision vectors, which is only possible
+   * for FLOAT32 encoded fields: a BYTE encoded field has no higher precision representation to
+   * re-rank against.
+   *
+   * @throws SolrException if oversampling was requested for a field that cannot support it
+   */
+  public void checkRerankOversampleSupported(String fieldName, int rerankOversample) {
+    if (rerankOversample > 1 && vectorEncoding != VectorEncoding.FLOAT32) {
+      throw new SolrException(
+          SolrException.ErrorCode.BAD_REQUEST,
+          "rerankOversample is only supported for FLOAT32 vector encoding; field '"
+              + fieldName
+              + "' uses "
+              + vectorEncoding);
+    }
+  }
+
   public Query getKnnVectorQuery(
       String fieldName,
       String vectorToSearch,
@@ -499,6 +519,28 @@ public Query getKnnVectorQuery(
       Query seedQuery,
       EarlyTerminationParams earlyTermination,
       Integer filteredSearchThreshold) {
+    return getKnnVectorQuery(
+        fieldName,
+        vectorToSearch,
+        topK,
+        efSearch,
+        filterQuery,
+        seedQuery,
+        earlyTermination,
+        filteredSearchThreshold,
+        1);
+  }
+
+  public Query getKnnVectorQuery(
+      String fieldName,
+      String vectorToSearch,
+      int topK,
+      int efSearch,
+      Query filterQuery,
+      Query seedQuery,
+      EarlyTerminationParams earlyTermination,
+      Integer filteredSearchThreshold,
+      int rerankOversample) {
 
     if (FLAT_ALGORITHM.equals(knnAlgorithm)) {
       throw new SolrException(
@@ -507,9 +549,13 @@ public Query getKnnVectorQuery(
               + "Use vectorSimilarity() function queries instead.");
     }
 
+    checkRerankOversampleSupported(fieldName, rerankOversample);
+
     DenseVectorParser vectorBuilder =
         getVectorBuilder(vectorToSearch, DenseVectorParser.BuilderPhase.QUERY);
 
+    final int candidateTopK = topK * rerankOversample;
+
     // Create KnnSearchStrategy if filteredSearchThreshold is provided
     KnnSearchStrategy searchStrategy = null;
     if (filteredSearchThreshold != null) {
@@ -524,25 +570,29 @@ public Query getKnnVectorQuery(
                 ? new SolrKnnFloatVectorQuery(
                     fieldName,
                     vectorBuilder.getFloatVector(),
-                    topK,
+                    candidateTopK,
                     efSearch,
                     filterQuery,
                     searchStrategy)
                 : new SolrKnnFloatVectorQuery(
-                    fieldName, vectorBuilder.getFloatVector(), topK, efSearch, filterQuery);
+                    fieldName,
+                    vectorBuilder.getFloatVector(),
+                    candidateTopK,
+                    efSearch,
+                    filterQuery);
         break;
       case BYTE:
         baseQuery =
             searchStrategy != null
                 ? new SolrKnnByteVectorQuery(
                     fieldName,
                     vectorBuilder.getByteVector(),
-                    topK,
+                    candidateTopK,
                     efSearch,
                     filterQuery,
                     searchStrategy)
                 : new SolrKnnByteVectorQuery(
-                    fieldName, vectorBuilder.getByteVector(), topK, efSearch, filterQuery);
+                    fieldName, vectorBuilder.getByteVector(), candidateTopK, efSearch, filterQuery);
         break;
       default:
         throw new SolrException(
@@ -560,6 +610,20 @@ public Query getKnnVectorQuery(
       baseQuery = getEarlyTerminationQuery(baseQuery, earlyTermination);
     }
 
+    // Re-rank the oversampled candidates against the raw full precision vectors and keep topK.
+    if (rerankOversample > 1) {
+      // The similarity function is passed explicitly instead of using
+      // RescoreTopNQuery#createFullPrecisionRescorerQuery, which leaves it null until the search
+      // resolves it lazily. That would make the query un-printable in the meantime, and
+      // debugQuery relies on Query#toString.
+      baseQuery =
+          new SolrRescoreTopNQuery(
+              baseQuery,
+              new FullPrecisionFloatVectorSimilarityValuesSource(
+                  vectorBuilder.getFloatVector(), fieldName, similarityFunction),
+              topK);
+    }
+
     return baseQuery;
   }
 

solr/core/src/java/org/apache/solr/search/vector/KnnQParser.java

@@ -20,6 +20,7 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.FullPrecisionFloatVectorSimilarityValuesSource;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.search.join.DiversifyingChildrenByteKnnVectorQuery;
@@ -43,6 +44,11 @@ public class KnnQParser extends AbstractVectorQParserBase {
   protected static final String SEED_QUERY = "seedQuery";
   protected static final String FILTERED_SEARCH_THRESHOLD = "filteredSearchThreshold";
 
+  // multiplier applied to topK to decide how many candidates to collect before re-ranking them
+  // down to topK results
+  protected static final String RERANK_OVERSAMPLE = "rerankOversample";
+  protected static final int DEFAULT_RERANK_OVERSAMPLE = 1;
+
   // parameters for PatienceKnnVectorQuery, a version of knn vector query that exits early when HNSW
   // queue saturates over a {@code #saturationThreshold} for more than {@code #patience} times.
   protected static final String EARLY_TERMINATION = "earlyTermination";
@@ -102,6 +108,16 @@ public EarlyTerminationParams getEarlyTerminationParams() {
     return new EarlyTerminationParams(enabled, saturationThreshold, patience);
   }
 
+  public int getRerankOversample() {
+    final int rerankOversample = localParams.getInt(RERANK_OVERSAMPLE, DEFAULT_RERANK_OVERSAMPLE);
+    if (rerankOversample < 1) {
+      throw new SolrException(
+          SolrException.ErrorCode.BAD_REQUEST,
+          "rerankOversample (" + rerankOversample + ") must be >= 1");
+    }
+    return rerankOversample;
+  }
+
   protected Query getSeedQuery() throws SolrException, SyntaxError {
     String seed = localParams.get(SEED_QUERY);
     if (seed == null) return null;
@@ -129,14 +145,24 @@ public Query parse() throws SyntaxError {
 
     final String vectorToSearch = getVectorToSearch();
     final int topK = localParams.getInt(TOP_K, DEFAULT_TOP_K);
+    final int rerankOversample = getRerankOversample();
+
+    final int candidateTopK;
+    try {
+      candidateTopK = Math.multiplyExact(topK, rerankOversample);
+    } catch (ArithmeticException e) {
+      throw new SolrException(
+          SolrException.ErrorCode.BAD_REQUEST,
+          "topK (" + topK + ") * rerankOversample (" + rerankOversample + ") overflows an integer");
+    }
 
     final double efSearchScaleFactor = localParams.getDouble("efSearchScaleFactor", 1.0);
     if (Double.isNaN(efSearchScaleFactor) || efSearchScaleFactor < 1.0) {
       throw new SolrException(
           SolrException.ErrorCode.BAD_REQUEST,
           "efSearchScaleFactor (" + efSearchScaleFactor + ") must be >= 1.0");
     }
-    final int efSearch = (int) Math.round(efSearchScaleFactor * topK);
+    final int efSearch = (int) Math.round(efSearchScaleFactor * candidateTopK);
 
     final Integer filteredSearchThreshold = localParams.getInt(FILTERED_SEARCH_THRESHOLD);
 
@@ -161,19 +187,36 @@ public Query parse() throws SyntaxError {
               req, subQuery(allParentsQuery, null).getQuery());
       final BooleanQuery acceptedParents = getParentsFilter(parentsFilterQueries);
 
+      denseVectorType.checkRerankOversampleSupported(vectorField, rerankOversample);
+
       Query acceptedChildren =
           getChildrenFilter(getFilterQuery(), acceptedParents, allParentsBitSet);
       switch (vectorEncoding) {
         case FLOAT32:
-          return new DiversifyingChildrenFloatKnnVectorQuery(
+          // The diversifying query returns the best matching child per parent, so collecting
+          // candidateTopK of them and re-ranking down to topK only ever narrows an already
+          // diversified set: at most one child per parent is preserved. Note that which child
+          // represents a parent is still picked using the (possibly quantized) approximate score,
+          // re-ranking only reorders the representatives that were chosen.
+          final float[] target = vectorBuilder.getFloatVector();
+          final Query diversified =
+              new DiversifyingChildrenFloatKnnVectorQuery(
+                  vectorField, target, acceptedChildren, candidateTopK, allParentsBitSet);
+          if (rerankOversample <= 1) {
+            return diversified;
+          }
+          return new SolrRescoreTopNQuery(
+              diversified,
+              new FullPrecisionFloatVectorSimilarityValuesSource(
+                  target, vectorField, denseVectorType.getSimilarityFunction()),
+              topK);
+        case BYTE:
+          return new DiversifyingChildrenByteKnnVectorQuery(
               vectorField,
-              vectorBuilder.getFloatVector(),
+              vectorBuilder.getByteVector(),
               acceptedChildren,
-              topK,
+              candidateTopK,
               allParentsBitSet);
-        case BYTE:
-          return new DiversifyingChildrenByteKnnVectorQuery(
-              vectorField, vectorBuilder.getByteVector(), acceptedChildren, topK, allParentsBitSet);
         default:
           throw new SolrException(
               SolrException.ErrorCode.SERVER_ERROR,
@@ -189,7 +232,8 @@ public Query parse() throws SyntaxError {
         getFilterQuery(),
         getSeedQuery(),
         getEarlyTerminationParams(),
-        filteredSearchThreshold);
+        filteredSearchThreshold,
+        rerankOversample);
   }
 
   private BooleanQuery getParentsFilter(String[] parentsFilterQueries) throws SyntaxError {

solr/core/src/java/org/apache/solr/search/vector/SolrRescoreTopNQuery.java

@@ -0,0 +1,63 @@
+/*
+ * 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.solr.search.vector;
+
+import java.io.IOException;
+import org.apache.lucene.search.DoubleValuesSource;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchNoDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.RescoreTopNQuery;
+
+/**
+ * A {@link RescoreTopNQuery} that tolerates an inner query matching no documents.
+ *
+ * <p>{@link RescoreTopNQuery#rewrite} unconditionally hands its collected hits to {@code
+ * DocAndScoreQuery#createDocAndScoreQuery}, which requires at least one hit: it asserts as much,
+ * and without assertions enabled it reads element zero of an empty array. {@code
+ * AbstractKnnVectorQuery#rewrite} guards the very same call by returning {@link MatchNoDocsQuery},
+ * but the re-ranking query has no equivalent guard, so an oversampled knn query that matches
+ * nothing (a restrictive {@code preFilter}, say) fails instead of returning no results.
+ *
+ * <p>TODO: remove this class once the guard is added upstream in Lucene and Solr picks up a release
+ * containing it.
+ */
+public class SolrRescoreTopNQuery extends RescoreTopNQuery {
+
+  private final Query innerQuery;
+  private final DoubleValuesSource valuesSource;
+  private final int n;
+
+  public SolrRescoreTopNQuery(Query query, DoubleValuesSource valuesSource, int n) {
+    super(query, valuesSource, n);
+    this.innerQuery = query;
+    this.valuesSource = valuesSource;
+    this.n = n;
+  }
+
+  @Override
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    final Query rewrittenInner = indexSearcher.rewrite(innerQuery);
+    if (rewrittenInner instanceof MatchNoDocsQuery) {
+      return rewrittenInner;
+    }
+    // Delegate using the already rewritten inner query rather than calling super.rewrite(), which
+    // would rewrite innerQuery a second time. That matters because rewriting a knn query runs the
+    // whole vector search; rewriting its result is a no-op.
+    return new RescoreTopNQuery(rewrittenInner, valuesSource, n).rewrite(indexSearcher);
+  }
+}