trying the each bucket method

Signed-off-by: bowenlan-amzn <[email protected]>
bowenlan-amzn · Feb 3, 2025 · 0a8ee9a · 0a8ee9a
1 parent a7f8a5f
commit 0a8ee9a
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 15 deletions.
diff --git a/...rc/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/...rc/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java
@@ -561,7 +561,7 @@ protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucket
             ctx,
             this::incrementBucketDocCount,
             segmentMatchAll(context, ctx),
-            sub
+            collectableSubAggregators
         );
         if (optimized) throw new CollectionTerminatedException();
 

diff --git a/...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java b/...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java
@@ -15,7 +15,9 @@
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.PointValues;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.DocIdSetBuilder;
 import org.opensearch.index.mapper.DocCountFieldMapper;
+import org.opensearch.search.aggregations.BucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.internal.SearchContext;
 
@@ -103,7 +105,7 @@ public boolean tryOptimize(
         final LeafReaderContext leafCtx,
         final BiConsumer<Long, Long> incrementDocCount,
         boolean segmentMatchAll,
-        LeafBucketCollector sub
+        BucketCollector collectableSubAggregators
     ) throws IOException {
         segments.incrementAndGet();
         if (!canOptimize) {
@@ -143,11 +145,29 @@ public boolean tryOptimize(
         // 1. List of Iterators per ranges
         // 2. Composite iterator
 
-        CompositeDocIdSetIterator iter = new CompositeDocIdSetIterator(debugInfo.iterators);
-        while (iter.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
-            int currentDoc = iter.docID();
-            int bucket = iter.getCurrentBucket();
-            sub.collect(currentDoc, bucket);
+        // CompositeDocIdSetIterator iter = new CompositeDocIdSetIterator(debugInfo.iterators);
+        // while (iter.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
+        // int currentDoc = iter.docID();
+        // int bucket = iter.getCurrentBucket();
+        // sub.collect(currentDoc, bucket);
+        // }
+
+        // let's not use composite disi
+        // try rebuilding the subagg leaf collector
+        // for each bucket ord
+
+        LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(leafCtx);
+        for (int bucketOrd = 0; bucketOrd < debugInfo.builders.length; bucketOrd++) {
+            DocIdSetIterator iterator = debugInfo.builders[bucketOrd].build().iterator();
+            if (iterator == null) {
+                continue;
+            }
+            while (iterator.nextDoc() != NO_MORE_DOCS) {
+                int currentDoc = iterator.docID();
+                sub.collect(currentDoc, bucketOrd);
+            }
+            // resetting the sub collector after processing each bucket
+            sub = collectableSubAggregators.getLeafCollector(leafCtx);
         }
 
         return true;
@@ -187,6 +207,7 @@ static class DebugInfo {
         private final AtomicInteger innerNodeVisited = new AtomicInteger(); // inner node visited
 
         public DocIdSetIterator[] iterators;
+        public DocIdSetBuilder[] builders;
 
         void visitLeaf() {
             leafNodeVisited.incrementAndGet();

diff --git a/...main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java b/...main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java
@@ -102,14 +102,21 @@ static FilterRewriteOptimizationContext.DebugInfo multiRangesTraverse(
         Map<Long, DocIdSetBuilder> ordinalToBuilder = collector.bucketOrdinalToDocIdSetBuilder;
         logger.debug("keys of bucketOrdinalToDocIdSetBuilder: {}", ordinalToBuilder.keySet());
         int maxOrdinal = ordinalToBuilder.keySet().stream().mapToInt(Long::intValue).max().orElse(0) + 1;
-        DocIdSetIterator[] iterators = new DocIdSetIterator[maxOrdinal];
+        // DocIdSetIterator[] iterators = new DocIdSetIterator[maxOrdinal];
+        // for (Map.Entry<Long, DocIdSetBuilder> entry : ordinalToBuilder.entrySet()) {
+        //     int ordinal = Math.toIntExact(entry.getKey());
+        //     DocIdSetBuilder builder = entry.getValue();
+        //     DocIdSet docIdSet = builder.build();
+        //     iterators[ordinal] = docIdSet.iterator();
+        // }
+        // debugInfo.iterators = iterators;
+
+        DocIdSetBuilder[] builder = new DocIdSetBuilder[maxOrdinal];
         for (Map.Entry<Long, DocIdSetBuilder> entry : ordinalToBuilder.entrySet()) {
             int ordinal = Math.toIntExact(entry.getKey());
-            DocIdSetBuilder builder = entry.getValue();
-            DocIdSet docIdSet = builder.build();
-            iterators[ordinal] = docIdSet.iterator();
+            builder[ordinal] = entry.getValue();
         }
-        debugInfo.iterators = iterators;
+        debugInfo.builders = builder;
 
         return debugInfo;
     }

diff --git a/...java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java b/...java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java
@@ -240,7 +240,7 @@ public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBuc
             ctx,
             this::incrementBucketDocCount,
             segmentMatchAll(context, ctx),
-            sub
+            collectableSubAggregators
         );
         if (optimized) throw new CollectionTerminatedException();
 

diff --git a/...ain/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/...ain/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java
@@ -193,7 +193,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol
             ctx,
             this::incrementBucketDocCount,
             segmentMatchAll(context, ctx),
-            sub
+            collectableSubAggregators
         );
         if (optimized) {
             throw new CollectionTerminatedException();

diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java
@@ -311,7 +311,8 @@ public ScoreMode scoreMode() {
 
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
-        if (segmentMatchAll(context, ctx) && filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false, sub)) {
+        if (segmentMatchAll(context, ctx)
+            && filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false, collectableSubAggregators)) {
             throw new CollectionTerminatedException();
         }