Add sub_searches to the search endpoint

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();