SOLR-17841: Use DocSetCollector for multithreaded DocSet(Faceted) search

changelog/unreleased/SOLR-17841-mt-docset-collector.yml

@@ -0,0 +1,9 @@
+title: >
+  Improved multiThreaded=true performance when a docset is needed (e.g. faceting).
+type: changed
+authors:
+  - name: Puneet Ahuja
+    nick: punAhuja
+links:
+  - name: SOLR-17841
+    url: https://issues.apache.org/jira/browse/SOLR-17841

solr/core/src/java/org/apache/solr/search/MultiThreadedSearcher.java

@@ -22,14 +22,11 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Iterator;
-import java.util.LinkedList;
 import java.util.concurrent.ExecutionException;
-import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.CollectorManager;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.SimpleCollector;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopDocsCollector;
 import org.apache.lucene.search.TopFieldDocs;
@@ -119,11 +116,14 @@ SearchResult searchCollectorManagers(
     return new SearchResult(scoreMode, ret);
   }
 
-  static boolean allowMT(DelegatingCollector postFilter, QueryCommand cmd) {
+  static boolean allowMT(DelegatingCollector postFilter, QueryCommand cmd, boolean hasExecutor) {
     // TODO: it's unclear if segmentTerminateEarly is truly incompatible but
     //  since it has to appropriately denote partial results this needs to be
     //  investigated/tested before we can remove this check (perhaps for 9.8).
-    return postFilter == null && !cmd.getSegmentTerminateEarly() && cmd.getMultiThreaded();
+    return postFilter == null
+        && !cmd.getSegmentTerminateEarly()
+        && cmd.getMultiThreaded()
+        && hasExecutor;
   }
 
   static class MaxScoreResult {
@@ -134,67 +134,6 @@ public MaxScoreResult(float maxScore) {
     }
   }
 
-  static class FixedBitSetCollector extends SimpleCollector {
-    @SuppressWarnings("JdkObsolete")
-    private final LinkedList<FixedBitSet> bitSets = new LinkedList<>();
-
-    @SuppressWarnings("JdkObsolete")
-    private final LinkedList<Integer> skipWords = new LinkedList<>();
-
-    @SuppressWarnings("JdkObsolete")
-    private final LinkedList<Integer> skipBits = new LinkedList<>();
-
-    FixedBitSetCollector() {}
-
-    @Override
-    protected void doSetNextReader(LeafReaderContext context) throws IOException {
-      this.bitSets.add(null); // lazy allocate when collecting document(s)
-      this.skipWords.add(context.docBase / 64);
-      this.skipBits.add(context.docBase % 64);
-    }
-
-    @Override
-    public void collect(int doc) throws IOException {
-      FixedBitSet bitSet = this.bitSets.getLast();
-      final int idx = this.skipBits.getLast() + doc;
-
-      final int numWords = FixedBitSet.bits2words(idx + 1); // +1 to ensure minimum 1 word
-
-      if (bitSet == null) {
-        this.bitSets.removeLast();
-        bitSet = new FixedBitSet(numWords * 64);
-        this.bitSets.addLast(bitSet);
-
-      } else if (bitSet.getBits().length < numWords) {
-        FixedBitSet smallerBitSet = this.bitSets.removeLast();
-        bitSet = new FixedBitSet(numWords * 64);
-        bitSet.xor(smallerBitSet);
-        this.bitSets.addLast(bitSet);
-      }
-
-      bitSet.set(idx);
-    }
-
-    void update(FixedBitSet allBitSet) {
-      final long[] allBits = allBitSet.getBits();
-      for (int bs_idx = 0; bs_idx < this.bitSets.size(); ++bs_idx) {
-        final FixedBitSet itBitSet = this.bitSets.get(bs_idx);
-        if (itBitSet != null) {
-          final int skipWords = this.skipWords.get(bs_idx);
-          final long[] itBits = itBitSet.getBits();
-          for (int idx = 0; idx < itBits.length && skipWords + idx < allBits.length; ++idx) {
-            allBits[skipWords + idx] ^= itBits[idx];
-          }
-        }
-      }
-    }
-
-    @Override
-    public ScoreMode scoreMode() {
-      return ScoreMode.COMPLETE_NO_SCORES;
-    }
-  }
-
   static class SearchResult {
     final ScoreMode scoreMode;
     private final Object[] result;
@@ -283,21 +222,34 @@ public DocSetCM(int maxDoc) {
 
     @Override
     public Collector newCollector() throws IOException {
-      // TODO: add to firstCollectors here? or if not have comment w.r.t. why not adding
-      return new FixedBitSetCollector();
+      return new DocSetCollector(maxDoc);
     }
 
     @Override
     @SuppressWarnings({"rawtypes"})
     public Object reduce(Collection collectors) throws IOException {
       final FixedBitSet reduced = new FixedBitSet(maxDoc);
       for (Object collector : collectors) {
-        if (collector instanceof FixedBitSetCollector fixedBitSetCollector) {
-          fixedBitSetCollector.update(reduced);
+        if (collector instanceof EarlyTerminatingCollector earlyTerminatingCollector) {
+          collector = earlyTerminatingCollector.getDelegate();
+        }
+        if (collector instanceof DocSetCollector docSetCollector) {
+          mergeDocSetIntoFixedBitSet(docSetCollector.getDocSet(), reduced);
         }
       }
       return reduced;
     }
+
+    private static void mergeDocSetIntoFixedBitSet(DocSet docSet, FixedBitSet reduced) {
+      if (docSet instanceof BitDocSet bitDocSet) {
+        reduced.or(bitDocSet.getBits());
+      } else {
+        DocIterator iter = docSet.iterator();
+        while (iter.hasNext()) {
+          reduced.set(iter.nextDoc());
+        }
+      }
+    }
   }
 
   private class TopDocsCM implements CollectorManager<Collector, Object> {

solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java

@@ -1982,7 +1982,7 @@ public ScoreMode scoreMode() {
       }
       final TopDocs topDocs;
       final ScoreMode scoreModeUsed;
-      if (!MultiThreadedSearcher.allowMT(pf.postFilter, cmd)) {
+      if (!MultiThreadedSearcher.allowMT(pf.postFilter, cmd, getTaskExecutor() != null)) {
         log.trace("SINGLE THREADED search, skipping collector manager in getDocListNC");
         final TopDocsCollector<?> topCollector = buildTopDocsCollector(len, cmd);
         MaxScoreCollector maxScoreCollector = null;
@@ -2093,7 +2093,7 @@ public ScoreMode scoreMode() {
       qr.setNextCursorMark(cmd.getCursorMark());
     } else {
       final TopDocs topDocs;
-      if (!MultiThreadedSearcher.allowMT(pf.postFilter, cmd)) {
+      if (!MultiThreadedSearcher.allowMT(pf.postFilter, cmd, getTaskExecutor() != null)) {
         log.trace("SINGLE THREADED search, skipping collector manager in getDocListAndSetNC");
 
         @SuppressWarnings({"rawtypes"})

solr/core/src/test/org/apache/solr/search/TestMultiThreadedSearcher.java

@@ -123,6 +123,37 @@ public void testReRankWithMultiThreadedSearch() throws Exception {
             });
   }
 
+  /** Multi-threaded DocSet collection must match single-threaded for the same query. */
+  public void testMultiThreadedDocSetMatchesSingleThreaded() throws Exception {
+    h.getCore()
+        .withSearcher(
+            searcher -> {
+              assertTrue(searcher.getSlices().length > 1);
+
+              final Query query = new TermQuery(new Term("field1_s", "xyzrareterm"));
+              final QueryCommand cmdSingle = new QueryCommand();
+              cmdSingle.setQuery(query);
+              cmdSingle.setNeedDocSet(true);
+              cmdSingle.setLen(10);
+              cmdSingle.setMultiThreaded(false);
+
+              final QueryCommand cmdMulti = new QueryCommand();
+              cmdMulti.setQuery(query);
+              cmdMulti.setNeedDocSet(true);
+              cmdMulti.setLen(10);
+              cmdMulti.setMultiThreaded(true);
+
+              final QueryResult singleThreaded = searcher.search(cmdSingle);
+              final QueryResult multiThreaded = searcher.search(cmdMulti);
+
+              final DocSet stSet = singleThreaded.getDocListAndSet().docSet;
+              final DocSet mtSet = multiThreaded.getDocListAndSet().docSet;
+              assertEquals(stSet.size(), mtSet.size());
+              assertTrue(DocSetUtil.equals(stSet, mtSet));
+              return null;
+            });
+  }
+
   private static final class SimpleReRankQuery extends RankQuery {
 
     private Query q;