added originalQuery in ShardSearchRequest

Author: piergm

server/src/main/java/org/elasticsearch/TransportVersion.java

@@ -16,13 +16,7 @@
 
 import java.io.IOException;
 import java.lang.reflect.Field;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.NavigableMap;
-import java.util.Set;
-import java.util.TreeMap;
+import java.util.*;
 
 /**
  * Represents the version of the wire protocol used to communicate between a pair of ES nodes.
@@ -163,9 +157,10 @@ private static TransportVersion registerTransportVersion(int id, String uniqueId
     public static final TransportVersion V_8_500_028 = registerTransportVersion(8_500_028, "a6592d08-15cb-4e1a-b9b4-b2ba24058444");
     public static final TransportVersion V_8_500_029 = registerTransportVersion(8_500_029, "f3bd98af-6187-e161-e315-718a2fecc2db");
     public static final TransportVersion V_8_500_030 = registerTransportVersion(8_500_030, "b72d7f12-8ed3-4a5b-8e6a-4910ea10e0d7");
+    public static final TransportVersion V_8_500_031 = registerTransportVersion(8_500_031, "2a229066-1707-4faf-83f8-a678a3395b89");
 
     private static class CurrentHolder {
-        private static final TransportVersion CURRENT = findCurrent(V_8_500_030);
+        private static final TransportVersion CURRENT = findCurrent(V_8_500_031);
 
         // finds the pluggable current version, or uses the given fallback
         private static TransportVersion findCurrent(TransportVersion fallback) {

server/src/main/java/org/elasticsearch/action/search/CanMatchNodeRequest.java

@@ -207,6 +207,7 @@ public ShardSearchRequest createShardSearchRequest(Shard r) {
             numberOfShards,
             searchType,
             source,
+            source.query(),
             requestCache,
             r.aliasFilter,
             r.indexBoost,

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

@@ -146,7 +146,7 @@ final class DefaultSearchContext extends SearchContext {
         this.shardTarget = shardTarget;
         this.indexService = readerContext.indexService();
         this.indexShard = readerContext.indexShard();
-        this.originalQueryBuilder = request.originalSource().query();
+        this.originalQueryBuilder = request.originalQuery();
 
         Engine.Searcher engineSearcher = readerContext.acquireSearcher("search");
         this.searcher = new ContextIndexSearcher(

server/src/main/java/org/elasticsearch/search/internal/ShardSearchRequest.java

@@ -88,7 +88,7 @@ public class ShardSearchRequest extends TransportRequest implements IndicesReque
     // these are the only mutable fields, as they are subject to rewriting
     private AliasFilter aliasFilter;
     private SearchSourceBuilder source;
-    private SearchSourceBuilder originalSource;
+    private QueryBuilder originalQuery;
     private final ShardSearchContextId readerId;
     private final TimeValue keepAlive;
 
@@ -148,6 +148,7 @@ public ShardSearchRequest(
             numberOfShards,
             searchRequest.searchType(),
             searchRequest.source(),
+            searchRequest.source().query(),
             searchRequest.requestCache(),
             aliasFilter,
             indexBoost,
@@ -190,6 +191,7 @@ public ShardSearchRequest(ShardId shardId, long nowInMillis, AliasFilter aliasFi
             SearchType.QUERY_THEN_FETCH,
             null,
             null,
+            null,
             aliasFilter,
             1.0f,
             true,
@@ -211,6 +213,7 @@ public ShardSearchRequest(
         int numberOfShards,
         SearchType searchType,
         SearchSourceBuilder source,
+        QueryBuilder originalQuery,
         Boolean requestCache,
         AliasFilter aliasFilter,
         float indexBoost,
@@ -229,7 +232,7 @@ public ShardSearchRequest(
         this.numberOfShards = numberOfShards;
         this.searchType = searchType;
         this.source(source);
-        this.originalSource(source);
+        this.originalQuery(originalQuery);
         this.requestCache = requestCache;
         this.aliasFilter = aliasFilter;
         this.indexBoost = indexBoost;
@@ -254,7 +257,7 @@ public ShardSearchRequest(ShardSearchRequest clone) {
         this.numberOfShards = clone.numberOfShards;
         this.scroll = clone.scroll;
         this.source(clone.source);
-        this.originalSource(clone.originalSource);
+        this.originalQuery(clone.originalQuery);
         this.aliasFilter = clone.aliasFilter;
         this.indexBoost = clone.indexBoost;
         this.nowInMillis = clone.nowInMillis;
@@ -280,7 +283,9 @@ public ShardSearchRequest(StreamInput in) throws IOException {
         numberOfShards = in.readVInt();
         scroll = in.readOptionalWriteable(Scroll::new);
         source = in.readOptionalWriteable(SearchSourceBuilder::new);
-        originalSource = source;
+        if (in.getTransportVersion().onOrAfter(TransportVersion.V_8_500_031)) {
+            originalQuery = in.readOptionalNamedWriteable(QueryBuilder.class);
+        }
         if (in.getTransportVersion().onOrAfter(TransportVersion.V_8_8_0) && in.getTransportVersion().before(TransportVersion.V_8_500_013)) {
             // to deserialize between the 8.8 and 8.500.013 version we need to translate
             // the rank queries into sub searches if we are ranking; if there are no rank queries
@@ -302,7 +307,6 @@ public ShardSearchRequest(StreamInput in) throws IOException {
                     subSearchSourceBuilders.add(new SubSearchSourceBuilder(queryBuilder));
                 }
                 source.subSearches(subSearchSourceBuilders);
-                originalSource.subSearches(subSearchSourceBuilders);
             }
         }
         if (in.getTransportVersion().before(TransportVersion.V_8_0_0)) {
@@ -375,6 +379,9 @@ protected final void innerWriteTo(StreamOutput out, boolean asKey) throws IOExce
         }
         out.writeOptionalWriteable(scroll);
         out.writeOptionalWriteable(source);
+        if (out.getTransportVersion().onOrAfter(TransportVersion.V_8_500_031)) {
+            out.writeOptionalNamedWriteable(originalQuery);
+        }
         if (out.getTransportVersion().onOrAfter(TransportVersion.V_8_8_0)
             && out.getTransportVersion().before(TransportVersion.V_8_500_013)) {
             // to serialize between the 8.8 and 8.500.013 version we need to translate
@@ -459,12 +466,12 @@ public SearchSourceBuilder source() {
         return source;
     }
 
-    public SearchSourceBuilder originalSource() {
-        return originalSource;
+    public QueryBuilder originalQuery() {
+        return originalQuery;
     }
 
-    public void originalSource(SearchSourceBuilder originalSource) {
-        this.originalSource = originalSource;
+    public void originalQuery(QueryBuilder originalQuery) {
+        this.originalQuery = originalQuery;
     }
 
     public AliasFilter getAliasFilter() {
@@ -641,7 +648,7 @@ public Rewriteable rewrite(QueryRewriteContext ctx) throws IOException {
                     newSource.size(0);
                 }
                 request.source(newSource);
-                request.originalSource(request.source());
+                request.originalQuery(request.source().query());
                 request.setBottomSortValues(null);
             }
 

test/framework/src/main/java/org/elasticsearch/test/TestSearchContext.java

@@ -88,7 +88,7 @@ public TestSearchContext(IndexService indexService) {
         this.indexShard = indexService.getShardOrNull(0);
         searchExecutionContext = indexService.newSearchExecutionContext(0, 0, null, () -> 0L, null, emptyMap());
         this.request = new ShardSearchRequest(indexShard.shardId(), 0L, AliasFilter.EMPTY);
-        this.originalQueryBuilder = request.originalSource().query();
+        this.originalQueryBuilder = request.originalQuery();
     }
 
     public TestSearchContext(SearchExecutionContext searchExecutionContext) {