SOLR-17319 : Combined Query Feature for Multi Query Execution

solr/core/src/java/org/apache/solr/handler/component/CombinedQueryComponent.java

@@ -0,0 +1,169 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.solr.handler.component;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.solr.common.SolrDocumentList;
+import org.apache.solr.common.params.CombinerParams;
+import org.apache.solr.common.params.CursorMarkParams;
+import org.apache.solr.common.params.ShardParams;
+import org.apache.solr.common.params.SolrParams;
+import org.apache.solr.response.BasicResultContext;
+import org.apache.solr.response.ResultContext;
+import org.apache.solr.response.SolrQueryResponse;
+import org.apache.solr.search.DocListAndSet;
+import org.apache.solr.search.QueryResult;
+import org.apache.solr.search.combine.QueryAndResponseCombiner;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The CombinedQueryComponent class extends QueryComponent and provides support for executing
+ * multiple queries and combining their results.
+ */
+public class CombinedQueryComponent extends QueryComponent {
+
+  public static final String COMPONENT_NAME = "combined_query";
+  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
+  /**
+   * Overrides the prepare method to handle combined queries.
+   *
+   * @param rb the ResponseBuilder to prepare
+   * @throws IOException if an I/O error occurs during preparation
+   */
+  @Override
+  public void prepare(ResponseBuilder rb) throws IOException {
+    if (rb instanceof CombinedQueryResponseBuilder crb) {
+      SolrParams params = crb.req.getParams();
+      String[] queriesToCombineKeys = params.getParams(CombinerParams.COMBINER_QUERY);
+      for (String queryKey : queriesToCombineKeys) {
+        final var unparsedQuery = params.get(queryKey);
+        ResponseBuilder rbNew = new ResponseBuilder(rb.req, new SolrQueryResponse(), rb.components);
+        rbNew.setQueryString(unparsedQuery);
+        super.prepare(rbNew);
+        crb.responseBuilders.add(rbNew);
+      }
+    }
+    super.prepare(rb);
+  }
+
+  /**
+   * Overrides the process method to handle CombinedQueryResponseBuilder instances. This method
+   * processes the responses from multiple shards, combines them using the specified
+   * QueryAndResponseCombiner strategy, and sets the appropriate results and metadata in the
+   * CombinedQueryResponseBuilder.
+   *
+   * @param rb the ResponseBuilder object to process
+   * @throws IOException if an I/O error occurs during processing
+   */
+  @Override
+  public void process(ResponseBuilder rb) throws IOException {
+    if (rb instanceof CombinedQueryResponseBuilder crb) {
+      boolean partialResults = false;
+      boolean segmentTerminatedEarly = false;
+      List<QueryResult> queryResults = new ArrayList<>();
+      for (ResponseBuilder rbNow : crb.responseBuilders) {
+        super.process(rbNow);
+        DocListAndSet docListAndSet = rbNow.getResults();
+        QueryResult queryResult = new QueryResult();
+        queryResult.setDocListAndSet(docListAndSet);
+        queryResults.add(queryResult);
+        partialResults |= SolrQueryResponse.isPartialResults(rbNow.rsp.getResponseHeader());
+        rbNow.setCursorMark(rbNow.getCursorMark());
+        if (rbNow.rsp.getResponseHeader() != null) {
+          segmentTerminatedEarly |=
+              (boolean)
+                  rbNow
+                      .rsp
+                      .getResponseHeader()
+                      .getOrDefault(
+                          SolrQueryResponse.RESPONSE_HEADER_SEGMENT_TERMINATED_EARLY_KEY, false);
+        }
+      }
+      QueryAndResponseCombiner combinerStrategy =
+          QueryAndResponseCombiner.getImplementation(rb.req.getParams());
+      QueryResult combinedQueryResult = combinerStrategy.combine(queryResults);
+      combinedQueryResult.setPartialResults(partialResults);
+      combinedQueryResult.setSegmentTerminatedEarly(segmentTerminatedEarly);
+      crb.setResult(combinedQueryResult);
+      ResultContext ctx = new BasicResultContext(crb);
+      crb.rsp.addResponse(ctx);
+      crb.rsp
+          .getToLog()
+          .add(
+              "hits",
+              crb.getResults() == null || crb.getResults().docList == null
+                  ? 0
+                  : crb.getResults().docList.matches());
+      if (!crb.req.getParams().getBool(ShardParams.IS_SHARD, false)) {
+        if (null != crb.getNextCursorMark()) {
+          crb.rsp.add(
+              CursorMarkParams.CURSOR_MARK_NEXT, crb.getNextCursorMark().getSerializedTotem());
+        }
+      }
+
+      if (crb.mergeFieldHandler != null) {
+        crb.mergeFieldHandler.handleMergeFields(crb, crb.req.getSearcher());
+      } else {
+        doFieldSortValues(rb, crb.req.getSearcher());
+      }
+      doPrefetch(crb);
+    } else {
+      super.process(rb);
+    }
+  }
+
+  @Override
+  protected Map<Object, ShardDoc> createShardResult(
+      ResponseBuilder rb,
+      int resultSize,
+      ShardFieldSortedHitQueue queue,
+      Map<String, List<ShardDoc>> shardDocMap,
+      SolrDocumentList responseDocs) {
+    QueryAndResponseCombiner combinerStrategy =
+        QueryAndResponseCombiner.getImplementation(rb.req.getParams());
+    List<ShardDoc> combinedShardDocs = combinerStrategy.combine(shardDocMap);
+    Map<String, ShardDoc> shardDocIdMap = new HashMap<>();
+    shardDocMap.forEach(
+        (shardKey, shardDocs) ->
+            shardDocs.forEach(shardDoc -> shardDocIdMap.put(shardDoc.id.toString(), shardDoc)));
+    Map<Object, ShardDoc> resultIds = new HashMap<>();
+    float maxScore = 0.0f;
+    for (int i = 0; i < resultSize; i++) {
+      ShardDoc shardDoc = combinedShardDocs.get(i);
+      shardDoc.positionInResponse = i;
+      maxScore = Math.max(maxScore, shardDoc.score);
+      if (Float.isNaN(shardDocIdMap.get(shardDoc.id.toString()).score)) {
+        shardDoc.score = Float.NaN;
+      }
+      resultIds.put(shardDoc.id.toString(), shardDoc);
+    }
+    responseDocs.setMaxScore(maxScore);
+    return resultIds;
+  }
+
+  @Override
+  public String getDescription() {
+    return "Combined Query Component to support multiple query execution";
+  }
+}

solr/core/src/java/org/apache/solr/handler/component/CombinedQueryResponseBuilder.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.solr.handler.component;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.solr.request.SolrQueryRequest;
+import org.apache.solr.response.SolrQueryResponse;
+
+/**
+ * The CombinedQueryResponseBuilder class extends the ResponseBuilder class and is responsible for
+ * building a combined response for multiple SearchComponent objects. It orchestrates the process of
+ * constructing the SolrQueryResponse by aggregating results from various components.
+ */
+public class CombinedQueryResponseBuilder extends ResponseBuilder {
+
+  public final List<ResponseBuilder> responseBuilders = new ArrayList<>();
+
+  /**
+   * Constructs a CombinedQueryResponseBuilder instance.
+   *
+   * @param req the SolrQueryRequest object containing the query parameters and context.
+   * @param rsp the SolrQueryResponse object to which the combined results will be added.
+   * @param components a list of SearchComponent objects that will be used to build the response.
+   */
+  public CombinedQueryResponseBuilder(
+      SolrQueryRequest req, SolrQueryResponse rsp, List<SearchComponent> components) {
+    super(req, rsp, components);
+  }
+}

solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java

@@ -963,6 +963,7 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     boolean maxHitsTerminatedEarly = false;
     long approximateTotalHits = 0;
     int failedShardCount = 0;
+    Map<String, List<ShardDoc>> shardDocMap = new HashMap<>();
     for (ShardResponse srsp : sreq.responses) {
       SolrDocumentList docs = null;
       NamedList<?> responseHeader = null;
@@ -1152,7 +1153,7 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
         }
 
         shardDoc.sortFieldValues = unmarshalledSortFieldValues;
-
+        shardDocMap.computeIfAbsent(srsp.getShard(), list -> new ArrayList<>()).add(shardDoc);
         queue.insertWithOverflow(shardDoc);
       } // end for-each-doc-in-response
     } // end for-each-response
@@ -1163,21 +1164,16 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     int resultSize = queue.size() - ss.getOffset();
     resultSize = Math.max(0, resultSize); // there may not be any docs in range
 
-    Map<Object, ShardDoc> resultIds = new HashMap<>();
-    for (int i = resultSize - 1; i >= 0; i--) {
-      ShardDoc shardDoc = queue.pop();
-      shardDoc.positionInResponse = i;
-      // Need the toString() for correlation with other lists that must
-      // be strings (like keys in highlighting, explain, etc)
-      resultIds.put(shardDoc.id.toString(), shardDoc);
-    }
+    SolrDocumentList responseDocs = new SolrDocumentList();
+    if (maxScore != null) responseDocs.setMaxScore(maxScore);
+
+    Map<Object, ShardDoc> resultIds =
+        createShardResult(rb, resultSize, queue, shardDocMap, responseDocs);
 
     // Add hits for distributed requests
     // https://issues.apache.org/jira/browse/SOLR-3518
     rb.rsp.addToLog("hits", numFound);
 
-    SolrDocumentList responseDocs = new SolrDocumentList();
-    if (maxScore != null) responseDocs.setMaxScore(maxScore);
     responseDocs.setNumFound(numFound);
     responseDocs.setNumFoundExact(hitCountIsExact);
     responseDocs.setStart(ss.getOffset());
@@ -1234,6 +1230,34 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     }
   }
 
+  /**
+   * Creates a map of shard results based on the provided parameters.
+   *
+   * @param rb ResponseBuilder to retrieve any request info
+   * @param resultSize the desired size of the result map
+   * @param queue the queue containing sorted hits from different shards
+   * @param shardDocMap a map of shard documents indexed by shard ID
+   * @param responseDocs the final SolrDocumentList to be returned or set in the responseBuilder
+   * @return a map of shard documents, where the keys are the shard IDs as strings, and the values
+   *     are the corresponding ShardDoc objects
+   */
+  protected Map<Object, ShardDoc> createShardResult(
+      ResponseBuilder rb,
+      int resultSize,
+      ShardFieldSortedHitQueue queue,
+      Map<String, List<ShardDoc>> shardDocMap,
+      SolrDocumentList responseDocs) {
+    Map<Object, ShardDoc> resultIds = new HashMap<>();
+    for (int i = resultSize - 1; i >= 0; i--) {
+      ShardDoc shardDoc = queue.pop();
+      shardDoc.positionInResponse = i;
+      // Need the toString() for correlation with other lists that must
+      // be strings (like keys in highlighting, explain, etc)
+      resultIds.put(shardDoc.id.toString(), shardDoc);
+    }
+    return resultIds;
+  }
+
   /**
    * Inspects the state of the {@link ResponseBuilder} and populates the next {@link
    * ResponseBuilder#setNextCursorMark} as appropriate based on the merged sort values from

solr/core/src/java/org/apache/solr/handler/component/SearchHandler.java

@@ -49,6 +49,7 @@
 import org.apache.solr.cloud.ZkController;
 import org.apache.solr.common.SolrDocumentList;
 import org.apache.solr.common.SolrException;
+import org.apache.solr.common.params.CombinerParams;
 import org.apache.solr.common.params.CommonParams;
 import org.apache.solr.common.params.CursorMarkParams;
 import org.apache.solr.common.params.ModifiableSolrParams;
@@ -240,7 +241,7 @@ public void changed(SolrPackageLoader.SolrPackage pkg, Ctx ctx) {
   }
 
   @SuppressWarnings({"unchecked"})
-  private void initComponents() {
+  private void initComponents(boolean isCombinedQuery) {
     Object declaredComponents = initArgs.get(INIT_COMPONENTS);
     List<String> first = (List<String>) initArgs.get(INIT_FIRST_COMPONENTS);
     List<String> last = (List<String>) initArgs.get(INIT_LAST_COMPONENTS);
@@ -251,6 +252,11 @@ private void initComponents() {
       // Use the default component list
       list = getDefaultComponents();
 
+      if (isCombinedQuery && list.getFirst().equals(QueryComponent.COMPONENT_NAME)) {
+        list.removeFirst();
+        list.addFirst(CombinedQueryComponent.COMPONENT_NAME);
+      }
+
       if (first != null) {
         List<String> clist = first;
         clist.addAll(list);
@@ -289,12 +295,12 @@ private void initComponents() {
     this.components = components;
   }
 
-  public List<SearchComponent> getComponents() {
+  public List<SearchComponent> getComponents(boolean isCombinedQuery) {
     List<SearchComponent> result = components; // volatile read
     if (result == null) {
       synchronized (this) {
         if (components == null) {
-          initComponents();
+          initComponents(isCombinedQuery);
         }
         result = components;
       }
@@ -354,6 +360,9 @@ public ShardHandler getAndPrepShardHandler(SolrQueryRequest req, ResponseBuilder
    */
   protected ResponseBuilder newResponseBuilder(
       SolrQueryRequest req, SolrQueryResponse rsp, List<SearchComponent> components) {
+    if (req.getParams().getBool(CombinerParams.COMBINER, false)) {
+      return new CombinedQueryResponseBuilder(req, rsp, components);
+    }
     return new ResponseBuilder(req, rsp, components);
   }
 
@@ -401,7 +410,8 @@ public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throw
           purpose, n -> shardPurposes.computeIfAbsent(n, name -> new Counter()).inc());
     }
 
-    List<SearchComponent> components = getComponents();
+    List<SearchComponent> components =
+        getComponents(req.getParams().getBool(CombinerParams.COMBINER, false));
     ResponseBuilder rb = newResponseBuilder(req, rsp, components);
     if (rb.requestInfo != null) {
       rb.requestInfo.setResponseBuilder(rb);