Optimize sub aggregation using bulk collection lucene apis

CHANGELOG.md

@@ -31,6 +31,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Support pull-based ingestion message mappers and raw payload support ([#19765](https://github.com/opensearch-project/OpenSearch/pull/19765)]
 
 ### Changed
+- Combining filter rewrite and skip list to optimize sub aggregation([#19573](https://github.com/opensearch-project/OpenSearch/pull/19573))
 - Faster `terms` query creation for `keyword` field with index and docValues enabled ([#19350](https://github.com/opensearch-project/OpenSearch/pull/19350))
 - Refactor to move prepareIndex and prepareDelete methods to Engine class ([#19551](https://github.com/opensearch-project/OpenSearch/pull/19551))
 - Omit maxScoreCollector in SimpleTopDocsCollectorContext when concurrent segment search enabled ([#19584](https://github.com/opensearch-project/OpenSearch/pull/19584))

server/src/main/java/org/opensearch/search/aggregations/LeafBucketCollector.java

@@ -126,6 +126,12 @@ public void collect(int doc) throws IOException {
         collect(doc, 0);
     }
 
+    public void collect(int[] docIds, long owningBucketOrd) throws IOException {
+        for (int doc : docIds) {
+            collect(doc, owningBucketOrd);
+        }
+    }
+
     @Override
     public void collect(DocIdStream stream) throws IOException {
         collect(stream, 0);

server/src/main/java/org/opensearch/search/aggregations/bucket/HistogramSkiplistLeafCollector.java

@@ -0,0 +1,173 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.aggregations.bucket;
+
+import org.apache.lucene.index.DocValuesSkipper;
+import org.apache.lucene.index.NumericDocValues;
+import org.apache.lucene.search.DocIdStream;
+import org.apache.lucene.search.Scorable;
+import org.opensearch.common.Rounding;
+import org.opensearch.search.aggregations.LeafBucketCollector;
+import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds;
+
+import java.io.IOException;
+
+/**
+ * Histogram collection logic using skip list.
+ *
+ * @opensearch.internal
+ */
+public class HistogramSkiplistLeafCollector extends LeafBucketCollector {
+
+    private final NumericDocValues values;
+    private final DocValuesSkipper skipper;
+    private final Rounding.Prepared preparedRounding;
+    private final LongKeyedBucketOrds bucketOrds;
+    private final LeafBucketCollector sub;
+    private final BucketsAggregator aggregator;
+
+    /**
+     * Max doc ID (inclusive) up to which all docs values may map to the same
+     * bucket.
+     */
+    private int upToInclusive = -1;
+
+    /**
+     * Whether all docs up to {@link #upToInclusive} values map to the same bucket.
+     */
+    private boolean upToSameBucket;
+
+    /**
+     * Index in bucketOrds for docs up to {@link #upToInclusive}.
+     */
+    private long upToBucketIndex;
+
+    public HistogramSkiplistLeafCollector(
+        NumericDocValues values,
+        DocValuesSkipper skipper,
+        Rounding.Prepared preparedRounding,
+        LongKeyedBucketOrds bucketOrds,
+        LeafBucketCollector sub,
+        BucketsAggregator aggregator
+    ) {
+        this.values = values;
+        this.skipper = skipper;
+        this.preparedRounding = preparedRounding;
+        this.bucketOrds = bucketOrds;
+        this.sub = sub;
+        this.aggregator = aggregator;
+    }
+
+    @Override
+    public void setScorer(Scorable scorer) throws IOException {
+        if (sub != null) {
+            sub.setScorer(scorer);
+        }
+    }
+
+    private void advanceSkipper(int doc, long owningBucketOrd) throws IOException {
+        if (doc > skipper.maxDocID(0)) {
+            skipper.advance(doc);
+        }
+        upToSameBucket = false;
+
+        if (skipper.minDocID(0) > doc) {
+            // Corner case which happens if `doc` doesn't have a value and is between two
+            // intervals of
+            // the doc-value skip index.
+            upToInclusive = skipper.minDocID(0) - 1;
+            return;
+        }
+
+        upToInclusive = skipper.maxDocID(0);
+
+        // Now find the highest level where all docs map to the same bucket.
+        for (int level = 0; level < skipper.numLevels(); ++level) {
+            int totalDocsAtLevel = skipper.maxDocID(level) - skipper.minDocID(level) + 1;
+            long minBucket = preparedRounding.round(skipper.minValue(level));
+            long maxBucket = preparedRounding.round(skipper.maxValue(level));
+
+            if (skipper.docCount(level) == totalDocsAtLevel && minBucket == maxBucket) {
+                // All docs at this level have a value, and all values map to the same bucket.
+                upToInclusive = skipper.maxDocID(level);
+                upToSameBucket = true;
+                upToBucketIndex = bucketOrds.add(owningBucketOrd, maxBucket);
+                if (upToBucketIndex < 0) {
+                    upToBucketIndex = -1 - upToBucketIndex;
+                }
+            } else {
+                break;
+            }
+        }
+    }
+
+    @Override
+    public void collect(int doc, long owningBucketOrd) throws IOException {
+        if (doc > upToInclusive) {
+            advanceSkipper(doc, owningBucketOrd);
+        }
+
+        if (upToSameBucket) {
+            aggregator.incrementBucketDocCount(upToBucketIndex, 1L);
+            sub.collect(doc, upToBucketIndex);
+        } else if (values.advanceExact(doc)) {
+            final long value = values.longValue();
+            long bucketIndex = bucketOrds.add(owningBucketOrd, preparedRounding.round(value));
+            if (bucketIndex < 0) {
+                bucketIndex = -1 - bucketIndex;
+                aggregator.collectExistingBucket(sub, doc, bucketIndex);
+            } else {
+                aggregator.collectBucket(sub, doc, bucketIndex);
+            }
+        }
+    }
+
+    @Override
+    public void collect(DocIdStream stream) throws IOException {
+        // This will only be called if its the top agg
+        collect(stream, 0);
+    }
+
+    @Override
+    public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+        // This will only be called if its the sub aggregation
+        for (;;) {
+            int upToExclusive = upToInclusive + 1;
+            if (upToExclusive < 0) { // overflow
+                upToExclusive = Integer.MAX_VALUE;
+            }
+
+            if (upToSameBucket) {
+                if (sub == NO_OP_COLLECTOR) {
+                    // stream.count maybe faster when we don't need to handle sub-aggs
+                    long count = stream.count(upToExclusive);
+                    aggregator.incrementBucketDocCount(upToBucketIndex, count);
+                } else {
+                    int count = 0;
+                    int[] docBuffer = new int[64];
+                    int cnt = Integer.MAX_VALUE;
+                    while (cnt != 0) {
+                        cnt = stream.intoArray(upToExclusive, docBuffer);
+                        sub.collect(docBuffer, upToBucketIndex);
+                        count += cnt;
+                    }
+                    aggregator.incrementBucketDocCount(upToBucketIndex, count);
+                }
+            } else {
+                stream.forEach(upToExclusive, doc -> collect(doc, owningBucketOrd));
+            }
+
+            if (stream.mayHaveRemaining()) {
+                advanceSkipper(upToExclusive, owningBucketOrd);
+            } else {
+                break;
+            }
+        }
+    }
+}

server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/BFSCollector.java

@@ -0,0 +1,24 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.aggregations.bucket.filterrewrite;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.opensearch.search.aggregations.LeafBucketCollector;
+
+import java.io.IOException;
+
+/**
+ * Workaround for collectors that cannot handle DFS travel, i.e. changing owningBucketOrd)
+ *
+ * @opensearch.internal
+ */
+public interface BFSCollector {
+
+    LeafBucketCollector getBFSLeafCollector(LeafReaderContext ctx) throws IOException;
+}

server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/rangecollector/SubAggRangeCollector.java

@@ -17,6 +17,7 @@
 import org.apache.lucene.util.FixedBitSet;
 import org.opensearch.search.aggregations.BucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollector;
+import org.opensearch.search.aggregations.bucket.filterrewrite.BFSCollector;
 import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext;
 import org.opensearch.search.aggregations.bucket.filterrewrite.Ranges;
 
@@ -105,7 +106,12 @@ public void finalizePreviousRange() {
         // trigger the sub agg collection for this range
         try {
             // build a new leaf collector for each bucket
-            LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(leafCtx);
+            LeafBucketCollector sub = null;
+            if (collectableSubAggregators instanceof BFSCollector bfsCollector) {
+                sub = bfsCollector.getBFSLeafCollector(leafCtx);
+            } else {
+                sub = collectableSubAggregators.getLeafCollector(leafCtx);
+            }
             sub.collect(DocIdStreamHelper.getDocIdStream(bitSet), bucketOrd);
             logger.trace("collected sub aggregation for bucket {}", bucketOrd);
         } catch (IOException e) {