SOLR-17319 : Combined Query Feature for Multi Query Execution

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

@@ -0,0 +1,62 @@
+/*
+ * 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);
+  }
+
+  /**
+   * Propagates all the properties from parent ResponseBuilder to the all the children which are
+   * being set later after the CombinedQueryComponent is prepared.
+   */
+  public final void propagate() {
+    responseBuilders.forEach(
+        thisRb -> {
+          thisRb.setNeedDocSet(isNeedDocSet());
+          thisRb.setNeedDocList(isNeedDocList());
+          thisRb.doFacets = doFacets;
+          thisRb.doHighlights = doHighlights;
+          thisRb.doExpand = doExpand;
+          thisRb.doTerms = doTerms;
+          thisRb.doStats = doStats;
+          thisRb.setDistribStatsDisabled(isDistribStatsDisabled());
+        });
+  }
+}

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

@@ -0,0 +1,89 @@
+/*
+ * 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.lang.invoke.MethodHandles;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.solr.common.params.CombinerParams;
+import org.apache.solr.request.SolrQueryRequest;
+import org.apache.solr.response.SolrQueryResponse;
+import org.apache.solr.search.facet.FacetModule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The CombinedQuerySearchHandler class extends the SearchHandler and provides custom behavior for
+ * handling combined queries. It overrides methods to create a response builder based on the {@link
+ * CombinerParams#COMBINER} parameter and to define the default components included in the search
+ * configuration.
+ */
+public class CombinedQuerySearchHandler extends SearchHandler {
+
+  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
+  /**
+   * Overrides the default response builder creation method. This method checks if the {@link
+   * CombinerParams#COMBINER} parameter is set to true in the request. If it is, it returns an
+   * instance of {@link CombinedQueryResponseBuilder}, otherwise, it returns an instance of {@link
+   * ResponseBuilder}.
+   *
+   * @param req the SolrQueryRequest object
+   * @param rsp the SolrQueryResponse object
+   * @param components the list of SearchComponent objects
+   * @return the appropriate ResponseBuilder instance based on the CombinerParams.COMBINER parameter
+   */
+  @Override
+  protected ResponseBuilder newResponseBuilder(
+      SolrQueryRequest req, SolrQueryResponse rsp, List<SearchComponent> components) {
+    if (req.getParams().getBool(CombinerParams.COMBINER, false)) {
+      return new CombinedQueryResponseBuilder(req, rsp, components);
+    }
+    return super.newResponseBuilder(req, rsp, components);
+  }
+
+  /**
+   * Overrides the default components and returns a list of component names that are included in the
+   * default configuration.
+   *
+   * @return a list of component names
+   */
+  @Override
+  @SuppressWarnings("unchecked")
+  protected List<String> getDefaultComponents() {
+    List<String> names = new ArrayList<>(9);
+    names.add(CombinedQueryComponent.COMPONENT_NAME);
+    names.add(FacetComponent.COMPONENT_NAME);
+    names.add(FacetModule.COMPONENT_NAME);
+    names.add(MoreLikeThisComponent.COMPONENT_NAME);
+    names.add(HighlightComponent.COMPONENT_NAME);
+    names.add(StatsComponent.COMPONENT_NAME);
+    names.add(DebugComponent.COMPONENT_NAME);
+    names.add(ExpandComponent.COMPONENT_NAME);
+    names.add(TermsComponent.COMPONENT_NAME);
+    return names;
+  }
+
+  @Override
+  protected void postPrepareComponents(ResponseBuilder rb) {
+    super.postPrepareComponents(rb);
+    // propagate the CombinedQueryResponseBuilder's state to all subBuilders after prepare
+    if (rb instanceof CombinedQueryResponseBuilder crb) {
+      crb.propagate();
+    }
+  }
+}

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

@@ -22,6 +22,7 @@
 import java.lang.invoke.MethodHandles;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
@@ -31,6 +32,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import net.jcip.annotations.NotThreadSafe;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.solr.client.solrj.SolrRequest;
 import org.apache.solr.client.solrj.SolrResponse;
 import org.apache.solr.client.solrj.impl.Http2SolrClient;
@@ -404,6 +406,9 @@ public void cancelAll() {
 
   @Override
   public void prepDistributed(ResponseBuilder rb) {
+    if (rb.isForcedDistrib()) {
+      forceDistributed(rb);
+    }
     final SolrQueryRequest req = rb.req;
     final SolrParams params = req.getParams();
     final String shards = params.get(ShardParams.SHARDS);
@@ -500,6 +505,28 @@ public void prepDistributed(ResponseBuilder rb) {
     }
   }
 
+  private static void forceDistributed(ResponseBuilder rb) {
+    SolrQueryRequest req = rb.req;
+    ModifiableSolrParams solrParams = new ModifiableSolrParams(req.getParams());
+    solrParams.set("shortCircuit", false);
+    req.setParams(solrParams);
+    if (req.getHttpSolrCall() != null
+        && StringUtils.isEmpty(req.getParams().get(ShardParams.SHARDS))) {
+      String scheme = req.getHttpSolrCall().getReq().getScheme();
+      String host = req.getHttpSolrCall().getReq().getServerName();
+      int port = req.getHttpSolrCall().getReq().getServerPort();
+      String context = req.getHttpSolrCall().getReq().getContextPath();
+      String core = req.getCore().getName();
+      String localShardUrl =
+          String.format(Locale.ROOT, "%s://%s:%d%s/%s", scheme, host, port, context, core);
+      solrParams.set(ShardParams.SHARDS, localShardUrl);
+      req.setParams(solrParams);
+      return;
+    }
+    throw new SolrException(
+        SolrException.ErrorCode.BAD_REQUEST, "Force Distributed cannot be enabled");
+  }
+
   private static String createSliceShardsStr(final List<String> shardUrls) {
     return String.join("|", shardUrls);
   }

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

@@ -905,6 +905,67 @@ protected boolean addFL(StringBuilder fl, String field, boolean additionalAdded)
     return true;
   }
 
+  protected abstract static class ShardDocQueue {
+    public abstract boolean push(ShardDoc shardDoc);
+
+    public abstract Map<Object, ShardDoc> resultIds(int offset);
+  }
+  ;
+
+  protected ShardDocQueue newShardDocQueue(
+      SolrIndexSearcher searcher, SortField[] sortFields, Integer size) {
+    return new ShardDocQueue() {
+
+      // id to shard mapping, to eliminate any accidental dups
+      private final HashMap<Object, String> uniqueDoc = new HashMap<>();
+
+      private final ShardFieldSortedHitQueue queue =
+          new ShardFieldSortedHitQueue(sortFields, size, searcher);
+
+      @Override
+      public boolean push(ShardDoc shardDoc) {
+        final String prevShard = uniqueDoc.put(shardDoc.id, shardDoc.shard);
+        if (prevShard != null) {
+          // duplicate detected
+
+          // For now, just always use the first encountered since we can't currently
+          // remove the previous one added to the priority queue.  If we switched
+          // to the Java5 PriorityQueue, this would be easier.
+          return false;
+          // make which duplicate is used deterministic based on shard
+          // if (prevShard.compareTo(shardDoc.shard) >= 0) {
+          //  TODO: remove previous from priority queue
+          //  return false;
+          // }
+        }
+
+        queue.insertWithOverflow(shardDoc);
+        return true;
+      }
+
+      @Override
+      public Map<Object, ShardDoc> resultIds(int offset) {
+        final Map<Object, ShardDoc> resultIds = new HashMap<>();
+
+        // The queue now has 0 -> queuesize docs, where queuesize <= start + rows
+        // So we want to pop the last documents off the queue to get
+        // the docs offset -> queuesize
+        int resultSize = queue.size() - offset;
+        resultSize = Math.max(0, resultSize); // there may not be any docs in range
+
+        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;
+      }
+    };
+  }
+
   protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     List<MergeStrategy> mergeStrategies = rb.getMergeStrategies();
     if (mergeStrategies != null) {
@@ -947,14 +1008,10 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     IndexSchema schema = rb.req.getSchema();
     SchemaField uniqueKeyField = schema.getUniqueKeyField();
 
-    // id to shard mapping, to eliminate any accidental dups
-    HashMap<Object, String> uniqueDoc = new HashMap<>();
-
     // Merge the docs via a priority queue so we don't have to sort *all* of the
     // documents... we only need to order the top (rows+start)
-    final ShardFieldSortedHitQueue queue =
-        new ShardFieldSortedHitQueue(
-            sortFields, ss.getOffset() + ss.getCount(), rb.req.getSearcher());
+    final ShardDocQueue shardDocQueue =
+        newShardDocQueue(rb.req.getSearcher(), sortFields, ss.getOffset() + ss.getCount());
 
     NamedList<Object> shardInfo = null;
     if (rb.req.getParams().getBool(ShardParams.SHARDS_INFO, false)) {
@@ -1125,23 +1182,6 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
       for (int i = 0; i < docs.size(); i++) {
         SolrDocument doc = docs.get(i);
         Object id = doc.getFieldValue(uniqueKeyField.getName());
-
-        String prevShard = uniqueDoc.put(id, srsp.getShard());
-        if (prevShard != null) {
-          // duplicate detected
-          numFound--;
-
-          // For now, just always use the first encountered since we can't currently
-          // remove the previous one added to the priority queue.  If we switched
-          // to the Java5 PriorityQueue, this would be easier.
-          continue;
-          // make which duplicate is used deterministic based on shard
-          // if (prevShard.compareTo(srsp.shard) >= 0) {
-          //  TODO: remove previous from priority queue
-          //  continue;
-          // }
-        }
-
         ShardDoc shardDoc = new ShardDoc();
         shardDoc.id = id;
         shardDoc.shard = srsp.getShard();
@@ -1160,42 +1200,18 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
 
         shardDoc.sortFieldValues = unmarshalledSortFieldValues;
 
-        queue.insertWithOverflow(shardDoc);
+        if (!shardDocQueue.push(shardDoc)) {
+          numFound--;
+        }
       } // end for-each-doc-in-response
     } // end for-each-response
 
-    // The queue now has 0 -> queuesize docs, where queuesize <= start + rows
-    // So we want to pop the last documents off the queue to get
-    // the docs offset -> queuesize
-    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);
-    }
-
     // 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());
-    // size appropriately
-    for (int i = 0; i < resultSize; i++) responseDocs.add(null);
-
-    // save these results in a private area so we can access them
-    // again when retrieving stored fields.
-    // TODO: use ResponseBuilder (w/ comments) or the request context?
-    rb.resultIds = resultIds;
-    rb.setResponseDocs(responseDocs);
+    setResultIdsAndResponseDocs(
+        rb, shardDocQueue, maxScore, numFound, hitCountIsExact, ss.getOffset());
 
     populateNextCursorMarkFromMergedShards(rb);
 
@@ -1241,6 +1257,30 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     }
   }
 
+  protected void setResultIdsAndResponseDocs(
+      ResponseBuilder rb,
+      ShardDocQueue shardDocQueue,
+      Float maxScore,
+      long numFound,
+      boolean hitCountIsExact,
+      int offset) {
+    final Map<Object, ShardDoc> resultIds = shardDocQueue.resultIds(offset);
+
+    final SolrDocumentList responseDocs = new SolrDocumentList();
+    if (maxScore != null) responseDocs.setMaxScore(maxScore);
+    responseDocs.setNumFound(numFound);
+    responseDocs.setNumFoundExact(hitCountIsExact);
+    responseDocs.setStart(offset);
+    // size appropriately
+    for (int i = 0; i < resultIds.size(); i++) responseDocs.add(null);
+
+    // save these results in a private area so we can access them
+    // again when retrieving stored fields.
+    // TODO: use ResponseBuilder (w/ comments) or the request context?
+    rb.resultIds = resultIds;
+    rb.setResponseDocs(responseDocs);
+  }
+
   /**
    * Inspects the state of the {@link ResponseBuilder} and populates the next {@link
    * ResponseBuilder#setNextCursorMark} as appropriate based on the merged sort values from