Add sub_searches to the search endpoint (#96224)

This change adds a new top-level element to the search endpoint called sub_searches. This top-level 
element allows for a list of additional searches where each "sub search" will have a query executed 
separately as part of ranking and later combined into a final single set of documents based on the 
ranking algorithm.

Author: jdconrad

docs/changelog/96224.yaml

@@ -0,0 +1,14 @@
+pr: 96224
+summary: Add multiple queries for ranking to the search endpoint
+area: Ranking
+type: enhancement
+issues: []
+highlight:
+  title: Add multiple queries for ranking to the search endpoint
+  body: "The search endpoint adds a new top-level element called `sub_searches`. \
+  This top-level element is a list of searches used for ranking where each \
+  \"sub search\" is executed independently. The `sub_searches` element is \
+  used instead of `query` to allow a search using ranking to execute \
+  multiple queries. The `sub_searches` and `query` elements cannot be used \
+  together."
+  notable: true

modules/data-streams/src/test/java/org/elasticsearch/datastreams/action/GetDataStreamsTransportActionTests.java

@@ -199,6 +199,7 @@ public void testGetTimeSeriesDataStream() {
         );
     }
 
+    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/96672")
     public void testGetTimeSeriesMixedDataStream() {
         Instant now = Instant.now().truncatedTo(ChronoUnit.SECONDS);
         String dataStream1 = "ds-1";

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

@@ -136,9 +136,10 @@ private static TransportVersion registerTransportVersion(int id, String uniqueId
     public static final TransportVersion V_8_500_010 = registerTransportVersion(8_500_010, "9818C628-1EEC-439B-B943-468F61460675");
     public static final TransportVersion V_8_500_011 = registerTransportVersion(8_500_011, "2209F28D-B52E-4BC4-9889-E780F291C32E");
     public static final TransportVersion V_8_500_012 = registerTransportVersion(8_500_012, "BB6F4AF1-A860-4FD4-A138-8150FFBE0ABD");
+    public static final TransportVersion V_8_500_013 = registerTransportVersion(8_500_013, "f65b85ac-db5e-4558-a487-a1dde4f6a33a");
 
     private static class CurrentHolder {
-        private static final TransportVersion CURRENT = findCurrent(V_8_500_012);
+        private static final TransportVersion CURRENT = findCurrent(V_8_500_013);
 
         // finds the pluggable current version, or uses the given fallback
         private static TransportVersion findCurrent(TransportVersion fallback) {

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

@@ -8,11 +8,10 @@
 package org.elasticsearch.action.search;
 
 import org.apache.lucene.search.ScoreDoc;
-import org.elasticsearch.index.query.BoolQueryBuilder;
-import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.builder.SubSearchSourceBuilder;
 import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.dfs.DfsKnnResults;
 import org.elasticsearch.search.dfs.DfsSearchResult;
@@ -138,75 +137,23 @@ ShardSearchRequest rewriteShardSearchRequest(ShardSearchRequest request) {
             return request;
         }
 
-        if (source.rankBuilder() == null) {
-            // this path will use linear combination if there are
-            // multiple knn queries to combine all knn queries into
-            // a single query per shard
+        List<SubSearchSourceBuilder> subSearchSourceBuilders = new ArrayList<>(source.subSearches());
 
+        for (DfsKnnResults dfsKnnResults : knnResults) {
             List<ScoreDoc> scoreDocs = new ArrayList<>();
-            for (DfsKnnResults dfsKnnResults : knnResults) {
-                for (ScoreDoc scoreDoc : dfsKnnResults.scoreDocs()) {
-                    if (scoreDoc.shardIndex == request.shardRequestIndex()) {
-                        scoreDocs.add(scoreDoc);
-                    }
+            for (ScoreDoc scoreDoc : dfsKnnResults.scoreDocs()) {
+                if (scoreDoc.shardIndex == request.shardRequestIndex()) {
+                    scoreDocs.add(scoreDoc);
                 }
             }
             scoreDocs.sort(Comparator.comparingInt(scoreDoc -> scoreDoc.doc));
-            // It is possible that the different results refer to the same doc.
-            for (int i = 0; i < scoreDocs.size() - 1; i++) {
-                ScoreDoc scoreDoc = scoreDocs.get(i);
-                int j = i + 1;
-                for (; j < scoreDocs.size(); j++) {
-                    ScoreDoc otherScoreDoc = scoreDocs.get(j);
-                    if (otherScoreDoc.doc != scoreDoc.doc) {
-                        break;
-                    }
-                    scoreDoc.score += otherScoreDoc.score;
-                }
-                if (j > i + 1) {
-                    scoreDocs.subList(i + 1, j).clear();
-                }
-            }
-
             KnnScoreDocQueryBuilder knnQuery = new KnnScoreDocQueryBuilder(scoreDocs.toArray(new ScoreDoc[0]));
-            SearchSourceBuilder newSource = source.shallowCopy().knnSearch(List.of());
-            if (source.query() == null) {
-                newSource.query(knnQuery);
-            } else {
-                newSource.query(new BoolQueryBuilder().should(knnQuery).should(source.query()));
-            }
-            request.source(newSource);
-        } else {
-            // this path will keep knn queries separate for ranking per shard
-            // if there are multiple knn queries
-
-            List<QueryBuilder> rankQueryBuilders = new ArrayList<>();
-            if (source.query() != null) {
-                rankQueryBuilders.add(source.query());
-            }
-
-            for (DfsKnnResults dfsKnnResults : knnResults) {
-                List<ScoreDoc> scoreDocs = new ArrayList<>();
-                for (ScoreDoc scoreDoc : dfsKnnResults.scoreDocs()) {
-                    if (scoreDoc.shardIndex == request.shardRequestIndex()) {
-                        scoreDocs.add(scoreDoc);
-                    }
-                }
-                scoreDocs.sort(Comparator.comparingInt(scoreDoc -> scoreDoc.doc));
-                KnnScoreDocQueryBuilder knnQuery = new KnnScoreDocQueryBuilder(scoreDocs.toArray(new ScoreDoc[0]));
-                rankQueryBuilders.add(knnQuery);
-            }
-
-            BoolQueryBuilder searchQuery = new BoolQueryBuilder();
-            for (QueryBuilder queryBuilder : rankQueryBuilders) {
-                searchQuery.should(queryBuilder);
-            }
-
-            SearchSourceBuilder newSource = source.shallowCopy().query(searchQuery).knnSearch(List.of());
-            request.source(newSource);
-            request.rankQueryBuilders(rankQueryBuilders);
+            subSearchSourceBuilders.add(new SubSearchSourceBuilder(knnQuery));
         }
 
+        source = source.shallowCopy().subSearches(subSearchSourceBuilders).knnSearch(List.of());
+        request.source(source);
+
         return request;
     }
 }

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

@@ -104,7 +104,9 @@ private void innerRun() throws Exception {
         final SearchPhaseController.ReducedQueryPhase reducedQueryPhase = resultConsumer.reduce();
         // Usually when there is a single shard, we force the search type QUERY_THEN_FETCH. But when there's kNN, we might
         // still use DFS_QUERY_THEN_FETCH, which does not perform the "query and fetch" optimization during the query phase.
-        final boolean queryAndFetchOptimization = queryResults.length() == 1 && context.getRequest().hasKnnSearch() == false;
+        final boolean queryAndFetchOptimization = queryResults.length() == 1
+            && context.getRequest().hasKnnSearch() == false
+            && reducedQueryPhase.rankCoordinatorContext() == null;
         final Runnable finishPhase = () -> moveToNextPhase(
             queryResults,
             reducedQueryPhase,

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

@@ -359,6 +359,9 @@ public ActionRequestValidationException validate() {
             }
         }
         if (source != null) {
+            if (source.subSearches().size() >= 2 && source.rankBuilder() == null) {
+                validationException = addValidationError("[sub_searches] requires [rank]", validationException);
+            }
             if (source.aggregations() != null) {
                 validationException = source.aggregations().validate(validationException);
             }
@@ -378,10 +381,10 @@ public ActionRequestValidationException validate() {
                         validationException
                     );
                 }
-                if (source.knnSearch().isEmpty() || source.query() == null && source.knnSearch().size() < 2) {
+                int queryCount = source.subSearches().size() + source.knnSearch().size();
+                if (queryCount < 2) {
                     validationException = addValidationError(
-                        "[rank] requires a minimum of [2] result sets which"
-                            + " either needs at minimum [a query and a knn search] or [multiple knn searches]",
+                        "[rank] requires a minimum of [2] result sets using a combination of sub searches and/or knn searches",
                         validationException
                     );
                 }

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

@@ -20,6 +20,7 @@
 import org.elasticsearch.search.aggregations.PipelineAggregationBuilder;
 import org.elasticsearch.search.builder.PointInTimeBuilder;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.builder.SubSearchSourceBuilder;
 import org.elasticsearch.search.collapse.CollapseBuilder;
 import org.elasticsearch.search.fetch.subphase.FieldAndFormat;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
@@ -165,6 +166,14 @@ public SearchRequestBuilder setQuery(QueryBuilder queryBuilder) {
         return this;
     }
 
+    /**
+     * Constructs a new search source builder with a list of sub searches.
+     */
+    public SearchRequestBuilder setSubSearches(List<SubSearchSourceBuilder> subSearches) {
+        sourceBuilder().subSearches(subSearches);
+        return this;
+    }
+
     /**
      * Sets a filter that will be executed after the query has been executed and only has affect on the search hits
      * (not aggregations). This filter is always executed as last filtering mechanism.

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

@@ -236,7 +236,8 @@ public void sendExecuteQuery(
     ) {
         // we optimize this and expect a QueryFetchSearchResult if we only have a single shard in the search request
         // this used to be the QUERY_AND_FETCH which doesn't exist anymore.
-        final boolean fetchDocuments = request.numberOfShards() == 1;
+        final boolean fetchDocuments = request.numberOfShards() == 1
+            && (request.source() == null || request.source().rankBuilder() == null);
         Writeable.Reader<SearchPhaseResult> reader = fetchDocuments ? QueryFetchSearchResult::new : in -> new QuerySearchResult(in, true);
 
         final ActionListener<? super SearchPhaseResult> handler = responseWrapper.apply(connection, listener);

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

@@ -77,6 +77,7 @@
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.AggregationContext.ProductionAggregationContext;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.builder.SubSearchSourceBuilder;
 import org.elasticsearch.search.collapse.CollapseContext;
 import org.elasticsearch.search.dfs.DfsPhase;
 import org.elasticsearch.search.dfs.DfsSearchResult;
@@ -626,7 +627,7 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, SearchSh
             } finally {
                 tracer.stopTrace();
             }
-            if (request.numberOfShards() == 1) {
+            if (request.numberOfShards() == 1 && (request.source() == null || request.source().rankBuilder() == null)) {
                 // we already have query results, but we can run fetch at the same time
                 context.addFetchResult();
                 return executeFetchPhase(readerContext, context, afterQueryTime);
@@ -992,13 +993,6 @@ protected SearchContext createContext(
             if (context.size() == -1) {
                 context.size(DEFAULT_SIZE);
             }
-            if (request.rankQueryBuilders().isEmpty() == false) {
-                List<Query> rankQueries = new ArrayList<>();
-                for (QueryBuilder queryBuilder : request.rankQueryBuilders()) {
-                    rankQueries.add(queryBuilder.toQuery(context.getSearchExecutionContext()));
-                }
-                context.rankShardContext(request.source().rankBuilder().buildRankShardContext(rankQueries, context.from()));
-            }
             context.setTask(task);
 
             context.preProcess();
@@ -1166,7 +1160,7 @@ private void processFailure(ReaderContext context, Exception exc) {
         }
     }
 
-    private void parseSource(DefaultSearchContext context, SearchSourceBuilder source, boolean includeAggregations) {
+    private void parseSource(DefaultSearchContext context, SearchSourceBuilder source, boolean includeAggregations) throws IOException {
         // nothing to parse...
         if (source == null) {
             return;
@@ -1176,9 +1170,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         context.from(source.from());
         context.size(source.size());
         Map<String, InnerHitContextBuilder> innerHitBuilders = new HashMap<>();
-        if (source.query() != null) {
-            InnerHitContextBuilder.extractInnerHits(source.query(), innerHitBuilders);
-            context.parsedQuery(searchExecutionContext.toQuery(source.query()));
+        QueryBuilder query = source.query();
+        if (query != null) {
+            InnerHitContextBuilder.extractInnerHits(query, innerHitBuilders);
+            context.parsedQuery(searchExecutionContext.toQuery(query));
         }
         if (source.postFilter() != null) {
             InnerHitContextBuilder.extractInnerHits(source.postFilter(), innerHitBuilders);
@@ -1374,6 +1369,14 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             final CollapseContext collapseContext = source.collapse().build(searchExecutionContext);
             context.collapse(collapseContext);
         }
+
+        if (source.rankBuilder() != null) {
+            List<Query> queries = new ArrayList<>();
+            for (SubSearchSourceBuilder subSearchSourceBuilder : source.subSearches()) {
+                queries.add(subSearchSourceBuilder.toSearchQuery(context.getSearchExecutionContext()));
+            }
+            context.rankShardContext(source.rankBuilder().buildRankShardContext(queries, context.from()));
+        }
     }
 
     /**
@@ -1623,14 +1626,12 @@ private static boolean canMatchAfterRewrite(final ShardSearchRequest request, fi
     @SuppressWarnings("unchecked")
     public static boolean queryStillMatchesAfterRewrite(ShardSearchRequest request, QueryRewriteContext context) throws IOException {
         Rewriteable.rewrite(request.getRewriteable(), context, false);
-        final boolean aliasFilterCanMatch = request.getAliasFilter().getQueryBuilder() instanceof MatchNoneQueryBuilder == false;
-        final boolean canMatch;
+        boolean canMatch = request.getAliasFilter().getQueryBuilder() instanceof MatchNoneQueryBuilder == false;
         if (canRewriteToMatchNone(request.source())) {
-            QueryBuilder queryBuilder = request.source().query();
-            canMatch = aliasFilterCanMatch && queryBuilder instanceof MatchNoneQueryBuilder == false;
-        } else {
-            // null query means match_all
-            canMatch = aliasFilterCanMatch;
+            canMatch &= request.source()
+                .subSearches()
+                .stream()
+                .anyMatch(sqwb -> sqwb.getQueryBuilder() instanceof MatchNoneQueryBuilder == false);
         }
         return canMatch;
     }
@@ -1641,7 +1642,11 @@ public static boolean queryStillMatchesAfterRewrite(ShardSearchRequest request,
      * a global aggregation is part of this request or if there is a suggest builder present.
      */
     public static boolean canRewriteToMatchNone(SearchSourceBuilder source) {
-        if (source == null || source.query() == null || source.query() instanceof MatchAllQueryBuilder || source.suggest() != null) {
+        if (source == null || source.suggest() != null) {
+            return false;
+        }
+        if (source.subSearches().isEmpty()
+            || source.subSearches().stream().anyMatch(sqwb -> sqwb.getQueryBuilder() instanceof MatchAllQueryBuilder)) {
             return false;
         }
         AggregatorFactories.Builder aggregations = source.aggregations();