diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java index 33fa0cf0e4ca8..9ca27a3459ecf 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java @@ -173,6 +173,7 @@ public final class CompositeAggregator extends BucketsAggregator { @Override public boolean canOptimize() { + if (parent != null || subAggregators.length != 0) return false; if (canOptimize(sourceConfigs)) { this.valuesSource = (RoundingValuesSource) sourceConfigs[0].valuesSource(); if (rawAfterKey != null) { diff --git a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/CompositeAggregatorBridge.java b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/CompositeAggregatorBridge.java index 1982793332605..9cd6c35ad9541 100644 --- a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/CompositeAggregatorBridge.java +++ b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/CompositeAggregatorBridge.java @@ -8,11 +8,20 @@ package org.opensearch.search.optimization.filterrewrite; +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.index.PointValues; import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.search.aggregations.LeafBucketCollector; import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceConfig; import org.opensearch.search.aggregations.bucket.composite.RoundingValuesSource; +import java.io.IOException; +import java.util.List; +import java.util.function.BiConsumer; + +import static org.opensearch.search.optimization.filterrewrite.TreeTraversal.multiRangesTraverse; + /** * For composite aggregation to do optimization when it only has a single date histogram source */ @@ -33,4 +42,17 @@ private boolean canOptimize(boolean missing, boolean hasScript, MappedFieldType } return false; } + + @Override + public final void tryOptimize(PointValues values, BiConsumer incrementDocCount, final LeafBucketCollector sub) throws IOException { + DateFieldMapper.DateFieldType fieldType = getFieldType(); + BiConsumer> collectRangeIDs = (activeIndex, docIDs) -> { + long rangeStart = LongPoint.decodeDimension(optimizationContext.getRanges().lowers[activeIndex], 0); + rangeStart = fieldType.convertNanosToMillis(rangeStart); + long ord = getBucketOrd(bucketOrdProducer().apply(rangeStart)); + incrementDocCount.accept(ord, (long) docIDs.size()); + }; + + optimizationContext.consumeDebugInfo(multiRangesTraverse(values.getPointTree(), optimizationContext.getRanges(), collectRangeIDs, getSize())); + } } diff --git a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java index 705e16b93264c..bff63137f8575 100644 --- a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java +++ b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java @@ -115,7 +115,7 @@ protected long[] processHardBounds(long[] bounds, LongBounds hardBounds) { return bounds; } - private DateFieldMapper.DateFieldType getFieldType() { + protected DateFieldMapper.DateFieldType getFieldType() { assert fieldType instanceof DateFieldMapper.DateFieldType; return (DateFieldMapper.DateFieldType) fieldType; } @@ -125,7 +125,7 @@ protected int getSize() { } @Override - public final void tryOptimize(PointValues values, BiConsumer incrementDocCount, final LeafBucketCollector sub) throws IOException { + public void tryOptimize(PointValues values, BiConsumer incrementDocCount, final LeafBucketCollector sub) throws IOException { int size = getSize(); DateFieldMapper.DateFieldType fieldType = getFieldType(); @@ -147,7 +147,7 @@ public final void tryOptimize(PointValues values, BiConsumer increme optimizationContext.consumeDebugInfo(multiRangesTraverse(values.getPointTree(), optimizationContext.getRanges(), collectRangeIDs, size)); } - private static long getBucketOrd(long bucketOrd) { + protected static long getBucketOrd(long bucketOrd) { if (bucketOrd < 0) { // already seen bucketOrd = -1 - bucketOrd; }