Lucene query improvements (#311)

* Rework lucene collectors, add executor service

Reworks collectors to use executor service, removes unnecceary collectors depending on provided params

* Simplify stats collector logic

* PR feedback - add indexing sorting

* Revert "Remove overrides for ramBuffer, compound file"

This reverts commit e82bddd.

* PR feedback

* Fix count discrepancy

* PR feedback

Author: bryanlb

kaldb/src/main/java/com/slack/kaldb/logstore/LuceneIndexStoreImpl.java

@@ -21,6 +21,8 @@
 import org.apache.lucene.index.KeepOnlyLastCommitDeletionPolicy;
 import org.apache.lucene.index.SnapshotDeletionPolicy;
 import org.apache.lucene.search.SearcherManager;
+import org.apache.lucene.search.Sort;
+import org.apache.lucene.search.SortField;
 import org.apache.lucene.store.FSDirectory;
 import org.apache.lucene.store.MMapDirectory;
 import org.slf4j.Logger;
@@ -139,6 +141,14 @@ private IndexWriterConfig buildIndexWriterConfig(
         new IndexWriterConfig(analyzer)
             .setOpenMode(IndexWriterConfig.OpenMode.CREATE)
             .setMergeScheduler(new KalDBMergeScheduler(metricsRegistry))
+            // we sort by timestamp descending, as that is the order we expect to return results the
+            // majority of the time
+            .setIndexSort(
+                new Sort(
+                    new SortField(
+                        LogMessage.SystemField.TIME_SINCE_EPOCH.fieldName,
+                        SortField.Type.LONG,
+                        true)))
             .setIndexDeletionPolicy(snapshotDeletionPolicy);
 
     if (config.enableTracing) {

kaldb/src/main/java/com/slack/kaldb/logstore/search/LogIndexSearcherImpl.java

@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
@@ -29,16 +30,17 @@
 import org.apache.lucene.queryparser.classic.QueryParser;
 import org.apache.lucene.search.BooleanClause.Occur;
 import org.apache.lucene.search.BooleanQuery.Builder;
-import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.CollectorManager;
 import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.MultiCollector;
+import org.apache.lucene.search.MultiCollectorManager;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.SearcherManager;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.SortField.Type;
 import org.apache.lucene.search.TopFieldCollector;
+import org.apache.lucene.search.TopFieldDocs;
 import org.apache.lucene.store.MMapDirectory;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -111,30 +113,41 @@ public SearchResult<LogMessage> search(
       // This is a useful optimization for indexes that are static.
       IndexSearcher searcher = searcherManager.acquire();
       try {
-        TopFieldCollector topFieldCollector = buildTopFieldCollector(howMany);
-        StatsCollector statsCollector =
+        List<LogMessage> results;
+        Histogram histogram = new NoOpHistogramImpl();
+
+        CollectorManager<StatsCollector, Histogram> statsCollector =
             buildStatsCollector(bucketCount, startTimeMsEpoch, endTimeMsEpoch);
-        Collector collectorChain = MultiCollector.wrap(topFieldCollector, statsCollector);
 
-        searcher.search(query, collectorChain);
-        List<LogMessage> results;
         if (howMany > 0) {
-          ScoreDoc[] hits = topFieldCollector.topDocs().scoreDocs;
+          CollectorManager<TopFieldCollector, TopFieldDocs> topFieldCollector =
+              buildTopFieldCollector(howMany, bucketCount > 0 ? Integer.MAX_VALUE : howMany);
+          MultiCollectorManager collectorManager;
+          if (bucketCount > 0) {
+            collectorManager = new MultiCollectorManager(topFieldCollector, statsCollector);
+          } else {
+            collectorManager = new MultiCollectorManager(topFieldCollector);
+          }
+          Object[] collector = searcher.search(query, collectorManager);
+
+          ScoreDoc[] hits = ((TopFieldDocs) collector[0]).scoreDocs;
           results = new ArrayList<>(hits.length);
           for (ScoreDoc hit : hits) {
             results.add(buildLogMessage(searcher, hit));
           }
+          if (bucketCount > 0) {
+            histogram = ((Histogram) collector[1]);
+          }
         } else {
           results = Collections.emptyList();
+          histogram = searcher.search(query, statsCollector);
         }
 
-        Histogram histogram = statsCollector.histogram;
-
         elapsedTime.stop();
         return new SearchResult<>(
             results,
             elapsedTime.elapsed(TimeUnit.MICROSECONDS),
-            histogram.count(),
+            bucketCount > 0 ? histogram.count() : results.size(),
             histogram.getBuckets(),
             0,
             0,
@@ -167,22 +180,50 @@ private LogMessage buildLogMessage(IndexSearcher searcher, ScoreDoc hit) {
     }
   }
 
-  private TopFieldCollector buildTopFieldCollector(int howMany) {
+  /**
+   * Builds a top field collector for the requested amount of results, with the option to set the
+   * totalHitsThreshold. If the totalHitsThreshold is set to Integer.MAX_VALUE it will force a
+   * ScoreMode.COMPLETE, iterating over all documents at the expense of a longer query time. This
+   * value can be set to equal howMany to allow early exiting (ScoreMode.TOP_SCORES), but should
+   * only be done when all collectors are tolerant of an early exit.
+   */
+  private CollectorManager<TopFieldCollector, TopFieldDocs> buildTopFieldCollector(
+      int howMany, int totalHitsThreshold) {
     if (howMany > 0) {
       SortField sortField = new SortField(SystemField.TIME_SINCE_EPOCH.fieldName, Type.LONG, true);
-      return TopFieldCollector.create(new Sort(sortField), howMany, howMany);
+      return TopFieldCollector.createSharedManager(
+          new Sort(sortField), howMany, null, totalHitsThreshold);
     } else {
       return null;
     }
   }
 
-  private StatsCollector buildStatsCollector(
+  private CollectorManager<StatsCollector, Histogram> buildStatsCollector(
       int bucketCount, long startTimeMsEpoch, long endTimeMsEpoch) {
     Histogram histogram =
         bucketCount > 0
             ? new FixedIntervalHistogramImpl(startTimeMsEpoch, endTimeMsEpoch, bucketCount)
             : new NoOpHistogramImpl();
-    return new StatsCollector(histogram);
+
+    return new CollectorManager<>() {
+      @Override
+      public StatsCollector newCollector() {
+        return new StatsCollector(histogram);
+      }
+
+      @Override
+      public Histogram reduce(Collection<StatsCollector> collectors) {
+        Histogram histogram = null;
+        for (StatsCollector collector : collectors) {
+          if (histogram == null) {
+            histogram = collector.getHistogram();
+          } else {
+            histogram.mergeHistogram(collector.getHistogram().getBuckets());
+          }
+        }
+        return histogram;
+      }
+    };
   }
 
   private Query buildQuery(

kaldb/src/main/java/com/slack/kaldb/logstore/search/StatsCollector.java

@@ -3,15 +3,12 @@
 import com.slack.kaldb.histogram.Histogram;
 import com.slack.kaldb.logstore.LogMessage.SystemField;
 import java.io.IOException;
-import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
-import org.apache.lucene.search.Collector;
-import org.apache.lucene.search.LeafCollector;
-import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.SimpleCollector;
 
-public class StatsCollector implements Collector {
+public class StatsCollector extends SimpleCollector {
 
   public final Histogram histogram;
   private NumericDocValues docValues;
@@ -23,26 +20,24 @@ public StatsCollector(Histogram histogram) {
   }
 
   @Override
-  public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
-    LeafReader reader = context.reader();
-    docValues = reader.getNumericDocValues(SystemField.TIME_SINCE_EPOCH.fieldName);
-
-    return new LeafCollector() {
-      @Override
-      public void setScorer(Scorable scorer) {}
-
-      @Override
-      public void collect(int doc) throws IOException {
-        if (docValues != null && docValues.advanceExact(doc)) {
-          long timestamp = docValues.longValue();
-          histogram.add(timestamp);
-        }
-      }
-    };
+  protected void doSetNextReader(final LeafReaderContext context) throws IOException {
+    docValues = context.reader().getNumericDocValues(SystemField.TIME_SINCE_EPOCH.fieldName);
+  }
+
+  public Histogram getHistogram() {
+    return histogram;
   }
 
   @Override
   public ScoreMode scoreMode() {
     return ScoreMode.COMPLETE_NO_SCORES;
   }
+
+  @Override
+  public void collect(int doc) throws IOException {
+    if (docValues != null && docValues.advanceExact(doc)) {
+      long timestamp = docValues.longValue();
+      histogram.add(timestamp);
+    }
+  }
 }