diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java index 4088e2ca5196d..87ebdb48f3064 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java @@ -22,7 +22,6 @@ 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; import java.io.IOException; @@ -162,21 +161,21 @@ public boolean tryOptimize( // try rebuilding the subagg leaf collector // for each bucket ord - LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(leafCtx); - for (int bucketOrd = 0; bucketOrd < debugInfo.builders.length; bucketOrd++) { - logger.debug("Collecting bucket {} for sub aggregation", bucketOrd); - DocIdSetBuilder builder = debugInfo.builders[bucketOrd]; - if (builder == null) { - continue; - } - DocIdSetIterator iterator = debugInfo.builders[bucketOrd].build().iterator(); - 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); - } + // LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(leafCtx); + // for (int bucketOrd = 0; bucketOrd < debugInfo.builders.length; bucketOrd++) { + // logger.debug("Collecting bucket {} for sub aggregation", bucketOrd); + // DocIdSetBuilder builder = debugInfo.builders[bucketOrd]; + // if (builder == null) { + // continue; + // } + // DocIdSetIterator iterator = debugInfo.builders[bucketOrd].build().iterator(); + // 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; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java index 57e76a4428a3e..4955197a946d1 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java @@ -82,25 +82,9 @@ static FilterRewriteOptimizationContext.DebugInfo multiRangesTraverse( } collector.finalizePreviousRange(); - // DocIdSetBuilder[] builders = collector.docIdSetBuilders; - // logger.debug("length of docIdSetBuilders: {}", builders.length); - // int totalCount = 0; - // for (int i = 0; i < builders.length; i++) { - // if (builders[i] != null) { - // int count = 0; - // DocIdSetIterator iterator = builders[i].build().iterator(); - // while (iterator.nextDoc() != NO_MORE_DOCS) { - // count++; - // } - // logger.trace(" docIdSetBuilder[{}] disi has documents: {}", i, count); - // totalCount += count; - // } - // } - // logger.debug("total count of documents from docIdSetBuilder: {}", totalCount); - - Map ordinalToBuilder = collector.bucketOrdinalToDocIdSetBuilder; - logger.debug("keys of bucketOrdinalToDocIdSetBuilder: {}", ordinalToBuilder.keySet()); - int maxOrdinal = ordinalToBuilder.keySet().stream().mapToInt(Long::intValue).max().orElse(0) + 1; + // Map 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]; // for (Map.Entry entry : ordinalToBuilder.entrySet()) { // int ordinal = Math.toIntExact(entry.getKey()); @@ -110,12 +94,12 @@ static FilterRewriteOptimizationContext.DebugInfo multiRangesTraverse( // } // debugInfo.iterators = iterators; - DocIdSetBuilder[] builder = new DocIdSetBuilder[maxOrdinal]; - for (Map.Entry entry : ordinalToBuilder.entrySet()) { - int ordinal = Math.toIntExact(entry.getKey()); - builder[ordinal] = entry.getValue(); - } - debugInfo.builders = builder; + // DocIdSetBuilder[] builder = new DocIdSetBuilder[maxOrdinal]; + // for (Map.Entry entry : ordinalToBuilder.entrySet()) { + // int ordinal = Math.toIntExact(entry.getKey()); + // builder[ordinal] = entry.getValue(); + // } + // debugInfo.builders = builder; return debugInfo; } @@ -131,7 +115,7 @@ private static void intersectWithRanges( switch (r) { case CELL_INSIDE_QUERY: collector.countNode((int) pointTree.size()); - pointTree.visitDocIDs(visitor); + // pointTree.visitDocIDs(visitor); debug.visitInner(); break; case CELL_CROSSES_QUERY: @@ -154,28 +138,28 @@ private static PointValues.IntersectVisitor getIntersectVisitor(RangeCollectorFo @Override public void grow(int count) { - collector.grow(count); + // collector.grow(count); } @Override public void visit(int docID) { // this branch should be unreachable - // throw new UnsupportedOperationException( - // "This IntersectVisitor does not perform any actions on a " + "docID=" + docID + " node being visited" - // ); - collector.collectDocId(docID); + throw new UnsupportedOperationException( + "This IntersectVisitor does not perform any actions on a " + "docID=" + docID + " node being visited" + ); + // collector.collectDocId(docID); } @Override public void visit(DocIdSetIterator iterator) throws IOException { - collector.collectDocIdSet(iterator); + // collector.collectDocIdSet(iterator); } @Override public void visit(int docID, byte[] packedValue) throws IOException { visitPoints(packedValue, collector::count); - collector.collectDocId(docID); + // collector.collectDocId(docID); } @Override @@ -186,7 +170,7 @@ public void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOExcept } }); - collector.collectDocIdSet(iterator); + // collector.collectDocIdSet(iterator); } private void visitPoints(byte[] packedValue, CheckedRunnable collect) throws IOException { @@ -261,21 +245,11 @@ private void count() { } private void collectDocId(int docId) { - // if (docIdSetBuilders[activeIndex] == null) { - // // TODO hard code for now, should be controlled by intersector grow - // docIdSetBuilders[activeIndex] = disBuilderSupplier.get(); - // currentAdder = docIdSetBuilders[activeIndex].grow(1000); - // } logger.trace("collect docId {}", docId); currentAdder.add(docId); } private void collectDocIdSet(DocIdSetIterator iter) throws IOException { - // if (docIdSetBuilders[activeIndex] == null) { - // // TODO hard code for now, should be controlled by intersector grow - // docIdSetBuilders[activeIndex] = disBuilderSupplier.get(); - // currentAdder = docIdSetBuilders[activeIndex].grow(1000); - // } logger.trace("collect disi {}", iter); currentAdder.add(iter); } @@ -298,11 +272,11 @@ private void finalizePreviousRange() { counter = 0; } - long bucketOrd = getBucketOrd.apply(activeIndex); - if (docIdSetBuilders[activeIndex] != null) { - logger.trace("finalize docIdSetBuilder[{}] with bucket ordinal {}", activeIndex, bucketOrd); - bucketOrdinalToDocIdSetBuilder.put(bucketOrd, docIdSetBuilders[activeIndex]); - } + // long bucketOrd = getBucketOrd.apply(activeIndex); + // if (docIdSetBuilders[activeIndex] != null) { + // logger.trace("finalize docIdSetBuilder[{}] with bucket ordinal {}", activeIndex, bucketOrd); + // bucketOrdinalToDocIdSetBuilder.put(bucketOrd, docIdSetBuilders[activeIndex]); + // } } /** 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 index 3b2655518293d..d2334c9e78c88 100644 --- 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 @@ -317,34 +317,41 @@ public ScoreMode scoreMode() { @Override public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException { - if (segmentMatchAll(context, ctx) && filterRewriteOptimizationContext.tryGetRanges(ctx, false, context)) { - List weights = filterRewriteOptimizationContext.getWeights(); - class SubLeafCollector implements LeafCollector { - LeafBucketCollector subCollector; - int filterOrd; - - @Override - public void setScorer(Scorable scorer) throws IOException {} - - @Override - public void collect(int docId) throws IOException { - collectBucket(sub, docId, filterOrd); - } + if (collectableSubAggregators == null) { + if (segmentMatchAll(context, ctx) + && filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false, collectableSubAggregators)) { + throw new CollectionTerminatedException(); } - Bits live = ctx.reader().getLiveDocs(); - SubLeafCollector subLeafCollector = new SubLeafCollector(); - for (int i = 0; i < weights.size(); i++) { - Weight weight = weights.get(i); - BulkScorer scorer = weight.bulkScorer(ctx); - if (scorer == null) { - continue; + } else { + if (segmentMatchAll(context, ctx) && filterRewriteOptimizationContext.tryGetRanges(ctx, false, context)) { + List weights = filterRewriteOptimizationContext.getWeights(); + class SubLeafCollector implements LeafCollector { + LeafBucketCollector subCollector; + int filterOrd; + + @Override + public void setScorer(Scorable scorer) throws IOException {} + + @Override + public void collect(int docId) throws IOException { + collectBucket(sub, docId, filterOrd); + } + } + Bits live = ctx.reader().getLiveDocs(); + SubLeafCollector subLeafCollector = new SubLeafCollector(); + for (int i = 0; i < weights.size(); i++) { + Weight weight = weights.get(i); + BulkScorer scorer = weight.bulkScorer(ctx); + if (scorer == null) { + continue; + } + subLeafCollector.subCollector = collectableSubAggregators.getLeafCollector(ctx); + subLeafCollector.filterOrd = i; + scorer.score(subLeafCollector, live, 0, NO_MORE_DOCS); } - subLeafCollector.subCollector = collectableSubAggregators.getLeafCollector(ctx); - subLeafCollector.filterOrd = i; - scorer.score(subLeafCollector, live, 0, NO_MORE_DOCS); - } - throw new CollectionTerminatedException(); + throw new CollectionTerminatedException(); + } } final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);