Add chunk_rescorer usage to output of explain and profile for text_similarity_rank_retriever (elastic#137249)

kderusso · elasticsearchmachine · chrisparrinello · commit b61bffb9dab1 · 2025-11-03T14:06:17.000-06:00
* Add chunk_rescorer usage to output of explain for text_similarity_rank_retriever

* Update docs/changelog/137249.yaml

* Update toString

* Add support for profile

* Update 137249.yaml

* That's what you get for resolving conflicts in the UI, fixed compile failure

* [CI] Auto commit changes from spotless

* Fix test compilation

---------

Co-authored-by: elasticsearchmachine &lt;infra-root+elasticsearchmachine@elastic.co&gt;
diff --git a/docs/changelog/137249.yaml b/docs/changelog/137249.yaml
@@ -0,0 +1,5 @@
+pr: 137249
+summary: Add `chunk_rescorer` usage to output of explain and profile for `text_similarity_rank_retriever`
+area: Relevance
+type: enhancement
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java
@@ -598,6 +598,7 @@ public void onFailure(Exception e) {
         });
 
         final boolean isExplain = source != null && source.explain() != null && source.explain();
+        final boolean isProfile = source != null && source.profile();
         Rewriteable.rewriteAndFetch(
             original,
             searchService.getRewriteContext(
@@ -607,7 +608,8 @@ public void onFailure(Exception e) {
                 resolvedIndices,
                 original.pointInTimeBuilder(),
                 shouldMinimizeRoundtrips(original),
-                isExplain
+                isExplain,
+                isProfile
             ),
             rewriteListener
         );
diff --git a/server/src/main/java/org/elasticsearch/index/IndexService.java b/server/src/main/java/org/elasticsearch/index/IndexService.java
@@ -835,6 +835,7 @@ public QueryRewriteContext newQueryRewriteContext(
             null,
             null,
             null,
+            false,
             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
@@ -123,6 +123,7 @@ public CoordinatorRewriteContext(
             null,
             null,
             null,
+            false,
             false
         );
         this.dateFieldRangeInfo = dateFieldRangeInfo;
diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
@@ -80,6 +80,7 @@ public class QueryRewriteContext {
     private QueryRewriteInterceptor queryRewriteInterceptor;
     private final Boolean ccsMinimizeRoundTrips;
     private final boolean isExplain;
+    private final boolean isProfile;
     private Long timeRangeFilterFromMillis;
     private boolean trackTimeRangeFilterFrom = true;
 
@@ -103,7 +104,8 @@ public QueryRewriteContext(
         final PointInTimeBuilder pit,
         final QueryRewriteInterceptor queryRewriteInterceptor,
         final Boolean ccsMinimizeRoundTrips,
-        final boolean isExplain
+        final boolean isExplain,
+        final boolean isProfile
     ) {
 
         this.parserConfiguration = parserConfiguration;
@@ -127,6 +129,7 @@ public QueryRewriteContext(
         this.queryRewriteInterceptor = queryRewriteInterceptor;
         this.ccsMinimizeRoundTrips = ccsMinimizeRoundTrips;
         this.isExplain = isExplain;
+        this.isProfile = isProfile;
     }
 
     public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
@@ -150,6 +153,7 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration
             null,
             null,
             null,
+            false,
             false
         );
     }
@@ -175,6 +179,7 @@ public QueryRewriteContext(
             pit,
             queryRewriteInterceptor,
             ccsMinimizeRoundTrips,
+            false,
             false
         );
     }
@@ -189,7 +194,8 @@ public QueryRewriteContext(
         final PointInTimeBuilder pit,
         final QueryRewriteInterceptor queryRewriteInterceptor,
         final Boolean ccsMinimizeRoundTrips,
-        final boolean isExplain
+        final boolean isExplain,
+        final boolean isProfile
     ) {
         this(
             parserConfiguration,
@@ -211,7 +217,8 @@ public QueryRewriteContext(
             pit,
             queryRewriteInterceptor,
             ccsMinimizeRoundTrips,
-            isExplain
+            isExplain,
+            isProfile
         );
     }
 
@@ -324,6 +331,10 @@ public boolean isExplain() {
         return this.isExplain;
     }
 
+    public boolean isProfile() {
+        return this.isProfile;
+    }
+
     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
@@ -282,6 +282,7 @@ private SearchExecutionContext(
             null,
             null,
             null,
+            false,
             false
         );
         this.shardId = shardId;
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java
@@ -1852,7 +1852,8 @@ public QueryRewriteContext getRewriteContext(
         ResolvedIndices resolvedIndices,
         PointInTimeBuilder pit,
         final Boolean ccsMinimizeRoundTrips,
-        final boolean isExplain
+        final boolean isExplain,
+        final boolean isProfile
     ) {
         return new QueryRewriteContext(
             parserConfig,
@@ -1864,7 +1865,8 @@ public QueryRewriteContext getRewriteContext(
             pit,
             queryRewriteInterceptor,
             ccsMinimizeRoundTrips,
-            isExplain
+            isExplain,
+            isProfile
         );
     }
 
diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java
@@ -2159,7 +2159,7 @@ public QueryRewriteContext getRewriteContext(
         PointInTimeBuilder pit,
         final Boolean ccsMinimizeRoundTrips
     ) {
-        return getRewriteContext(nowInMillis, minTransportVersion, clusterAlias, resolvedIndices, pit, ccsMinimizeRoundTrips, false);
+        return getRewriteContext(nowInMillis, minTransportVersion, clusterAlias, resolvedIndices, pit, ccsMinimizeRoundTrips, false, false);
     }
 
     /**
@@ -2172,7 +2172,8 @@ public QueryRewriteContext getRewriteContext(
         ResolvedIndices resolvedIndices,
         PointInTimeBuilder pit,
         final Boolean ccsMinimizeRoundTrips,
-        final boolean isExplain
+        final boolean isExplain,
+        final boolean isProfile
     ) {
         return indicesService.getRewriteContext(
             nowInMillis,
@@ -2181,7 +2182,8 @@ public QueryRewriteContext getRewriteContext(
             resolvedIndices,
             pit,
             ccsMinimizeRoundTrips,
-            isExplain
+            isExplain,
+            isProfile
         );
     }
 
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java
@@ -191,7 +191,8 @@ public void onResponse(MultiSearchResponse items) {
                         failures.forEach(ex::addSuppressed);
                         listener.onFailure(ex);
                     } else {
-                        results.set(combineInnerRetrieverResults(topDocs, ctx.isExplain()));
+                        boolean enrichResults = ctx.isExplain() || ctx.isProfile();
+                        results.set(combineInnerRetrieverResults(topDocs, enrichResults));
                         listener.onResponse(null);
                     }
                 }
diff --git a/server/src/main/resources/transport/definitions/referable/text_similarity_rank_docs_explain_chunks.csv b/server/src/main/resources/transport/definitions/referable/text_similarity_rank_docs_explain_chunks.csv
@@ -0,0 +1 @@
+9205000
diff --git a/server/src/main/resources/transport/upper_bounds/9.3.csv b/server/src/main/resources/transport/upper_bounds/9.3.csv
@@ -1 +1 @@
-esql_resolve_fields_response_used,9204000
+text_similarity_rank_docs_explain_chunks,9205000
diff --git a/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java b/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java
@@ -1784,7 +1784,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
             NodeClient client = new NodeClient(settings, threadPool, TestProjectResolvers.alwaysThrow());
 
             SearchService searchService = mock(SearchService.class);
-            when(searchService.getRewriteContext(any(), any(), any(), any(), any(), anyBoolean(), anyBoolean())).thenReturn(
+            when(searchService.getRewriteContext(any(), any(), any(), any(), any(), anyBoolean(), anyBoolean(), anyBoolean())).thenReturn(
                 new QueryRewriteContext(null, null, null, 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
@@ -57,6 +57,7 @@ public void testGetTierPreference() {
                 null,
                 null,
                 null,
+                false,
                 false
             );
 
@@ -89,6 +90,7 @@ public void testGetTierPreference() {
                 null,
                 null,
                 null,
+                false,
                 false
             );
 
diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java
@@ -641,6 +641,7 @@ QueryRewriteContext createQueryRewriteContext() {
                 null,
                 createMockQueryRewriteInterceptor(),
                 null,
+                false,
                 false
             );
         }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java
@@ -29,6 +29,7 @@
 import static org.elasticsearch.xpack.inference.queries.LegacySemanticKnnVectorQueryRewriteInterceptor.SEMANTIC_KNN_VECTOR_QUERY_REWRITE_INTERCEPTION_SUPPORTED;
 import static org.elasticsearch.xpack.inference.queries.LegacySemanticMatchQueryRewriteInterceptor.SEMANTIC_MATCH_QUERY_REWRITE_INTERCEPTION_SUPPORTED;
 import static org.elasticsearch.xpack.inference.queries.LegacySemanticSparseVectorQueryRewriteInterceptor.SEMANTIC_SPARSE_VECTOR_QUERY_REWRITE_INTERCEPTION_SUPPORTED;
+import static org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankDoc.TEXT_SIMILARITY_RANK_DOC_EXPLAIN_CHUNKS;
 import static org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder.TEXT_SIMILARITY_RERANKER_SNIPPETS;
 
 /**
@@ -100,7 +101,8 @@ public Set<NodeFeature> getTestFeatures() {
                 InterceptedInferenceQueryBuilder.NEW_SEMANTIC_QUERY_INTERCEPTORS,
                 TEXT_SIMILARITY_RERANKER_SNIPPETS,
                 ModelStats.SEMANTIC_TEXT_USAGE,
-                SEARCH_USAGE_EXTENDED_DATA
+                SEARCH_USAGE_EXTENDED_DATA,
+                TEXT_SIMILARITY_RANK_DOC_EXPLAIN_CHUNKS
             )
         );
         testFeatures.addAll(getFeatures());
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/ChunkScorerConfig.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/ChunkScorerConfig.java
@@ -77,6 +77,10 @@ public Integer size() {
         return size;
     }
 
+    public int sizeOrDefault() {
+        return size != null ? size : DEFAULT_SIZE;
+    }
+
     public String inferenceText() {
         return inferenceText;
     }
@@ -100,6 +104,19 @@ public int hashCode() {
         return Objects.hash(size, inferenceText, chunkingSettings);
     }
 
+    @Override
+    public String toString() {
+        return "ChunkScorerConfig{"
+            + "size="
+            + sizeOrDefault()
+            + ", inferenceText=["
+            + inferenceText
+            + ']'
+            + ", chunkingSettings="
+            + chunkingSettings
+            + "}";
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
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/inference/src/test/java/org/elasticsearch/xpack/inference/queries/AbstractInterceptedInferenceQueryBuilderTestCase.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/AbstractInterceptedInferenceQueryBuilderTestCase.java
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankDocTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankDocTests.java
diff --git a/x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/70_text_similarity_rank_retriever.yml b/x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/70_text_similarity_rank_retriever.yml

Original file line number	Diff line number	Diff line change
`@@ -835,6 +835,7 @@ public QueryRewriteContext newQueryRewriteContext(`
`835`	`835`	`null,`
`836`	`836`	`null,`
`837`	`837`	`null,`
	`838`	`+ false,`
`838`	`839`	`false`
`839`	`840`	`);`
`840`	`841`	`}`
-Original file line number
+Diff line change
             null,
             null,
             null,
 +            false,
             false
         );
         this.dateFieldRangeInfo = dateFieldRangeInfo;
Original file line number	Diff line number	Diff line change
`@@ -191,7 +191,8 @@ public void onResponse(MultiSearchResponse items) {`
`191`	`191`	`failures.forEach(ex::addSuppressed);`
`192`	`192`	`listener.onFailure(ex);`
`193`	`193`	`} else {`
`194`		`- results.set(combineInnerRetrieverResults(topDocs, ctx.isExplain()));`
	`194`	`+ boolean enrichResults = ctx.isExplain() \|\| ctx.isProfile();`
	`195`	`+ results.set(combineInnerRetrieverResults(topDocs, enrichResults));`
`195`	`196`	`listener.onResponse(null);`
`196`	`197`	`}`
`197`	`198`	`}`