Fail fast filter optimization in any case where segment is not match all

Signed-off-by: Finn Carroll <carrofin@amazon.com>
opensearch-project · Aug 9, 2024 · ae267fa · ae267fa
1 parent 252742b
commit ae267fa
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 13 deletions.
diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/40_range.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/40_range.yml
@@ -673,3 +673,49 @@ setup:
   - match: { aggregations.my_range.buckets.3.from: 1.5 }
   - is_false:  aggregations.my_range.buckets.3.to
   - match: { aggregations.my_range.buckets.3.doc_count: 2 }
+
+---
+"Range query and aggregation test":
+  - do:
+      bulk:
+        refresh: true
+        index: range-agg-w-query
+        body:
+          - '{"index": {}}'
+          - '{"routing": "route1", "v": -10}'
+          - '{"index": {}}'
+          - '{"routing": "route1", "v": -5}'
+          - '{"index": {}}'
+          - '{"routing": "route1", "v": 10}'
+          - '{"index": {}}'
+          - '{"routing": "route2", "v": 15}'
+          - '{"index": {}}'
+          - '{"routing": "route2", "v": 20}'
+
+  - do:
+      search:
+        index: range-agg-w-query
+        body:
+          query:
+            bool:
+              must:
+                match_all: {}
+              filter:
+                - terms:
+                    routing:
+                      - "route1"
+          aggregations:
+            aggregationName:
+              range:
+                field: v
+                keyed: true
+                ranges:
+                  - to: 0
+                    key: "0"
+                  - from: 0
+                    key: "1"
+          _source: false
+
+  - match: { hits.total.value: 3 }
+  - match: { aggregations.aggregationName.buckets.0.doc_count: 2 }
+  - match: { aggregations.aggregationName.buckets.1.doc_count: 1 }
diff --git a/...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java b/...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java
@@ -99,12 +99,10 @@ void setRanges(Ranges ranges) {
     public boolean tryOptimize(final LeafReaderContext leafCtx, final BiConsumer<Long, Long> incrementDocCount, boolean segmentMatchAll)
         throws IOException {
         segments.incrementAndGet();
-        if (!canOptimize) {
-            return false;
-        }
 
+        if (!canOptimize) return false;
+        if (!segmentMatchAll) return false;
         if (leafCtx.reader().hasDeletions()) return false;
-
         PointValues values = leafCtx.reader().getPointValues(aggregatorBridge.fieldType.name());
         if (values == null) return false;
         // only proceed if every document corresponds to exactly one point
@@ -120,7 +118,7 @@ public boolean tryOptimize(final LeafReaderContext leafCtx, final BiConsumer<Lon
             return false;
         }
 
-        Ranges ranges = getRanges(leafCtx, segmentMatchAll);
+        Ranges ranges = getRanges(leafCtx);
         if (ranges == null) return false;
 
         consumeDebugInfo(aggregatorBridge.tryOptimize(values, incrementDocCount, ranges));
@@ -132,10 +130,10 @@ public boolean tryOptimize(final LeafReaderContext leafCtx, final BiConsumer<Lon
         return true;
     }
 
-    Ranges getRanges(LeafReaderContext leafCtx, boolean segmentMatchAll) {
+    Ranges getRanges(LeafReaderContext leafCtx) {
         if (!preparedAtShardLevel) {
             try {
-                return getRangesFromSegment(leafCtx, segmentMatchAll);
+                return getRangesFromSegment(leafCtx);
             } catch (IOException e) {
                 logger.warn("Failed to build ranges from segment.", e);
                 return null;
@@ -148,11 +146,7 @@ Ranges getRanges(LeafReaderContext leafCtx, boolean segmentMatchAll) {
      * Even when ranges cannot be built at shard level, we can still build ranges
      * at segment level when it's functionally match-all at segment level
      */
-    private Ranges getRangesFromSegment(LeafReaderContext leafCtx, boolean segmentMatchAll) throws IOException {
-        if (!segmentMatchAll) {
-            return null;
-        }
-
+    private Ranges getRangesFromSegment(LeafReaderContext leafCtx) throws IOException {
         logger.debug("Shard {} segment {} functionally match all documents. Build the fast filter", shardId, leafCtx.ord);
         return aggregatorBridge.tryBuildRangesFromSegment(leafCtx);
     }

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
@@ -70,6 +70,7 @@
 import java.util.function.Function;
 
 import static org.opensearch.core.xcontent.ConstructingObjectParser.optionalConstructorArg;
+import static org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll;
 
 /**
  * Aggregate all docs that match given ranges.
@@ -310,7 +311,7 @@ public ScoreMode scoreMode() {
 
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
-        boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false);
+        boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, segmentMatchAll(context, ctx));
         if (optimized) throw new CollectionTerminatedException();
 
         final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);