Feature/solr 17928 finalize pr

changelog/unreleased/SOLR-17928_added_efSearch_parameter.yml

@@ -0,0 +1,11 @@
+# See https://github.com/apache/solr/blob/main/dev-docs/changelog.adoc
+title: Added efSearch parameter to knn query, exposed efSearchScaleFactor that is used to calculate efSearch internally
+type: added
+authors:
+  - name: Puneet Ahuja
+  - name: Elia Porciani
+links:
+  - name: SOLR-17928
+    url: https://issues.apache.org/jira/browse/SOLR-17928
+issues:
+  - 17928

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

@@ -37,8 +37,6 @@
 import org.apache.lucene.queries.function.ValueSource;
 import org.apache.lucene.queries.function.valuesource.ByteKnnVectorFieldSource;
 import org.apache.lucene.queries.function.valuesource.FloatKnnVectorFieldSource;
-import org.apache.lucene.search.KnnByteVectorQuery;
-import org.apache.lucene.search.KnnFloatVectorQuery;
 import org.apache.lucene.search.PatienceKnnVectorQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.SeededKnnVectorQuery;
@@ -49,6 +47,8 @@
 import org.apache.solr.common.SolrException;
 import org.apache.solr.search.QParser;
 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.uninverting.UninvertingReader;
 import org.apache.solr.util.vector.ByteDenseVectorParser;
 import org.apache.solr.util.vector.DenseVectorParser;
@@ -502,6 +502,7 @@ public Query getKnnVectorQuery(
       String fieldName,
       String vectorToSearch,
       int topK,
+      int efSearch,
       Query filterQuery,
       Query seedQuery,
       EarlyTerminationParams earlyTermination,
@@ -516,22 +517,32 @@ public Query getKnnVectorQuery(
             if (filteredSearchThreshold != null) {
               KnnSearchStrategy knnSearchStrategy =
                   new KnnSearchStrategy.Hnsw(filteredSearchThreshold);
-              yield new KnnFloatVectorQuery(
-                  fieldName, vectorBuilder.getFloatVector(), topK, filterQuery, knnSearchStrategy);
+              yield new SolrKnnFloatVectorQuery(
+                  fieldName,
+                  vectorBuilder.getFloatVector(),
+                  topK,
+                  efSearch,
+                  filterQuery,
+                  knnSearchStrategy);
             } else {
-              yield new KnnFloatVectorQuery(
-                  fieldName, vectorBuilder.getFloatVector(), topK, filterQuery);
+              yield new SolrKnnFloatVectorQuery(
+                  fieldName, vectorBuilder.getFloatVector(), topK, efSearch, filterQuery);
             }
           }
           case BYTE -> {
             if (filteredSearchThreshold != null) {
               KnnSearchStrategy knnSearchStrategy =
                   new KnnSearchStrategy.Hnsw(filteredSearchThreshold);
-              yield new KnnByteVectorQuery(
-                  fieldName, vectorBuilder.getByteVector(), topK, filterQuery, knnSearchStrategy);
+              yield new SolrKnnByteVectorQuery(
+                  fieldName,
+                  vectorBuilder.getByteVector(),
+                  topK,
+                  efSearch,
+                  filterQuery,
+                  knnSearchStrategy);
             } else {
-              yield new KnnByteVectorQuery(
-                  fieldName, vectorBuilder.getByteVector(), topK, filterQuery);
+              yield new SolrKnnByteVectorQuery(
+                  fieldName, vectorBuilder.getByteVector(), topK, efSearch, filterQuery);
             }
           }
         };
@@ -586,9 +597,9 @@ public SortField getSortField(SchemaField field, boolean top) {
 
   private Query getSeededQuery(Query knnQuery, Query seed) {
     return switch (knnQuery) {
-      case KnnFloatVectorQuery knnFloatQuery -> SeededKnnVectorQuery.fromFloatQuery(
+      case SolrKnnFloatVectorQuery knnFloatQuery -> SeededKnnVectorQuery.fromFloatQuery(
           knnFloatQuery, seed);
-      case KnnByteVectorQuery knnByteQuery -> SeededKnnVectorQuery.fromByteQuery(
+      case SolrKnnByteVectorQuery knnByteQuery -> SeededKnnVectorQuery.fromByteQuery(
           knnByteQuery, seed);
       default -> throw new SolrException(
           SolrException.ErrorCode.SERVER_ERROR, "Invalid type of knn query");
@@ -600,13 +611,13 @@ private Query getEarlyTerminationQuery(Query knnQuery, EarlyTerminationParams ea
         (earlyTermination.getSaturationThreshold() != null
             && earlyTermination.getPatience() != null);
     return switch (knnQuery) {
-      case KnnFloatVectorQuery knnFloatQuery -> useExplicitParams
+      case SolrKnnFloatVectorQuery knnFloatQuery -> useExplicitParams
           ? PatienceKnnVectorQuery.fromFloatQuery(
               knnFloatQuery,
               earlyTermination.getSaturationThreshold(),
               earlyTermination.getPatience())
           : PatienceKnnVectorQuery.fromFloatQuery(knnFloatQuery);
-      case KnnByteVectorQuery knnByteQuery -> useExplicitParams
+      case SolrKnnByteVectorQuery knnByteQuery -> useExplicitParams
           ? PatienceKnnVectorQuery.fromByteQuery(
               knnByteQuery,
               earlyTermination.getSaturationThreshold(),

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

@@ -108,12 +108,22 @@ public Query parse() throws SyntaxError {
     final DenseVectorField denseVectorType = getCheckedFieldType(schemaField);
     final String vectorToSearch = getVectorToSearch();
     final int topK = localParams.getInt(TOP_K, DEFAULT_TOP_K);
+
+    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 Integer filteredSearchThreshold = localParams.getInt(FILTERED_SEARCH_THRESHOLD);
 
     return denseVectorType.getKnnVectorQuery(
         schemaField.getName(),
         vectorToSearch,
         topK,
+        efSearch,
         getFilterQuery(),
         getSeedQuery(),
         getEarlyTerminationParams(),

solr/core/src/java/org/apache/solr/search/vector/SolrKnnByteVectorQuery.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 org.apache.lucene.search.KnnByteVectorQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
+public class SolrKnnByteVectorQuery extends KnnByteVectorQuery {
+  private final int topK;
+
+  public SolrKnnByteVectorQuery(String field, byte[] target, int topK, int efSearch, Query filter) {
+    // efSearch is used as 'k' to explore this many vectors in HNSW, then topK results are returned
+    // to the user
+    super(field, target, efSearch, filter);
+    this.topK = topK;
+  }
+
+  public SolrKnnByteVectorQuery(
+      String field,
+      byte[] target,
+      int topK,
+      int efSearch,
+      Query filter,
+      KnnSearchStrategy searchStrategy) {
+    super(field, target, efSearch, filter, searchStrategy);
+    this.topK = topK;
+  }
+
+  @Override
+  protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
+    return TopDocs.merge(topK, perLeafResults);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) return true;
+    if (!super.equals(obj)) return false;
+    if (getClass() != obj.getClass()) return false;
+    SolrKnnByteVectorQuery other = (SolrKnnByteVectorQuery) obj;
+    return this.topK == other.topK;
+  }
+
+  @Override
+  public int hashCode() {
+    return 31 * super.hashCode() + Integer.hashCode(topK);
+  }
+}

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

@@ -0,0 +1,64 @@
+/*
+ * 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 org.apache.lucene.search.KnnFloatVectorQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
+public class SolrKnnFloatVectorQuery extends KnnFloatVectorQuery {
+  private final int topK;
+
+  public SolrKnnFloatVectorQuery(
+      String field, float[] target, int topK, int efSearch, Query filter) {
+    // efSearch is used as 'k' to explore this many vectors in HNSW then topK results are returned
+    // to the user
+    super(field, target, efSearch, filter);
+    this.topK = topK;
+  }
+
+  public SolrKnnFloatVectorQuery(
+      String field,
+      float[] target,
+      int topK,
+      int efSearch,
+      Query filter,
+      KnnSearchStrategy searchStrategy) {
+    super(field, target, efSearch, filter, searchStrategy);
+    this.topK = topK;
+  }
+
+  @Override
+  protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
+    return TopDocs.merge(topK, perLeafResults);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) return true;
+    if (!super.equals(obj)) return false;
+    if (getClass() != obj.getClass()) return false;
+    SolrKnnFloatVectorQuery other = (SolrKnnFloatVectorQuery) obj;
+    return this.topK == other.topK;
+  }
+
+  @Override
+  public int hashCode() {
+    return 31 * super.hashCode() + Integer.hashCode(topK);
+  }
+}

solr/core/src/test/org/apache/solr/schema/DenseVectorFieldTest.java

@@ -868,7 +868,7 @@ public void testFilteredSearchThreshold_floatNoThresholdInInput_shouldSetDefault
       DenseVectorField type = (DenseVectorField) vectorField.getType();
       KnnFloatVectorQuery vectorQuery =
           (KnnFloatVectorQuery)
-              type.getKnnVectorQuery("vector", "[2, 1, 3, 4]", 3, null, null, null, null);
+              type.getKnnVectorQuery("vector", "[2, 1, 3, 4]", 3, 3, null, null, null, null);
       KnnSearchStrategy.Hnsw strategy = (KnnSearchStrategy.Hnsw) vectorQuery.getSearchStrategy();
       Integer threshold = strategy.filteredSearchThreshold();
 
@@ -892,7 +892,7 @@ public void testFilteredSearchThreshold_floatThresholdInInput_shouldSetCustomThr
       KnnFloatVectorQuery vectorQuery =
           (KnnFloatVectorQuery)
               type.getKnnVectorQuery(
-                  "vector", "[2, 1, 3, 4]", 3, null, null, null, expectedThreshold);
+                  "vector", "[2, 1, 3, 4]", 3, 3, null, null, null, expectedThreshold);
       KnnSearchStrategy.Hnsw strategy = (KnnSearchStrategy.Hnsw) vectorQuery.getSearchStrategy();
       Integer threshold = strategy.filteredSearchThreshold();
 
@@ -917,7 +917,7 @@ public void testFilteredSearchThreshold_seededFloatThresholdInInput_shouldSetCus
       SeededKnnVectorQuery vectorQuery =
           (SeededKnnVectorQuery)
               type.getKnnVectorQuery(
-                  "vector", "[2, 1, 3, 4]", 3, null, seedQuery, null, expectedThreshold);
+                  "vector", "[2, 1, 3, 4]", 3, 3, null, seedQuery, null, expectedThreshold);
       KnnSearchStrategy.Hnsw strategy = (KnnSearchStrategy.Hnsw) vectorQuery.getSearchStrategy();
       Integer threshold = strategy.filteredSearchThreshold();
 
@@ -944,7 +944,7 @@ public void testFilteredSearchThreshold_seededFloatThresholdInInput_shouldSetCus
       PatienceKnnVectorQuery vectorQuery =
           (PatienceKnnVectorQuery)
               type.getKnnVectorQuery(
-                  "vector", "[2, 1, 3, 4]", 3, null, null, earlyTermination, expectedThreshold);
+                  "vector", "[2, 1, 3, 4]", 3, 3, null, null, earlyTermination, expectedThreshold);
       KnnSearchStrategy.Hnsw strategy = (KnnSearchStrategy.Hnsw) vectorQuery.getSearchStrategy();
       Integer threshold = strategy.filteredSearchThreshold();
 
@@ -975,6 +975,7 @@ public void testFilteredSearchThreshold_seededFloatThresholdInInput_shouldSetCus
                   "vector",
                   "[2, 1, 3, 4]",
                   3,
+                  3,
                   null,
                   seedQuery,
                   earlyTermination,
@@ -1002,7 +1003,7 @@ public void testFilteredSearchThreshold_byteNoThresholdInInput_shouldSetDefaultT
       KnnByteVectorQuery vectorQuery =
           (KnnByteVectorQuery)
               type.getKnnVectorQuery(
-                  "vector_byte_encoding", "[2, 1, 3, 4]", 3, null, null, null, null);
+                  "vector_byte_encoding", "[2, 1, 3, 4]", 3, 3, null, null, null, null);
       KnnSearchStrategy.Hnsw strategy = (KnnSearchStrategy.Hnsw) vectorQuery.getSearchStrategy();
       Integer threshold = strategy.filteredSearchThreshold();
 
@@ -1026,7 +1027,14 @@ public void testFilteredSearchThreshold_byteThresholdInInput_shouldSetCustomThre
       KnnByteVectorQuery vectorQuery =
           (KnnByteVectorQuery)
               type.getKnnVectorQuery(
-                  "vector_byte_encoding", "[2, 1, 3, 4]", 3, null, null, null, expectedThreshold);
+                  "vector_byte_encoding",
+                  "[2, 1, 3, 4]",
+                  3,
+                  3,
+                  null,
+                  null,
+                  null,
+                  expectedThreshold);
       KnnSearchStrategy.Hnsw strategy = (KnnSearchStrategy.Hnsw) vectorQuery.getSearchStrategy();
       Integer threshold = strategy.filteredSearchThreshold();
 
@@ -1054,6 +1062,7 @@ public void testFilteredSearchThreshold_seededByteThresholdInInput_shouldSetCust
                   "vector_byte_encoding",
                   "[2, 1, 3, 4]",
                   3,
+                  3,
                   null,
                   seedQuery,
                   null,
@@ -1087,6 +1096,7 @@ public void testFilteredSearchThreshold_seededByteThresholdInInput_shouldSetCust
                   "vector_byte_encoding",
                   "[2, 1, 3, 4]",
                   3,
+                  3,
                   null,
                   null,
                   earlyTermination,
@@ -1121,6 +1131,7 @@ public void testFilteredSearchThreshold_seededByteThresholdInInput_shouldSetCust
                   "vector_byte_encoding",
                   "[2, 1, 3, 4]",
                   3,
+                  3,
                   null,
                   seedQuery,
                   earlyTermination,