handle segment level bounds separately

Signed-off-by: bowenlan-amzn <bowenlan23@gmail.com>

Author: bowenlan-amzn

server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java

@@ -31,7 +31,6 @@
 import org.opensearch.index.mapper.DocCountFieldMapper;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.query.DateRangeIncludingNowQuery;
-import org.opensearch.search.aggregations.bucket.composite.CompositeAggregator;
 import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceConfig;
 import org.opensearch.search.aggregations.bucket.composite.RoundingValuesSource;
 import org.opensearch.search.aggregations.bucket.histogram.LongBounds;
@@ -90,11 +89,11 @@ private static Query unwrapIntoConcreteQuery(Query query) {
     }
 
     /**
-     * Finds the global min and max bounds of the field for the shard from each segment
+     * Finds the global min and max bounds of the field for the shard across all segments
      *
      * @return null if the field is empty or not indexed
      */
-    private static long[] getIndexBounds(final SearchContext context, final String fieldName) throws IOException {
+    private static long[] getShardBounds(final SearchContext context, final String fieldName) throws IOException {
         final List<LeafReaderContext> leaves = context.searcher().getIndexReader().leaves();
         long min = Long.MAX_VALUE, max = Long.MIN_VALUE;
         for (LeafReaderContext leaf : leaves) {
@@ -111,6 +110,25 @@ private static long[] getIndexBounds(final SearchContext context, final String f
         return new long[] { min, max };
     }
 
+    /**
+     * Finds the min and max bounds of the field for the segment
+     *
+     * @return null if the field is empty or not indexed
+     */
+    private static long[] getSegmentBounds(final LeafReaderContext context, final String fieldName) throws IOException {
+        long min = Long.MAX_VALUE, max = Long.MIN_VALUE;
+        final PointValues values = context.reader().getPointValues(fieldName);
+        if (values != null) {
+            min = Math.min(min, NumericUtils.sortableBytesToLong(values.getMinPackedValue(), 0));
+            max = Math.max(max, NumericUtils.sortableBytesToLong(values.getMaxPackedValue(), 0));
+        }
+
+        if (min == Long.MAX_VALUE || max == Long.MIN_VALUE) {
+            return null;
+        }
+        return new long[] { min, max };
+    }
+
     /**
      * This method also acts as a pre-condition check for the optimization
      *
@@ -120,32 +138,21 @@ public static long[] getDateHistoAggBounds(final SearchContext context, final St
         final Query cq = unwrapIntoConcreteQuery(context.query());
         if (cq instanceof PointRangeQuery) {
             final PointRangeQuery prq = (PointRangeQuery) cq;
-            final long[] indexBounds = getIndexBounds(context, fieldName);
+            final long[] indexBounds = getShardBounds(context, fieldName);
             if (indexBounds == null) return null;
             return getBoundsWithRangeQuery(prq, fieldName, indexBounds);
         } else if (cq instanceof MatchAllDocsQuery) {
-            return getIndexBounds(context, fieldName);
+            return getShardBounds(context, fieldName);
         } else if (cq instanceof FieldExistsQuery) {
             // when a range query covers all values of a shard, it will be rewrite field exists query
             if (((FieldExistsQuery) cq).getField().equals(fieldName)) {
-                return getIndexBounds(context, fieldName);
+                return getShardBounds(context, fieldName);
             }
         }
 
         return null;
     }
 
-    private static long[] getDateHistoAggBoundsSegLevel(final SearchContext context, final String fieldName) throws IOException {
-        final long[] indexBounds = getIndexBounds(context, fieldName);
-        if (indexBounds == null) return null;
-        final Query cq = unwrapIntoConcreteQuery(context.query());
-        if (cq instanceof PointRangeQuery) {
-            final PointRangeQuery prq = (PointRangeQuery) cq;
-            return getBoundsWithRangeQuery(prq, fieldName, indexBounds);
-        }
-        return indexBounds;
-    }
-
     private static long[] getBoundsWithRangeQuery(PointRangeQuery prq, String fieldName, long[] indexBounds) {
         // Ensure that the query and aggregation are on the same field
         if (prq.getField().equals(fieldName)) {
@@ -158,6 +165,7 @@ private static long[] getBoundsWithRangeQuery(PointRangeQuery prq, String fieldN
             }
             return new long[] { lower, upper };
         }
+
         return null;
     }
 
@@ -230,7 +238,6 @@ public static class FastFilterContext {
         private boolean rewriteable = false;
         private Weight[] filters = null;
         private boolean filtersBuiltAtShardLevel = false;
-        private boolean shouldBuildFiltersAtSegmentLevel = true;
 
         private AggregationType aggregationType;
         private final SearchContext context;
@@ -278,6 +285,10 @@ interface AggregationType {
 
         Weight[] buildFastFilter(SearchContext ctx, CheckedBiFunction<SearchContext, String, long[], IOException> getBounds)
             throws IOException;
+
+        default int getSize() {
+            return Integer.MAX_VALUE;
+        }
     }
 
     /**
@@ -314,8 +325,23 @@ public boolean isRewriteable(Object parent, int subAggLength) {
         public Weight[] buildFastFilter(SearchContext context, CheckedBiFunction<SearchContext, String, long[], IOException> getBounds)
             throws IOException {
             long[] bounds = getBounds.apply(context, fieldType.name());
-            bounds = processHardBounds(bounds);
             logger.debug("Bounds are {} for shard {}", bounds, context.indexShard().shardId());
+            return buildFastFilter(context, bounds);
+        }
+
+        private Weight[] buildFastFilterWithSegBounds(
+            SearchContext context,
+            CheckedBiFunction<LeafReaderContext, String, long[], IOException> getBounds,
+            LeafReaderContext leaf
+        ) throws IOException {
+            long[] bounds = getBounds.apply(leaf, fieldType.name());
+            logger.debug("Bounds are {} for shard {} segment {}", bounds, context.indexShard().shardId(), leaf.ord);
+            return buildFastFilter(context, bounds);
+        }
+
+        private Weight[] buildFastFilter(SearchContext context, long[] bounds) throws IOException {
+            bounds = processHardBounds(bounds);
+            logger.debug("Bounds are {} for shard {} with hard bound", bounds, context.indexShard().shardId());
             if (bounds == null) {
                 return null;
             }
@@ -394,7 +420,7 @@ public static boolean tryFastFilterAggregation(
         final BiConsumer<Long, Integer> incrementDocCount
     ) throws IOException {
         if (fastFilterContext == null) return false;
-        if (!fastFilterContext.rewriteable || !fastFilterContext.shouldBuildFiltersAtSegmentLevel) {
+        if (!fastFilterContext.rewriteable) {
             return false;
         }
 
@@ -420,11 +446,15 @@ public static boolean tryFastFilterAggregation(
                 fastFilterContext.context.indexShard().shardId(),
                 ctx.ord
             );
-            filters = fastFilterContext.buildFastFilter(FastFilterRewriteHelper::getDateHistoAggBoundsSegLevel);
+            if (fastFilterContext.aggregationType instanceof AbstractDateHistogramAggregationType) {
+                filters = ((AbstractDateHistogramAggregationType) fastFilterContext.aggregationType).buildFastFilterWithSegBounds(
+                    fastFilterContext.context,
+                    FastFilterRewriteHelper::getSegmentBounds,
+                    ctx
+                );
+            }
             if (filters == null) {
-                // At segment level, build filter should only be called once
-                // since the conditions for build filter won't change for other segments
-                fastFilterContext.shouldBuildFiltersAtSegmentLevel = false;
+
                 return false;
             }
         }
@@ -441,7 +471,7 @@ public static boolean tryFastFilterAggregation(
         }
 
         int s = 0;
-        int size = Integer.MAX_VALUE;
+        int size = fastFilterContext.aggregationType.getSize();
         for (i = 0; i < filters.length; i++) {
             if (counts[i] > 0) {
                 long bucketKey = i; // the index of filters is the key for filters aggregation
@@ -451,9 +481,6 @@ public static boolean tryFastFilterAggregation(
                     bucketKey = fieldType.convertNanosToMillis(
                         NumericUtils.sortableBytesToLong(((PointRangeQuery) filters[i].getQuery()).getLowerPoint(), 0)
                     );
-                    if (fastFilterContext.aggregationType instanceof CompositeAggregator.CompositeAggregationType) {
-                        size = ((CompositeAggregator.CompositeAggregationType) fastFilterContext.aggregationType).getSize();
-                    }
                 }
                 incrementDocCount.accept(bucketKey, counts[i]);
                 s++;

server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -97,7 +97,7 @@
  *
  * @opensearch.internal
  */
-public final class CompositeAggregator extends BucketsAggregator {
+final class CompositeAggregator extends BucketsAggregator {
     private final int size;
     private final List<String> sourceNames;
     private final int[] reverseMuls;
@@ -178,7 +178,7 @@ public final class CompositeAggregator extends BucketsAggregator {
     /**
      * Currently the filter rewrite is only supported for date histograms
      */
-    public class CompositeAggregationType extends FastFilterRewriteHelper.AbstractDateHistogramAggregationType {
+    private class CompositeAggregationType extends FastFilterRewriteHelper.AbstractDateHistogramAggregationType {
         private final RoundingValuesSource valuesSource;
         private long afterKey = -1L;
 
@@ -210,6 +210,7 @@ protected void processAfterKey(long[] bound, long interval) {
             }
         }
 
+        @Override
         public int getSize() {
             return size;
         }

server/src/test/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregatorTests.java

@@ -1254,31 +1254,37 @@ public void testHardBoundsNotOverlapping() throws IOException {
         );
     }
 
-    public void testRangeQuery() throws IOException {
+    public void testFilterRewriteOptimizationWithRangeQuery() throws IOException {
         testSearchCase(
-            LongPoint.newRangeQuery(SEARCHABLE_DATE, asLong("2018-01-01"), asLong("2020-01-01")),
+            LongPoint.newRangeQuery(AGGREGABLE_DATE, asLong("2018-01-01"), asLong("2020-01-01")),
             Arrays.asList("2017-02-01", "2017-02-02", "2017-02-02", "2017-02-03", "2017-02-03", "2017-02-03", "2017-02-05"),
             aggregation -> aggregation.calendarInterval(DateHistogramInterval.DAY).field(AGGREGABLE_DATE),
             histogram -> {
                 List<? extends Histogram.Bucket> buckets = histogram.getBuckets();
                 assertEquals(0, buckets.size());
             },
-            false
+            10000,
+            false,
+            false,
+            true // force AGGREGABLE_DATE field to be searchable to test the filter rewrite optimization path
         );
 
         testSearchCase(
-            LongPoint.newRangeQuery(SEARCHABLE_DATE, asLong("2016-01-01"), asLong("2017-01-01")),
+            LongPoint.newRangeQuery(AGGREGABLE_DATE, asLong("2016-01-01"), asLong("2017-01-01")),
             Arrays.asList("2017-02-01", "2017-02-02", "2017-02-02", "2017-02-03", "2017-02-03", "2017-02-03", "2017-02-05"),
             aggregation -> aggregation.calendarInterval(DateHistogramInterval.DAY).field(AGGREGABLE_DATE),
             histogram -> {
                 List<? extends Histogram.Bucket> buckets = histogram.getBuckets();
                 assertEquals(0, buckets.size());
             },
-            false
+            10000,
+            false,
+            false,
+            true
         );
 
         testSearchCase(
-            LongPoint.newRangeQuery(SEARCHABLE_DATE, asLong("2016-01-01"), asLong("2017-02-02")),
+            LongPoint.newRangeQuery(AGGREGABLE_DATE, asLong("2016-01-01"), asLong("2017-02-02")),
             Arrays.asList("2017-02-01", "2017-02-02", "2017-02-02", "2017-02-03", "2017-02-03", "2017-02-03", "2017-02-05"),
             aggregation -> aggregation.calendarInterval(DateHistogramInterval.DAY).field(AGGREGABLE_DATE),
             histogram -> {
@@ -1293,11 +1299,14 @@ public void testRangeQuery() throws IOException {
                 assertEquals("2017-02-02T00:00:00.000Z", bucket.getKeyAsString());
                 assertEquals(2, bucket.getDocCount());
             },
-            false
+            10000,
+            false,
+            false,
+            true
         );
 
         testSearchCase(
-            LongPoint.newRangeQuery(SEARCHABLE_DATE, asLong("2017-02-03"), asLong("2020-01-01")),
+            LongPoint.newRangeQuery(AGGREGABLE_DATE, asLong("2017-02-03"), asLong("2020-01-01")),
             Arrays.asList("2017-02-01", "2017-02-02", "2017-02-02", "2017-02-03", "2017-02-03", "2017-02-03", "2017-02-05"),
             aggregation -> aggregation.calendarInterval(DateHistogramInterval.DAY).field(AGGREGABLE_DATE),
             histogram -> {
@@ -1316,7 +1325,10 @@ public void testRangeQuery() throws IOException {
                 assertEquals("2017-02-05T00:00:00.000Z", bucket.getKeyAsString());
                 assertEquals(1, bucket.getDocCount());
             },
-            false
+            10000,
+            false,
+            false,
+            true
         );
     }
 
@@ -1388,7 +1400,19 @@ private void testSearchCase(
         boolean useNanosecondResolution,
         boolean useDocCountField
     ) throws IOException {
-        boolean aggregableDateIsSearchable = randomBoolean();
+        testSearchCase(query, dataset, configure, verify, maxBucket, useNanosecondResolution, useDocCountField, randomBoolean());
+    }
+
+    private void testSearchCase(
+        Query query,
+        List<String> dataset,
+        Consumer<DateHistogramAggregationBuilder> configure,
+        Consumer<InternalDateHistogram> verify,
+        int maxBucket,
+        boolean useNanosecondResolution,
+        boolean useDocCountField,
+        boolean aggregableDateIsSearchable
+    ) throws IOException {
         logger.debug("Aggregable date is searchable {}", aggregableDateIsSearchable);
         DateFieldMapper.DateFieldType fieldType = aggregableDateFieldType(useNanosecondResolution, aggregableDateIsSearchable);