SOLR-17447 : Support for maxHitsPerShard.

Adds the capability to limit hits per shard. "maxHitsPerShard" request parameter controls how many hits the searcher should run over per shard. Once the searcher runs over the specfied number of documents, it will terminate the search with EarlyTerminatingCollectorException. This will be indicated by a new response header "terminatedEarly" also the "partialResults" will indicate that the results are partial. This parameter is supported in MT mode as well. Though there are other mechanisms to control runaway queries with CPU usage limits and time limits, this is simpler for certain use cases esp in case high recall queries and rerank use cases.
apache · Jan 10, 2025 · e48c636 · e48c636
1 parent 440e70d
commit e48c636
Show file tree

Hide file tree

Showing 11 changed files with 167 additions and 51 deletions.
diff --git a/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java b/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java
@@ -407,6 +407,17 @@ public void process(ResponseBuilder rb) throws IOException {
         params.getBool(
             CommonParams.SEGMENT_TERMINATE_EARLY, CommonParams.SEGMENT_TERMINATE_EARLY_DEFAULT));
 
+    // max hits per shard
+    final Integer maxHits = params.getInt(CommonParams.MAX_HITS_PER_SHARD);
+
+    if (maxHits != null) {
+      int hitsPerShard = Math.max(maxHits, cmd.getLen());
+      if (cmd.getMinExactCount() < Integer.MAX_VALUE) {
+        hitsPerShard = Math.max(cmd.getMinExactCount(), hitsPerShard);
+      }
+      cmd.setMaxHits(hitsPerShard);
+    }
+
     //
     // grouping / field collapsing
     //

diff --git a/solr/core/src/java/org/apache/solr/handler/component/ResponseBuilder.java b/solr/core/src/java/org/apache/solr/handler/component/ResponseBuilder.java
@@ -460,6 +460,13 @@ public void setResult(QueryResult result) {
               SolrQueryResponse.RESPONSE_HEADER_SEGMENT_TERMINATED_EARLY_KEY,
               segmentTerminatedEarly);
     }
+    final Boolean terminatedEarly = result.getTerminatedEarly();
+    if (terminatedEarly != null) {
+      rsp.getResponseHeader()
+        .add(
+          SolrQueryResponse.RESPONSE_HEADER_TERMINATED_EARLY_KEY,
+          terminatedEarly);
+    }
     if (null != cursorMark) {
       assert null != result.getNextCursorMark() : "using cursor but no next cursor set";
       this.setNextCursorMark(result.getNextCursorMark());

diff --git a/solr/core/src/java/org/apache/solr/response/SolrQueryResponse.java b/solr/core/src/java/org/apache/solr/response/SolrQueryResponse.java
@@ -68,6 +68,8 @@ public class SolrQueryResponse {
   public static final String RESPONSE_HEADER_PARTIAL_RESULTS_DETAILS_KEY = "partialResultsDetails";
   public static final String RESPONSE_HEADER_SEGMENT_TERMINATED_EARLY_KEY =
       "segmentTerminatedEarly";
+  public static final String RESPONSE_HEADER_TERMINATED_EARLY_KEY =
+    "terminatedEarly";
   public static final String RESPONSE_HEADER_KEY = "responseHeader";
   private static final String RESPONSE_KEY = "response";
 

diff --git a/solr/core/src/java/org/apache/solr/search/EarlyTerminatingCollector.java b/solr/core/src/java/org/apache/solr/search/EarlyTerminatingCollector.java
@@ -17,6 +17,8 @@
 package org.apache.solr.search;
 
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
+
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.FilterCollector;
@@ -29,11 +31,16 @@
  */
 public class EarlyTerminatingCollector extends FilterCollector {
 
+  private final int chunkSize = 100; // Check across threads only at a chunk size
+
   private final int maxDocsToCollect;
 
   private int numCollected = 0;
   private int prevReaderCumulativeSize = 0;
   private int currentReaderSize = 0;
+  private final AtomicInteger pendingDocsToCollect;
+  private boolean terminatedEarly = false;
+
 
   /**
    * Wraps a {@link Collector}, throwing {@link EarlyTerminatingCollectorException} once the
@@ -43,13 +50,14 @@ public class EarlyTerminatingCollector extends FilterCollector {
    * @param maxDocsToCollect - the maximum number of documents to Collect
    */
   public EarlyTerminatingCollector(Collector delegate, int maxDocsToCollect) {
-    super(delegate);
-    assert 0 < maxDocsToCollect;
-    assert null != delegate;
+    this(delegate,maxDocsToCollect, new AtomicInteger(maxDocsToCollect));
+  }
 
+  public EarlyTerminatingCollector(Collector delegate, int maxDocsToCollect,AtomicInteger docsToCollect) {
+    super(delegate);
     this.maxDocsToCollect = maxDocsToCollect;
+    this.pendingDocsToCollect = docsToCollect;
   }
-
   @Override
   public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
     prevReaderCumulativeSize += currentReaderSize; // not current any more
@@ -61,11 +69,22 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept
       public void collect(int doc) throws IOException {
         super.collect(doc);
         numCollected++;
-        if (maxDocsToCollect <= numCollected) {
-          throw new EarlyTerminatingCollectorException(
-              numCollected, prevReaderCumulativeSize + (doc + 1));
+        terminatedEarly = maxDocsToCollect <= numCollected;
+        if (numCollected % chunkSize == 0) {
+          terminatedEarly = pendingDocsToCollect.addAndGet(-1 * chunkSize) < 0;
+        }
+        if (terminatedEarly) {
+          throw new EarlyTerminatingCollectorException(maxDocsToCollect,
+            prevReaderCumulativeSize + (doc + 1));
         }
       }
     };
   }
+  public boolean isTerminatedEarly() {
+    return terminatedEarly;
+  }
+
+  public Collector getDelegate() {
+    return super.in;
+  }
 }
diff --git a/solr/core/src/java/org/apache/solr/search/EarlyTerminatingCollectorException.java b/solr/core/src/java/org/apache/solr/search/EarlyTerminatingCollectorException.java
@@ -22,13 +22,12 @@
  */
 public class EarlyTerminatingCollectorException extends RuntimeException {
   private static final long serialVersionUID = 5939241340763428118L;
-  private int numberScanned;
-  private int numberCollected;
+  private final int numberScanned;
+  private final int numberCollected;
 
   public EarlyTerminatingCollectorException(int numberCollected, int numberScanned) {
     assert numberCollected <= numberScanned : numberCollected + "<=" + numberScanned;
     assert 0 < numberCollected;
-    assert 0 < numberScanned;
 
     this.numberCollected = numberCollected;
     this.numberScanned = numberScanned;

diff --git a/solr/core/src/java/org/apache/solr/search/MultiThreadedSearcher.java b/solr/core/src/java/org/apache/solr/search/MultiThreadedSearcher.java
@@ -90,12 +90,16 @@ SearchResult searchCollectorManagers(
 
     @SuppressWarnings({"unchecked", "rawtypes"})
     CollectorManager<Collector, Object>[] colls = collectors.toArray(new CollectorManager[0]);
-    SolrMultiCollectorManager manager = new SolrMultiCollectorManager(colls);
+    final SolrMultiCollectorManager manager = new SolrMultiCollectorManager(cmd, colls);
+    boolean isTerminatedEarly = false;
     Object[] ret;
     try {
       ret = searcher.search(query, manager);
     } catch (Exception ex) {
-      if (ex instanceof RuntimeException
+      if (ex instanceof EarlyTerminatingCollectorException) {
+        ret = manager.reduce();
+        isTerminatedEarly = true;
+      } else if (ex instanceof RuntimeException
           && ex.getCause() != null
           && ex.getCause() instanceof ExecutionException
           && ex.getCause().getCause() != null
@@ -108,7 +112,7 @@ SearchResult searchCollectorManagers(
 
     ScoreMode scoreMode = SolrMultiCollectorManager.scoreMode(firstCollectors);
 
-    return new SearchResult(scoreMode, ret);
+    return new SearchResult(scoreMode, ret, isTerminatedEarly);
   }
 
   static boolean allowMT(DelegatingCollector postFilter, QueryCommand cmd) {
@@ -190,10 +194,16 @@ public ScoreMode scoreMode() {
   static class SearchResult {
     final ScoreMode scoreMode;
     private final Object[] result;
+    final boolean isTerminatedEarly;
+
 
     public SearchResult(ScoreMode scoreMode, Object[] result) {
+      this(scoreMode, result, false);
+    }
+    public SearchResult(ScoreMode scoreMode, Object[] result,boolean isTerminatedEarly) {
       this.scoreMode = scoreMode;
       this.result = result;
+      this.isTerminatedEarly = isTerminatedEarly;
     }
 
     public TopDocsResult getTopDocsResult() {
@@ -255,7 +265,13 @@ public Object reduce(Collection collectors) throws IOException {
       for (Iterator var4 = collectors.iterator();
           var4.hasNext();
           maxScore = Math.max(maxScore, collector.getMaxScore())) {
-        collector = (MaxScoreCollector) var4.next();
+        Collector next = (Collector) var4.next();
+        if (next instanceof EarlyTerminatingCollector) {
+          EarlyTerminatingCollector earlyTerminatingCollector = (EarlyTerminatingCollector) next;
+          next = earlyTerminatingCollector.getDelegate();
+        }
+        collector = (MaxScoreCollector) next;
+
       }
 
       return new MaxScoreResult(maxScore);
@@ -324,6 +340,10 @@ public Object reduce(Collection collectors) throws IOException {
       Collector collector;
       for (Object o : collectors) {
         collector = (Collector) o;
+        if (collector instanceof EarlyTerminatingCollector) {
+          EarlyTerminatingCollector earlyTerminatingCollector = (EarlyTerminatingCollector) collector;
+          collector = earlyTerminatingCollector.getDelegate();
+        }
         if (collector instanceof TopDocsCollector) {
           TopDocs td = ((TopDocsCollector) collector).topDocs(0, len);
           assert td != null : Arrays.asList(topDocs);

diff --git a/solr/core/src/java/org/apache/solr/search/QueryCommand.java b/solr/core/src/java/org/apache/solr/search/QueryCommand.java
@@ -42,6 +42,7 @@ public class QueryCommand {
   private int minExactCount = Integer.MAX_VALUE;
   private CursorMark cursorMark;
   private boolean distribStatsDisabled;
+  private int maxHits = Integer.MAX_VALUE;
 
   public CursorMark getCursorMark() {
     return cursorMark;
@@ -194,7 +195,7 @@ public QueryCommand setNeedDocSet(boolean needDocSet) {
   }
 
   public boolean getTerminateEarly() {
-    return (flags & SolrIndexSearcher.TERMINATE_EARLY) != 0;
+    return (flags & SolrIndexSearcher.TERMINATE_EARLY) != 0 || maxHits < Integer.MAX_VALUE;
   }
 
   public QueryCommand setTerminateEarly(boolean segmentTerminateEarly) {
@@ -245,4 +246,12 @@ public boolean isDistribStatsDisabled() {
   public QueryResult search(SolrIndexSearcher searcher) throws IOException {
     return searcher.search(this);
   }
+
+  public int getMaxHits() {
+    return maxHits;
+  }
+
+  public void setMaxHits(int maxHits) {
+    this.maxHits = maxHits;
+  }
 }
diff --git a/solr/core/src/java/org/apache/solr/search/QueryResult.java b/solr/core/src/java/org/apache/solr/search/QueryResult.java
@@ -22,6 +22,7 @@ public class QueryResult {
   // Object for back compatibility so that we render true not "true" in json
   private Object partialResults;
   private Boolean segmentTerminatedEarly;
+  private Boolean terminatedEarly;
   private DocListAndSet docListAndSet;
   private CursorMark nextCursorMark;
 
@@ -85,4 +86,12 @@ public void setNextCursorMark(CursorMark next) {
   public CursorMark getNextCursorMark() {
     return nextCursorMark;
   }
+
+  public Boolean getTerminatedEarly() {
+    return terminatedEarly;
+  }
+
+  public void setTerminatedEarly(boolean terminatedEarly) {
+    this.terminatedEarly = terminatedEarly;
+  }
 }
diff --git a/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java b/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java
@@ -295,7 +295,7 @@ private Collector buildAndRunCollectorChain(
 
     final boolean terminateEarly = cmd.getTerminateEarly();
     if (terminateEarly) {
-      collector = new EarlyTerminatingCollector(collector, cmd.getLen());
+      collector = new EarlyTerminatingCollector(collector, cmd.getMaxHits());
     }
 
     final long timeAllowed = cmd.getTimeAllowed();
@@ -329,7 +329,8 @@ private Collector buildAndRunCollectorChain(
       if (collector instanceof DelegatingCollector) {
         ((DelegatingCollector) collector).complete();
       }
-      throw etce;
+      qr.setPartialResults(true);
+
     } finally {
       if (earlyTerminatingSortingCollector != null) {
         qr.setSegmentTerminatedEarly(earlyTerminatingSortingCollector.terminatedEarly());
@@ -338,6 +339,9 @@ private Collector buildAndRunCollectorChain(
       if (cmd.isQueryCancellable()) {
         core.getCancellableQueryTracker().removeCancellableQuery(cmd.getQueryID());
       }
+      if (collector instanceof EarlyTerminatingCollector) {
+        qr.setTerminatedEarly(((EarlyTerminatingCollector) collector).isTerminatedEarly());
+      }
     }
     if (collector instanceof DelegatingCollector) {
       ((DelegatingCollector) collector).complete();
@@ -1979,7 +1983,10 @@ public ScoreMode scoreMode() {
         MultiThreadedSearcher.TopDocsResult topDocsResult = searchResult.getTopDocsResult();
         totalHits = topDocsResult.totalHits;
         topDocs = topDocsResult.topDocs;
-
+        if(searchResult.isTerminatedEarly){
+          qr.setTerminatedEarly(true);
+          qr.setPartialResults(Boolean.TRUE);
+        }
         maxScore = searchResult.getMaxScore(totalHits);
       }