Avoid populating rank docs metadata if explain is not specified (#120536) (#120766)

pmpailis · web-flow · commit 9d5c474f93b2 · 2025-01-24T23:32:19.000+11:00
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -170,6 +170,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_ROLLOVER_LEGACY_INDICES = def(8_830_00_0);
     public static final TransportVersion ADD_INCLUDE_FAILURE_INDICES_OPTION = def(8_831_00_0);
     public static final TransportVersion ESQL_RESPONSE_PARTIAL = def(8_832_00_0);
+    public static final TransportVersion RANK_DOC_OPTIONAL_METADATA_FOR_EXPLAIN = def(8_833_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java
@@ -507,6 +507,7 @@ void executeRequest(
         });
 
         final SearchSourceBuilder source = original.source();
+        final boolean isExplain = source != null && source.explain() != null && source.explain();
         if (shouldOpenPIT(source)) {
             // disabling shard reordering for request
             original.setPreFilterShardSize(Integer.MAX_VALUE);
@@ -536,7 +537,12 @@ public void onFailure(Exception e) {
         } else {
             Rewriteable.rewriteAndFetch(
                 original,
-                searchService.getRewriteContext(timeProvider::absoluteStartMillis, resolvedIndices, original.pointInTimeBuilder()),
+                searchService.getRewriteContext(
+                    timeProvider::absoluteStartMillis,
+                    resolvedIndices,
+                    original.pointInTimeBuilder(),
+                    isExplain
+                ),
                 rewriteListener
             );
         }
diff --git a/server/src/main/java/org/elasticsearch/index/IndexService.java b/server/src/main/java/org/elasticsearch/index/IndexService.java
@@ -807,7 +807,8 @@ public QueryRewriteContext newQueryRewriteContext(
             scriptService,
             null,
             null,
-            null
+            null,
+            false
         );
     }
 
diff --git a/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java
@@ -105,7 +105,8 @@ public CoordinatorRewriteContext(
             null,
             null,
             null,
-            null
+            null,
+            false
         );
         this.dateFieldRangeInfo = dateFieldRangeInfo;
         this.tier = tier;
diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
@@ -72,6 +72,7 @@ public class QueryRewriteContext {
     private final ResolvedIndices resolvedIndices;
     private final PointInTimeBuilder pit;
     private QueryRewriteInterceptor queryRewriteInterceptor;
+    private final boolean isExplain;
 
     public QueryRewriteContext(
         final XContentParserConfiguration parserConfiguration,
@@ -89,7 +90,8 @@ public QueryRewriteContext(
         final ScriptCompiler scriptService,
         final ResolvedIndices resolvedIndices,
         final PointInTimeBuilder pit,
-        final QueryRewriteInterceptor queryRewriteInterceptor
+        final QueryRewriteInterceptor queryRewriteInterceptor,
+        final boolean isExplain
     ) {
 
         this.parserConfiguration = parserConfiguration;
@@ -109,6 +111,7 @@ public QueryRewriteContext(
         this.resolvedIndices = resolvedIndices;
         this.pit = pit;
         this.queryRewriteInterceptor = queryRewriteInterceptor;
+        this.isExplain = isExplain;
     }
 
     public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
@@ -128,7 +131,8 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration
             null,
             null,
             null,
-            null
+            null,
+            false
         );
     }
 
@@ -139,6 +143,18 @@ public QueryRewriteContext(
         final ResolvedIndices resolvedIndices,
         final PointInTimeBuilder pit,
         final QueryRewriteInterceptor queryRewriteInterceptor
+    ) {
+        this(parserConfiguration, client, nowInMillis, resolvedIndices, pit, queryRewriteInterceptor, false);
+    }
+
+    public QueryRewriteContext(
+        final XContentParserConfiguration parserConfiguration,
+        final Client client,
+        final LongSupplier nowInMillis,
+        final ResolvedIndices resolvedIndices,
+        final PointInTimeBuilder pit,
+        final QueryRewriteInterceptor queryRewriteInterceptor,
+        final boolean isExplain
     ) {
         this(
             parserConfiguration,
@@ -156,7 +172,8 @@ public QueryRewriteContext(
             null,
             resolvedIndices,
             pit,
-            queryRewriteInterceptor
+            queryRewriteInterceptor,
+            isExplain
         );
     }
 
@@ -262,6 +279,10 @@ public void setMapUnmappedFieldAsString(boolean mapUnmappedFieldAsString) {
         this.mapUnmappedFieldAsString = mapUnmappedFieldAsString;
     }
 
+    public boolean isExplain() {
+        return this.isExplain;
+    }
+
     public NamedWriteableRegistry getWriteableRegistry() {
         return writeableRegistry;
     }
diff --git a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java
@@ -272,7 +272,8 @@ private SearchExecutionContext(
             scriptService,
             null,
             null,
-            null
+            null,
+            false
         );
         this.shardId = shardId;
         this.shardRequestIndex = shardRequestIndex;
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java
@@ -1745,7 +1745,19 @@ public AliasFilter buildAliasFilter(ClusterState state, String index, Set<Resolv
      * Returns a new {@link QueryRewriteContext} with the given {@code now} provider
      */
     public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, ResolvedIndices resolvedIndices, PointInTimeBuilder pit) {
-        return new QueryRewriteContext(parserConfig, client, nowInMillis, resolvedIndices, pit, queryRewriteInterceptor);
+        return getRewriteContext(nowInMillis, resolvedIndices, pit, false);
+    }
+
+    /**
+     * Returns a new {@link QueryRewriteContext} with the given {@code now} provider
+     */
+    public QueryRewriteContext getRewriteContext(
+        LongSupplier nowInMillis,
+        ResolvedIndices resolvedIndices,
+        PointInTimeBuilder pit,
+        final boolean isExplain
+    ) {
+        return new QueryRewriteContext(parserConfig, client, nowInMillis, resolvedIndices, pit, queryRewriteInterceptor, isExplain);
     }
 
     public DataRewriteContext getDataRewriteContext(LongSupplier nowInMillis) {
diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java
@@ -1950,7 +1950,19 @@ private void rewriteAndFetchShardRequest(IndexShard shard, ShardSearchRequest re
      * Returns a new {@link QueryRewriteContext} with the given {@code now} provider
      */
     public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, ResolvedIndices resolvedIndices, PointInTimeBuilder pit) {
-        return indicesService.getRewriteContext(nowInMillis, resolvedIndices, pit);
+        return getRewriteContext(nowInMillis, resolvedIndices, pit, false);
+    }
+
+    /**
+     * Returns a new {@link QueryRewriteContext} with the given {@code now} provider
+     */
+    public QueryRewriteContext getRewriteContext(
+        LongSupplier nowInMillis,
+        ResolvedIndices resolvedIndices,
+        PointInTimeBuilder pit,
+        final boolean isExplain
+    ) {
+        return indicesService.getRewriteContext(nowInMillis, resolvedIndices, pit, isExplain);
     }
 
     public CoordinatorRewriteContextProvider getCoordinatorRewriteContextProvider(LongSupplier nowInMillis) {
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java
@@ -78,7 +78,7 @@ public T addChild(RetrieverBuilder retrieverBuilder) {
     /**
      * Combines the provided {@code rankResults} to return the final top documents.
      */
-    protected abstract RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults);
+    protected abstract RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, boolean explain);
 
     @Override
     public final boolean isCompound() {
@@ -181,7 +181,7 @@ public void onResponse(MultiSearchResponse items) {
                         failures.forEach(ex::addSuppressed);
                         listener.onFailure(ex);
                     } else {
-                        results.set(combineInnerRetrieverResults(topDocs));
+                        results.set(combineInnerRetrieverResults(topDocs, ctx.isExplain()));
                         listener.onResponse(null);
                     }
                 }
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java
@@ -148,7 +148,7 @@ protected RescorerRetrieverBuilder clone(List<RetrieverSource> newChildRetriever
     }
 
     @Override
-    protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults) {
+    protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, boolean explain) {
         assert rankResults.size() == 1;
         ScoreDoc[] scoreDocs = rankResults.get(0);
         RankDoc[] rankDocs = new RankDoc[scoreDocs.length];
diff --git a/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java b/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java
@@ -131,6 +131,7 @@
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.hasSize;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -1744,7 +1745,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
             NodeClient client = new NodeClient(settings, threadPool);
 
             SearchService searchService = mock(SearchService.class);
-            when(searchService.getRewriteContext(any(), any(), any())).thenReturn(
+            when(searchService.getRewriteContext(any(), any(), any(), anyBoolean())).thenReturn(
                 new QueryRewriteContext(null, null, null, null, null, null)
             );
             ClusterService clusterService = new ClusterService(
diff --git a/server/src/test/java/org/elasticsearch/index/query/QueryRewriteContextTests.java b/server/src/test/java/org/elasticsearch/index/query/QueryRewriteContextTests.java
@@ -53,7 +53,8 @@ public void testGetTierPreference() {
                 null,
                 null,
                 null,
-                null
+                null,
+                false
             );
 
             assertThat(context.getTierPreference(), is("data_cold"));
@@ -81,7 +82,8 @@ public void testGetTierPreference() {
                 null,
                 null,
                 null,
-                null
+                null,
+                false
             );
 
             assertThat(context.getTierPreference(), is(nullValue()));
diff --git a/test/framework/src/main/java/org/elasticsearch/search/retriever/TestCompoundRetrieverBuilder.java b/test/framework/src/main/java/org/elasticsearch/search/retriever/TestCompoundRetrieverBuilder.java
@@ -38,7 +38,7 @@ protected TestCompoundRetrieverBuilder clone(List<RetrieverSource> newChildRetri
     }
 
     @Override
-    protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults) {
+    protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, boolean explain) {
         return new RankDoc[0];
     }
 
diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java
@@ -636,7 +636,8 @@ QueryRewriteContext createQueryRewriteContext() {
                 scriptService,
                 createMockResolvedIndices(),
                 null,
-                createMockQueryRewriteInterceptor()
+                createMockQueryRewriteInterceptor(),
+                false
             );
         }
 
diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java
@@ -169,13 +169,17 @@ protected QueryRuleRetrieverBuilder clone(List<RetrieverSource> newChildRetrieve
     }
 
     @Override
-    protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults) {
+    protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, boolean explain) {
         assert rankResults.size() == 1;
         ScoreDoc[] scoreDocs = rankResults.get(0);
         RankDoc[] rankDocs = new RuleQueryRankDoc[scoreDocs.length];
         for (int i = 0; i < scoreDocs.length; i++) {
             ScoreDoc scoreDoc = scoreDocs[i];
-            rankDocs[i] = new RuleQueryRankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex, rulesetIds, matchCriteria);
+            if (explain) {
+                rankDocs[i] = new RuleQueryRankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex, rulesetIds, matchCriteria);
+            } else {
+                rankDocs[i] = new RuleQueryRankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex);
+            }
             rankDocs[i].rank = i + 1;
         }
         return rankDocs;
diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/RuleQueryRankDoc.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/RuleQueryRankDoc.java
@@ -16,6 +16,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -27,6 +28,10 @@ public class RuleQueryRankDoc extends RankDoc {
     public final List<String> rulesetIds;
     public final Map<String, Object> matchCriteria;
 
+    public RuleQueryRankDoc(int doc, float score, int shardIndex) {
+        this(doc, score, shardIndex, null, null);
+    }
+
     public RuleQueryRankDoc(int doc, float score, int shardIndex, List<String> rulesetIds, Map<String, Object> matchCriteria) {
         super(doc, score, shardIndex);
         this.rulesetIds = rulesetIds;
@@ -35,13 +40,20 @@ public RuleQueryRankDoc(int doc, float score, int shardIndex, List<String> rules
 
     public RuleQueryRankDoc(StreamInput in) throws IOException {
         super(in);
-        rulesetIds = in.readStringCollectionAsImmutableList();
-        matchCriteria = in.readGenericMap();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.RANK_DOC_OPTIONAL_METADATA_FOR_EXPLAIN)) {
+            List<String> inRulesetIds = in.readOptionalStringCollectionAsList();
+            this.rulesetIds = inRulesetIds == null ? null : Collections.unmodifiableList(inRulesetIds);
+            boolean matchCriteriaExists = in.readBoolean();
+            this.matchCriteria = matchCriteriaExists ? in.readGenericMap() : null;
+        } else {
+            rulesetIds = in.readStringCollectionAsImmutableList();
+            matchCriteria = in.readGenericMap();
+        }
     }
 
     @Override
     public Explanation explain(Explanation[] sources, String[] queryNames) {
-
+        assert rulesetIds != null && matchCriteria != null;
         return Explanation.match(
             score,
             "query rules evaluated rules from rulesets " + rulesetIds + " and match criteria " + matchCriteria,
@@ -51,8 +63,16 @@ public Explanation explain(Explanation[] sources, String[] queryNames) {
 
     @Override
     public void doWriteTo(StreamOutput out) throws IOException {
-        out.writeStringCollection(rulesetIds);
-        out.writeGenericMap(matchCriteria);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.RANK_DOC_OPTIONAL_METADATA_FOR_EXPLAIN)) {
+            out.writeOptionalStringCollection(rulesetIds);
+            out.writeBoolean(matchCriteria != null);
+            if (matchCriteria != null) {
+                out.writeGenericMap(matchCriteria);
+            }
+        } else {
+            out.writeStringCollection(rulesetIds == null ? Collections.emptyList() : rulesetIds);
+            out.writeGenericMap(matchCriteria == null ? Collections.emptyMap() : matchCriteria);
+        }
     }
 
     @Override
@@ -89,10 +109,14 @@ public String getWriteableName() {
 
     @Override
     protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.array("rulesetIds", rulesetIds.toArray());
-        builder.startObject("matchCriteria");
-        builder.mapContents(matchCriteria);
-        builder.endObject();
+        if (rulesetIds != null) {
+            builder.array("rulesetIds", rulesetIds.toArray());
+        }
+        if (matchCriteria != null) {
+            builder.startObject("matchCriteria");
+            builder.mapContents(matchCriteria);
+            builder.endObject();
+        }
     }
 
     @Override
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankDoc.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankDoc.java
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java
diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFQueryPhaseRankCoordinatorContext.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFQueryPhaseRankCoordinatorContext.java
diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFQueryPhaseRankShardContext.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFQueryPhaseRankShardContext.java
diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankDoc.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankDoc.java
diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java

Original file line number	Diff line number	Diff line change
`@@ -807,7 +807,8 @@ public QueryRewriteContext newQueryRewriteContext(`
`807`	`807`	`scriptService,`
`808`	`808`	`null,`
`809`	`809`	`null,`
`810`		`- null`
	`810`	`+ null,`
	`811`	`+ false`
`811`	`812`	`);`
`812`	`813`	`}`
`813`	`814`
Original file line number	Diff line number	Diff line change
`@@ -148,7 +148,7 @@ protected RescorerRetrieverBuilder clone(List<RetrieverSource> newChildRetriever`
`148`	`148`	`}`
`149`	`149`
`150`	`150`	`@Override`
`151`		`- protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults) {`
	`151`	`+ protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, boolean explain) {`
`152`	`152`	`assert rankResults.size() == 1;`
`153`	`153`	`ScoreDoc[] scoreDocs = rankResults.get(0);`
`154`	`154`	`RankDoc[] rankDocs = new RankDoc[scoreDocs.length];`