Add segment match all guard for range aggregator

Signed-off-by: Finn Carroll <[email protected]>
opensearch-project · Aug 16, 2024 · 490d442 · 490d442
1 parent 374099e
commit 490d442
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java
@@ -543,7 +543,7 @@ public static boolean isCompositeAggRewriteable(CompositeValuesSourceConfig[] so
         return sourceConfigs.length == 1 && sourceConfigs[0].valuesSource() instanceof RoundingValuesSource;
     }
 
-    private static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leafCtx) throws IOException {
+    public static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leafCtx) throws IOException {
         Weight weight = ctx.searcher().createWeight(ctx.query(), ScoreMode.COMPLETE_NO_SCORES, 1f);
         return weight != null && weight.count(leafCtx) == leafCtx.reader().numDocs();
     }

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
@@ -68,6 +68,7 @@
 import java.util.function.BiConsumer;
 
 import static org.opensearch.core.xcontent.ConstructingObjectParser.optionalConstructorArg;
+import static org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper.segmentMatchAll;
 
 /**
  * Aggregate all docs that match given ranges.
@@ -298,12 +299,14 @@ public ScoreMode scoreMode() {
 
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
-        boolean optimized = fastFilterContext.tryFastFilterAggregation(
-            ctx,
-            this::incrementBucketDocCount,
-            (activeIndex) -> subBucketOrdinal(0, (int) activeIndex)
-        );
-        if (optimized) throw new CollectionTerminatedException();
+        if (segmentMatchAll(context, ctx)
+            && fastFilterContext.tryFastFilterAggregation(
+                ctx,
+                this::incrementBucketDocCount,
+                (activeIndex) -> subBucketOrdinal(0, (int) activeIndex)
+            )) {
+            throw new CollectionTerminatedException();
+        }
 
         final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
         return new LeafBucketCollectorBase(sub, values) {