Refactor InnerHitsRewriteContext into a more generic PerDocumentQueryRewriteContext

The DFS and highlight phases require rewriting the Lucene query outside of the query phase.
However, if the query contains a k-NN query, this triggers a nearest neighbor search on the entire shard,
which is unnecessary in these phases since computing top-N results is not required.

This change builds upon #104006, applying the same transformation used for nested inner hits.
As a result, DFS and highlight phases avoid wasting time and resources on costly nearest neighbor searches.

Note: The explain and matched query phases are also affected but still require the nearest neighbor search
for accurate results, so they remain unchanged for now.

Author: jimczi

modules/parent-join/src/main/java/org/elasticsearch/join/query/ParentChildInnerHitContextBuilder.java

@@ -70,6 +70,7 @@ protected void doBuild(SearchContext context, InnerHitsContext innerHitsContext)
             String name = innerHitBuilder.getName() != null ? innerHitBuilder.getName() : typeName;
             JoinFieldInnerHitSubContext joinFieldInnerHits = new JoinFieldInnerHitSubContext(
                 name,
+                query,
                 context,
                 typeName,
                 fetchChildInnerHits,
@@ -89,8 +90,15 @@ static final class JoinFieldInnerHitSubContext extends InnerHitsContext.InnerHit
         private final boolean fetchChildInnerHits;
         private final Joiner joiner;
 
-        JoinFieldInnerHitSubContext(String name, SearchContext context, String typeName, boolean fetchChildInnerHits, Joiner joiner) {
-            super(name, context);
+        JoinFieldInnerHitSubContext(
+            String name,
+            QueryBuilder query,
+            SearchContext context,
+            String typeName,
+            boolean fetchChildInnerHits,
+            Joiner joiner
+        ) {
+            super(name, query, context);
             this.typeName = typeName;
             this.fetchChildInnerHits = fetchChildInnerHits;
             this.joiner = joiner;

server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java

@@ -151,7 +151,7 @@ public QueryBuilder topDocsQuery() {
         @Override
         public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
             assertNull(ctx.getPointInTimeBuilder());
-            assertNull(ctx.convertToInnerHitsRewriteContext());
+            assertNull(ctx.convertToPerDocumentQueryRewriteContext());
             assertNull(ctx.convertToCoordinatorRewriteContext());
             assertNull(ctx.convertToIndexMetadataContext());
             assertNull(ctx.convertToSearchExecutionContext());
@@ -215,7 +215,7 @@ public QueryBuilder topDocsQuery() {
         @Override
         public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
             assertNotNull(ctx.getPointInTimeBuilder());
-            assertNull(ctx.convertToInnerHitsRewriteContext());
+            assertNull(ctx.convertToPerDocumentQueryRewriteContext());
             assertNull(ctx.convertToCoordinatorRewriteContext());
             assertNull(ctx.convertToIndexMetadataContext());
             assertNull(ctx.convertToSearchExecutionContext());

server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java

@@ -306,7 +306,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
         if (queryRewriteContext == null) {
             return this;
         }
-        final InnerHitsRewriteContext ihrc = queryRewriteContext.convertToInnerHitsRewriteContext();
+        final PerDocumentQueryRewriteContext ihrc = queryRewriteContext.convertToPerDocumentQueryRewriteContext();
         if (ihrc != null) {
             return doInnerHitsRewrite(ihrc);
         }
@@ -358,11 +358,11 @@ protected QueryBuilder doIndexMetadataRewrite(final QueryRewriteContext context)
 
     /**
      * Optional rewrite logic that allows for optimization for extracting inner hits
-     * @param context an {@link InnerHitsRewriteContext} instance
+     * @param context an {@link PerDocumentQueryRewriteContext} instance
      * @return A {@link QueryBuilder} representing the rewritten query optimized for inner hit extraction
      * @throws IOException if an error occurs while rewriting the query
      */
-    protected QueryBuilder doInnerHitsRewrite(final InnerHitsRewriteContext context) throws IOException {
+    protected QueryBuilder doInnerHitsRewrite(final PerDocumentQueryRewriteContext context) throws IOException {
         return this;
     }
 

server/src/main/java/org/elasticsearch/index/query/InnerHitsRewriteContext.java

@@ -380,8 +380,16 @@ protected void doBuild(SearchContext parentSearchContext, InnerHitsContext inner
             }
             String name = innerHitBuilder.getName() != null ? innerHitBuilder.getName() : nestedMapper.fullPath();
             NestedObjectMapper parentObjectMapper = searchExecutionContext.nestedScope().nextLevel(nestedMapper);
+
+            // We rewrite the query to ensure that child documents can receive scores from approximate nearest neighbor queries.
+            PerDocumentQueryRewriteContext innerHitsRewriteContext = new PerDocumentQueryRewriteContext(
+                searchExecutionContext.getParserConfig(),
+                searchExecutionContext.nowInMillis
+            );
+            var childQuery = Rewriteable.rewrite(query, innerHitsRewriteContext, true);
             NestedInnerHitSubContext nestedInnerHits = new NestedInnerHitSubContext(
                 name,
+                childQuery,
                 parentSearchContext,
                 parentObjectMapper,
                 nestedMapper
@@ -399,11 +407,12 @@ static final class NestedInnerHitSubContext extends InnerHitsContext.InnerHitSub
 
         NestedInnerHitSubContext(
             String name,
+            QueryBuilder query,
             SearchContext context,
             NestedObjectMapper parentObjectMapper,
             NestedObjectMapper childObjectMapper
         ) {
-            super(name, context);
+            super(name, query, context);
             this.parentObjectMapper = parentObjectMapper;
             this.childObjectMapper = childObjectMapper;
         }

server/src/main/java/org/elasticsearch/index/query/NestedQueryBuilder.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+package org.elasticsearch.index.query;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+
+import java.util.function.LongSupplier;
+
+/**
+ * A context object for rewriting {@link QueryBuilder} instances to optimize query execution
+ * by reducing the global operations required for processing Lucene queries on a selection of top documents.
+ *
+ * This rewrite context ensures that expensive operations, such as approximate nearest neighbor (ANN) searches,
+ * are avoided when evaluating documents in isolation. Instead, ANN queries are rewritten into exact k-NN queries,
+ * which execute efficiently when processing only a small subset of documents (or none at all).
+ *
+ * Additionally, this context is beneficial for nested inner hits, allowing child documents to receive
+ * independent similarity scores regardless of whether they matched the approximate nearest neighbor query.
+ */
+public final class PerDocumentQueryRewriteContext extends QueryRewriteContext {
+    public PerDocumentQueryRewriteContext(final XContentParserConfiguration parserConfiguration, final LongSupplier nowInMillis) {
+        super(parserConfiguration, null, nowInMillis);
+    }
+
+    @Override
+    public PerDocumentQueryRewriteContext convertToPerDocumentQueryRewriteContext() {
+        return this;
+    }
+
+    @Override
+    public void executeAsyncActions(ActionListener<Void> listener) {
+        // PerDocumentQueryRewriteContext does not support async actions at all, and doesn't supply a valid `client` object
+        throw new UnsupportedOperationException("PerDocumentQueryRewriteContext does not support async actions");
+    }
+
+}

server/src/main/java/org/elasticsearch/index/query/PerDocumentQueryRewriteContext.java

@@ -223,7 +223,7 @@ public DataRewriteContext convertToDataRewriteContext() {
         return null;
     }
 
-    public InnerHitsRewriteContext convertToInnerHitsRewriteContext() {
+    public PerDocumentQueryRewriteContext convertToPerDocumentQueryRewriteContext() {
         return null;
     }
 

server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -45,6 +45,7 @@
 import org.elasticsearch.index.mapper.NestedLookup;
 import org.elasticsearch.index.mapper.SourceLoader;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.ParsedQuery;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.SearchExecutionContext;
@@ -486,6 +487,11 @@ public String source() {
         return "search";
     }
 
+    @Override
+    public QueryBuilder userQueryBuilder() {
+        return request.source() != null ? request.source().query() : new MatchAllQueryBuilder();
+    }
+
     @Override
     public ShardSearchRequest request() {
         return this.request;

server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java

@@ -60,7 +60,6 @@
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.index.query.CoordinatorRewriteContextProvider;
 import org.elasticsearch.index.query.InnerHitContextBuilder;
-import org.elasticsearch.index.query.InnerHitsRewriteContext;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -1341,22 +1340,16 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         context.size(source.size());
         Map<String, InnerHitContextBuilder> innerHitBuilders = new HashMap<>();
         QueryBuilder query = source.query();
-        InnerHitsRewriteContext innerHitsRewriteContext = new InnerHitsRewriteContext(
-            context.getSearchExecutionContext().getParserConfig(),
-            context::getRelativeTimeInMillis
-        );
         if (query != null) {
-            QueryBuilder rewrittenForInnerHits = Rewriteable.rewrite(query, innerHitsRewriteContext, true);
             if (false == source.skipInnerHits()) {
-                InnerHitContextBuilder.extractInnerHits(rewrittenForInnerHits, innerHitBuilders);
+                InnerHitContextBuilder.extractInnerHits(query, innerHitBuilders);
             }
             searchExecutionContext.setAliasFilter(context.request().getAliasFilter().getQueryBuilder());
             context.parsedQuery(searchExecutionContext.toQuery(query));
         }
         if (source.postFilter() != null) {
-            QueryBuilder rewrittenForInnerHits = Rewriteable.rewrite(source.postFilter(), innerHitsRewriteContext, true);
             if (false == source.skipInnerHits()) {
-                InnerHitContextBuilder.extractInnerHits(rewrittenForInnerHits, innerHitBuilders);
+                InnerHitContextBuilder.extractInnerHits(query, innerHitBuilders);
             }
             context.parsedPostFilter(searchExecutionContext.toQuery(source.postFilter()));
         }

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -20,6 +20,8 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopScoreDocCollectorManager;
 import org.elasticsearch.index.query.ParsedQuery;
+import org.elasticsearch.index.query.PerDocumentQueryRewriteContext;
+import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.ContextIndexSearcher;
@@ -119,9 +121,19 @@ public CollectionStatistics collectionStatistics(String field) throws IOExceptio
         }
 
         try {
+            // We apply the per-document rewrite context to prevent costly nearest neighbor searches
+            // during query rewriting. Since we are only extracting term statistics and not executing
+            // the query, these expensive operations are unnecessary.
+            PerDocumentQueryRewriteContext rewriteContext = new PerDocumentQueryRewriteContext(
+                context.getSearchExecutionContext().getParserConfig(),
+                context.getSearchExecutionContext()::nowInMillis
+            );
+            Query query = Rewriteable.rewrite(context.userQueryBuilder(), rewriteContext, true)
+                .toQuery(context.getSearchExecutionContext());
+
             Timer timer = maybeStartTimer(profiler, DfsTimingType.CREATE_WEIGHT);
             try {
-                searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
+                searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE, 1);
             } finally {
                 if (timer != null) {
                     timer.stop();