replace the check from dfs phase to KnnVectorQueryBuilder doToQuery

Author: weizijun

server/src/main/java/org/elasticsearch/search/dfs/DfsPhase.java

@@ -44,7 +44,6 @@
 import java.util.Map;
 
 import static org.elasticsearch.index.query.AbstractQueryBuilder.DEFAULT_BOOST;
-import static org.elasticsearch.search.vectors.KnnVectorQueryBuilder.NUM_CANDS_FIELD;
 
 /**
  * DFS phase of a search request, used to make scoring 100% accurate by collecting additional info from each shard before the query phase.
@@ -178,7 +177,7 @@ private static Timer maybeStartTimer(DfsProfiler profiler, DfsTimingType dtt) {
         return null;
     };
 
-    static void executeKnnVectorQuery(SearchContext context) throws IOException {
+    private static void executeKnnVectorQuery(SearchContext context) throws IOException {
         SearchSourceBuilder source = context.request().source();
         if (source == null || source.knnSearch().isEmpty()) {
             return;
@@ -187,15 +186,6 @@ static void executeKnnVectorQuery(SearchContext context) throws IOException {
         SearchExecutionContext searchExecutionContext = context.getSearchExecutionContext();
         List<KnnSearchBuilder> knnSearch = source.knnSearch();
         List<KnnVectorQueryBuilder> knnVectorQueryBuilders = knnSearch.stream().map(KnnSearchBuilder::toQueryBuilder).toList();
-        int maxKnnNumCandidates = context.indexShard().indexSettings().getMaxKnnNumCandidates();
-        for (KnnVectorQueryBuilder knnVectorQueryBuilder : knnVectorQueryBuilders) {
-            if (knnVectorQueryBuilder.numCands() != null && knnVectorQueryBuilder.numCands() > maxKnnNumCandidates) {
-                throw new IllegalArgumentException(
-                    "[" + NUM_CANDS_FIELD.getPreferredName() + "] cannot exceed [" + maxKnnNumCandidates + "]"
-                );
-            }
-        }
-
         // Since we apply boost during the DfsQueryPhase, we should not apply boost here:
         knnVectorQueryBuilders.forEach(knnVectorQueryBuilder -> knnVectorQueryBuilder.boost(DEFAULT_BOOST));
 

server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java

@@ -482,7 +482,6 @@ protected QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException {
 
     @Override
     protected Query doToQuery(SearchExecutionContext context) throws IOException {
-        MappedFieldType fieldType = context.getFieldType(fieldName);
         int k;
         if (this.k != null) {
             k = this.k;
@@ -492,7 +491,15 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
                 k = Math.min(k, numCands);
             }
         }
-        int adjustedNumCands = numCands == null ? Math.round(NUM_CANDS_MULTIPLICATIVE_FACTOR * k) : numCands;
+
+        int maxKnnNumCandidates = context.getIndexSettings().getMaxKnnNumCandidates();
+        if (numCands != null && numCands > maxKnnNumCandidates) {
+            throw new IllegalArgumentException("[" + NUM_CANDS_FIELD.getPreferredName() + "] cannot exceed [" + maxKnnNumCandidates + "]");
+        }
+
+        int adjustedNumCands = numCands == null ? Math.round(Math.min(NUM_CANDS_MULTIPLICATIVE_FACTOR * k, maxKnnNumCandidates)) : numCands;
+
+        MappedFieldType fieldType = context.getFieldType(fieldName);
         if (fieldType == null) {
             return new MatchNoDocsQuery();
         }

server/src/test/java/org/elasticsearch/search/dfs/DfsPhaseTests.java

@@ -17,20 +17,12 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
-import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.index.IndexSettings;
-import org.elasticsearch.index.shard.IndexShard;
-import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.ContextIndexSearcher;
-import org.elasticsearch.search.internal.SearchContext;
-import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.profile.Profilers;
 import org.elasticsearch.search.profile.SearchProfileDfsPhaseResult;
 import org.elasticsearch.search.profile.query.CollectorResult;
 import org.elasticsearch.search.profile.query.QueryProfileShardResult;
-import org.elasticsearch.search.vectors.KnnSearchBuilder;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.test.IndexSettingsModule;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.junit.After;
@@ -40,9 +32,6 @@
 import java.util.List;
 import java.util.concurrent.ThreadPoolExecutor;
 
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 public class DfsPhaseTests extends ESTestCase {
 
     ThreadPoolExecutor threadPoolExecutor;
@@ -113,35 +102,4 @@ public void testSingleKnnSearch() throws IOException {
             reader.close();
         }
     }
-
-    public void testNumCandidatesExceedsMax() {
-        Settings settings = Settings.builder().put("index.max_knn_num_candidates", 100).build();
-        IndexSettings indexSettings = IndexSettingsModule.newIndexSettings("test", settings);
-
-        SearchContext context = mock(SearchContext.class);
-        when(context.indexShard()).thenAnswer(invocation -> {
-            IndexShard mockIndexShard = mock(IndexShard.class);
-            when(mockIndexShard.indexSettings()).thenReturn(indexSettings);
-            return mockIndexShard;
-        });
-
-        KnnSearchBuilder queryBuilder = new KnnSearchBuilder(
-            "float_vector",
-            new float[] { 0, 0, 0 },
-            10,
-            150, // 超过maxKnnNumCandidates的值
-            null,
-            null
-        );
-        SearchSourceBuilder source = new SearchSourceBuilder();
-        source.knnSearch(List.of(queryBuilder));
-        ShardSearchRequest searchRequest = mock(ShardSearchRequest.class);
-        when(searchRequest.source()).thenReturn(source);
-        when(context.request()).thenReturn(searchRequest);
-
-        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> DfsPhase.executeKnnVectorQuery(context));
-        assertEquals("[num_candidates] cannot exceed [100]", e.getMessage());
-
-    }
-
 }

server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java

@@ -21,6 +21,8 @@
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
 import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 import org.elasticsearch.index.query.InnerHitsRewriteContext;
@@ -34,6 +36,7 @@
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.test.AbstractBuilderTestCase;
 import org.elasticsearch.test.AbstractQueryTestCase;
+import org.elasticsearch.test.IndexSettingsModule;
 import org.elasticsearch.test.TransportVersionUtils;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
@@ -53,6 +56,8 @@
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.nullValue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 abstract class AbstractKnnVectorQueryBuilderTestCase extends AbstractQueryTestCase<KnnVectorQueryBuilder> {
     private static final String VECTOR_FIELD = "vector";
@@ -458,4 +463,17 @@ public void testRewriteWithQueryVectorBuilder() throws Exception {
         assertThat(rewritten.filterQueries(), hasSize(numFilters));
         assertThat(rewritten.filterQueries(), equalTo(filters));
     }
+
+    public void testMaxNumCandidatesExceeded() {
+        Settings settings = Settings.builder().put("index.max_knn_num_candidates", 100).build();
+        IndexSettings indexSettings = IndexSettingsModule.newIndexSettings("test", settings);
+
+        SearchExecutionContext context = mock(SearchExecutionContext.class);
+        when(context.getIndexSettings()).thenReturn(indexSettings);
+
+        KnnVectorQueryBuilder query = new KnnVectorQueryBuilder(VECTOR_FIELD, new float[] { 1.0f, 2.0f, 3.0f }, 5, 150, null, null);
+
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> query.doToQuery(context));
+        assertThat(e.getMessage(), containsString("[num_candidates] cannot exceed [100]"));
+    }
 }