Skip to content

Commit

Permalink
branch if no subaggregation
Browse files Browse the repository at this point in the history 
Signed-off-by: bowenlan-amzn <[email protected]>
  • Loading branch information
@bowenlan-amzn
bowenlan-amzn committed Feb 5, 2025
1 parent b1154d2 commit 4e15957
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 90 deletions.
31 changes: 15 additions & 16 deletions ...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down
72 changes: 23 additions & 49 deletions ...main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Long, DocIdSetBuilder> ordinalToBuilder = collector.bucketOrdinalToDocIdSetBuilder;
logger.debug("keys of bucketOrdinalToDocIdSetBuilder: {}", ordinalToBuilder.keySet());
int maxOrdinal = ordinalToBuilder.keySet().stream().mapToInt(Long::intValue).max().orElse(0) + 1;
// 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];
// for (Map.Entry<Long, DocIdSetBuilder> entry : ordinalToBuilder.entrySet()) {
// int ordinal = Math.toIntExact(entry.getKey());
Expand All @@ -110,12 +94,12 @@ static FilterRewriteOptimizationContext.DebugInfo multiRangesTraverse(
// }
// debugInfo.iterators = iterators;

DocIdSetBuilder[] builder = new DocIdSetBuilder[maxOrdinal];
for (Map.Entry<Long, DocIdSetBuilder> entry : ordinalToBuilder.entrySet()) {
int ordinal = Math.toIntExact(entry.getKey());
builder[ordinal] = entry.getValue();
}
debugInfo.builders = builder;
// DocIdSetBuilder[] builder = new DocIdSetBuilder[maxOrdinal];
// for (Map.Entry<Long, DocIdSetBuilder> entry : ordinalToBuilder.entrySet()) {
// int ordinal = Math.toIntExact(entry.getKey());
// builder[ordinal] = entry.getValue();
// }
// debugInfo.builders = builder;

return debugInfo;
}
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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<IOException> collect) throws IOException {
Expand Down Expand Up @@ -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);
}
Expand All @@ -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]);
// }
}

/**
Expand Down
57 changes: 32 additions & 25 deletions server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Weight> 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<Weight> 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);
Expand Down

0 comments on commit 4e15957

Please sign in to comment.