Add early termination support for min/max aggregations (#33375)

This commit adds the support to early terminate the collection of a leaf
in the min/max aggregator. If the query matches all documents the min and max value
for a numeric field can be retrieved efficiently in the points reader.
This change applies this optimization when possible.

Author: jimczi

server/src/main/java/org/elasticsearch/index/mapper/NumberFieldMapper.java

@@ -187,6 +187,11 @@ public Float parse(Object value, boolean coerce) {
                 return result;
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return HalfFloatPoint.decodeDimension(value, 0);
+            }
+
             @Override
             public Float parse(XContentParser parser, boolean coerce) throws IOException {
                 float parsed = parser.floatValue(coerce);
@@ -279,6 +284,11 @@ public Float parse(Object value, boolean coerce) {
                 return result;
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return FloatPoint.decodeDimension(value, 0);
+            }
+
             @Override
             public Float parse(XContentParser parser, boolean coerce) throws IOException {
                 float parsed = parser.floatValue(coerce);
@@ -360,6 +370,11 @@ public Double parse(Object value, boolean coerce) {
                 return parsed;
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return DoublePoint.decodeDimension(value, 0);
+            }
+
             @Override
             public Double parse(XContentParser parser, boolean coerce) throws IOException {
                 double parsed = parser.doubleValue(coerce);
@@ -452,6 +467,11 @@ public Byte parse(Object value, boolean coerce) {
                 return (byte) doubleValue;
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return INTEGER.parsePoint(value).byteValue();
+            }
+
             @Override
             public Short parse(XContentParser parser, boolean coerce) throws IOException {
                 int value = parser.intValue(coerce);
@@ -508,6 +528,11 @@ public Short parse(Object value, boolean coerce) {
                 return (short) doubleValue;
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return INTEGER.parsePoint(value).shortValue();
+            }
+
             @Override
             public Short parse(XContentParser parser, boolean coerce) throws IOException {
                 return parser.shortValue(coerce);
@@ -560,6 +585,11 @@ public Integer parse(Object value, boolean coerce) {
                 return (int) doubleValue;
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return IntPoint.decodeDimension(value, 0);
+            }
+
             @Override
             public Integer parse(XContentParser parser, boolean coerce) throws IOException {
                 return parser.intValue(coerce);
@@ -674,6 +704,11 @@ public Long parse(Object value, boolean coerce) {
                 return Numbers.toLong(stringValue, coerce);
             }
 
+            @Override
+            public Number parsePoint(byte[] value) {
+                return LongPoint.decodeDimension(value, 0);
+            }
+
             @Override
             public Long parse(XContentParser parser, boolean coerce) throws IOException {
                 return parser.longValue(coerce);
@@ -790,6 +825,7 @@ public abstract Query rangeQuery(String field, Object lowerTerm, Object upperTer
                                   boolean hasDocValues);
         public abstract Number parse(XContentParser parser, boolean coerce) throws IOException;
         public abstract Number parse(Object value, boolean coerce);
+        public abstract Number parsePoint(byte[] value);
         public abstract List<Field> createFields(String name, Number value, boolean indexed,
                                                  boolean docValued, boolean stored);
         Number valueForSearch(Number value) {
@@ -938,6 +974,10 @@ public DocValueFormat docValueFormat(String format, DateTimeZone timeZone) {
             }
         }
 
+        public Number parsePoint(byte[] value) {
+            return type.parsePoint(value);
+        }
+
         @Override
         public boolean equals(Object o) {
             if (super.equals(o) == false) {

server/src/main/java/org/elasticsearch/search/aggregations/metrics/max/MaxAggregator.java

@@ -18,7 +18,12 @@
  */
 package org.elasticsearch.search.aggregations.metrics.max;
 
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.PointValues;
+import org.apache.lucene.search.CollectionTerminatedException;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.FutureArrays;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.DoubleArray;
@@ -33,30 +38,45 @@
 import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregator;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
+import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
+
+import static org.elasticsearch.search.aggregations.metrics.min.MinAggregator.getPointReaderOrNull;
 
 public class MaxAggregator extends NumericMetricsAggregator.SingleValue {
 
     final ValuesSource.Numeric valuesSource;
     final DocValueFormat formatter;
 
+    final String pointField;
+    final Function<byte[], Number> pointConverter;
+
     DoubleArray maxes;
 
-    public MaxAggregator(String name, ValuesSource.Numeric valuesSource, DocValueFormat formatter,
-            SearchContext context,
-            Aggregator parent, List<PipelineAggregator> pipelineAggregators,
-            Map<String, Object> metaData) throws IOException {
+    MaxAggregator(String name,
+                    ValuesSourceConfig<ValuesSource.Numeric> config,
+                    ValuesSource.Numeric valuesSource,
+                    SearchContext context,
+                    Aggregator parent, List<PipelineAggregator> pipelineAggregators,
+                    Map<String, Object> metaData) throws IOException {
         super(name, context, parent, pipelineAggregators, metaData);
         this.valuesSource = valuesSource;
-        this.formatter = formatter;
         if (valuesSource != null) {
             maxes = context.bigArrays().newDoubleArray(1, false);
             maxes.fill(0, maxes.size(), Double.NEGATIVE_INFINITY);
         }
+        this.formatter = config.format();
+        this.pointConverter = getPointReaderOrNull(context, parent, config);
+        if (pointConverter != null) {
+            pointField = config.fieldContext().field();
+        } else {
+            pointField = null;
+        }
     }
 
     @Override
@@ -68,8 +88,28 @@ public boolean needsScores() {
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
             final LeafBucketCollector sub) throws IOException {
         if (valuesSource == null) {
-            return LeafBucketCollector.NO_OP_COLLECTOR;
-    }
+            if (parent != null) {
+                return LeafBucketCollector.NO_OP_COLLECTOR;
+            } else {
+                // we have no parent and the values source is empty so we can skip collecting hits.
+                throw new CollectionTerminatedException();
+            }
+        }
+        if (pointConverter != null) {
+            Number segMax = findLeafMaxValue(ctx.reader(), pointField, pointConverter);
+            if (segMax != null) {
+                /**
+                 * There is no parent aggregator (see {@link MinAggregator#getPointReaderOrNull}
+                 * so the ordinal for the bucket is always 0.
+                 */
+                assert maxes.size() == 1;
+                double max = maxes.get(0);
+                max = Math.max(max, segMax.doubleValue());
+                maxes.set(0, max);
+                // the maximum value has been extracted, we don't need to collect hits on this segment.
+                throw new CollectionTerminatedException();
+            }
+        }
         final BigArrays bigArrays = context.bigArrays();
         final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
         final NumericDoubleValues values = MultiValueMode.MAX.select(allValues);
@@ -118,4 +158,48 @@ public InternalAggregation buildEmptyAggregation() {
     public void doClose() {
         Releasables.close(maxes);
     }
+
+    /**
+     * Returns the maximum value indexed in the <code>fieldName</code> field or <code>null</code>
+     * if the value cannot be inferred from the indexed {@link PointValues}.
+     */
+    public static Number findLeafMaxValue(LeafReader reader, String fieldName, Function<byte[], Number> converter) throws IOException {
+        final PointValues pointValues = reader.getPointValues(fieldName);
+        if (pointValues == null) {
+            return null;
+        }
+        final Bits liveDocs = reader.getLiveDocs();
+        if (liveDocs == null) {
+            return converter.apply(pointValues.getMaxPackedValue());
+        }
+        int numBytes = pointValues.getBytesPerDimension();
+        final byte[] maxValue = pointValues.getMaxPackedValue();
+        final Number[] result = new Number[1];
+        pointValues.intersect(new PointValues.IntersectVisitor() {
+            @Override
+            public void visit(int docID) {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public void visit(int docID, byte[] packedValue) {
+                if (liveDocs.get(docID)) {
+                    // we need to collect all values in this leaf (the sort is ascending) where
+                    // the last live doc is guaranteed to contain the max value for the segment.
+                    result[0] = converter.apply(packedValue);
+                }
+            }
+
+            @Override
+            public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
+                if (FutureArrays.equals(maxValue, 0, numBytes, maxPackedValue, 0, numBytes)) {
+                    // we only check leaves that contain the max value for the segment.
+                    return PointValues.Relation.CELL_CROSSES_QUERY;
+                } else {
+                    return PointValues.Relation.CELL_OUTSIDE_QUERY;
+                }
+            }
+        });
+        return result[0];
+    }
 }

server/src/main/java/org/elasticsearch/search/aggregations/metrics/max/MaxAggregatorFactory.java

@@ -43,13 +43,13 @@ public MaxAggregatorFactory(String name, ValuesSourceConfig<Numeric> config, Sea
     @Override
     protected Aggregator createUnmapped(Aggregator parent,
             List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
-        return new MaxAggregator(name, null, config.format(), context, parent, pipelineAggregators, metaData);
+        return new MaxAggregator(name, config, null, context, parent, pipelineAggregators, metaData);
     }
 
     @Override
     protected Aggregator doCreateInternal(ValuesSource.Numeric valuesSource, Aggregator parent,
             boolean collectsFromSingleBucket, List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData)
                     throws IOException {
-        return new MaxAggregator(name, valuesSource, config.format(), context, parent, pipelineAggregators, metaData);
+        return new MaxAggregator(name, config, valuesSource, context, parent, pipelineAggregators, metaData);
     }
 }