Save memory when histogram agg is not on top (backport of #57277) (#57377)

This saves some memory when the `histogram` aggregation is not a top
level aggregation by dropping `asMultiBucketAggregator` in favor of
natively implementing multi-bucket storage in the aggregator. For the
most part this just uses the `LongKeyedBucketOrds` that we built the
first time we did this.

Author: nik9000

rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/10_histogram.yml

@@ -492,7 +492,65 @@ setup:
   - match: { aggregations.histo.buckets.0.doc_count: 1 }
 
 ---
-"profiler":
+"histogram profiler":
+  - skip:
+      version: " - 7.8.99"
+      reason:  debug info added in 7.9.0
+
+  - do:
+      indices.create:
+          index: test_2
+          body:
+            settings:
+              number_of_replicas: 0
+              number_of_shards: 1
+            mappings:
+              properties:
+                n:
+                  type: long
+
+  - do:
+      bulk:
+        index: test_2
+        refresh: true
+        body:
+            - '{"index": {}}'
+            - '{"n": "1"}'
+            - '{"index": {}}'
+            - '{"n": "2"}'
+            - '{"index": {}}'
+            - '{"n": "10"}'
+            - '{"index": {}}'
+            - '{"n": "17"}'
+
+  - do:
+      search:
+        index: test_2
+        body:
+          size: 0
+          profile: true
+          aggs:
+            histo:
+              histogram:
+                field: n
+                interval: 5
+  - match: { hits.total.value: 4 }
+  - length: { aggregations.histo.buckets: 4 }
+  - match: { aggregations.histo.buckets.0.key: 0 }
+  - match: { aggregations.histo.buckets.0.doc_count: 2 }
+  - match: { aggregations.histo.buckets.1.key: 5 }
+  - match: { aggregations.histo.buckets.1.doc_count: 0 }
+  - match: { aggregations.histo.buckets.2.key: 10 }
+  - match: { aggregations.histo.buckets.2.doc_count: 1 }
+  - match: { aggregations.histo.buckets.3.key: 15 }
+  - match: { aggregations.histo.buckets.3.doc_count: 1 }
+  - match: { profile.shards.0.aggregations.0.type: NumericHistogramAggregator }
+  - match: { profile.shards.0.aggregations.0.description: histo }
+  - match: { profile.shards.0.aggregations.0.breakdown.collect_count: 4 }
+  - match: { profile.shards.0.aggregations.0.debug.total_buckets: 3 }
+
+---
+"date_histogram profiler":
   - skip:
       version: " - 7.8.99"
       reason:  debug info added in 7.9.0

server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/AbstractHistogramAggregator.java

@@ -0,0 +1,124 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.aggregations.bucket.histogram;
+
+import org.apache.lucene.util.CollectionUtil;
+import org.elasticsearch.common.lease.Releasables;
+import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.Aggregator;
+import org.elasticsearch.search.aggregations.AggregatorFactories;
+import org.elasticsearch.search.aggregations.BucketOrder;
+import org.elasticsearch.search.aggregations.InternalAggregation;
+import org.elasticsearch.search.aggregations.bucket.BucketsAggregator;
+import org.elasticsearch.search.aggregations.bucket.histogram.InternalHistogram.EmptyBucketInfo;
+import org.elasticsearch.search.aggregations.bucket.terms.LongKeyedBucketOrds;
+import org.elasticsearch.search.internal.SearchContext;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.function.BiConsumer;
+
+/**
+ * Base class for functionality shared between aggregators for this
+ * {@code histogram} aggregation.
+ */
+public abstract class AbstractHistogramAggregator extends BucketsAggregator {
+    protected final DocValueFormat formatter;
+    protected final double interval;
+    protected final double offset;
+    protected final BucketOrder order;
+    protected final boolean keyed;
+    protected final long minDocCount;
+    protected final double minBound;
+    protected final double maxBound;
+    protected final LongKeyedBucketOrds bucketOrds;
+
+    public AbstractHistogramAggregator(
+        String name,
+        AggregatorFactories factories,
+        double interval,
+        double offset,
+        BucketOrder order,
+        boolean keyed,
+        long minDocCount,
+        double minBound,
+        double maxBound,
+        DocValueFormat formatter,
+        SearchContext context,
+        Aggregator parent,
+        boolean collectsFromSingleBucket,
+        Map<String, Object> metadata
+    ) throws IOException {
+        super(name, factories, context, parent, metadata);
+        if (interval <= 0) {
+            throw new IllegalArgumentException("interval must be positive, got: " + interval);
+        }
+        this.interval = interval;
+        this.offset = offset;
+        this.order = order;
+        order.validate(this);
+        this.keyed = keyed;
+        this.minDocCount = minDocCount;
+        this.minBound = minBound;
+        this.maxBound = maxBound;
+        this.formatter = formatter;
+        bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), collectsFromSingleBucket);
+    }
+
+    @Override
+    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        return buildAggregationsForVariableBuckets(owningBucketOrds, bucketOrds,
+            (bucketValue, docCount, subAggregationResults) -> {
+                double roundKey = Double.longBitsToDouble(bucketValue);
+                double key = roundKey * interval + offset;
+                return new InternalHistogram.Bucket(key, docCount, keyed, formatter, subAggregationResults);
+            }, buckets -> {
+                // the contract of the histogram aggregation is that shards must return buckets ordered by key in ascending order
+                CollectionUtil.introSort(buckets, BucketOrder.key(true).comparator());
+
+                EmptyBucketInfo emptyBucketInfo = null;
+                if (minDocCount == 0) {
+                    emptyBucketInfo = new EmptyBucketInfo(interval, offset, minBound, maxBound, buildEmptySubAggregations());
+                }
+                return new InternalHistogram(name, buckets, order, minDocCount, emptyBucketInfo, formatter, keyed, metadata());
+            });
+    }
+
+    @Override
+    public InternalAggregation buildEmptyAggregation() {
+        InternalHistogram.EmptyBucketInfo emptyBucketInfo = null;
+        if (minDocCount == 0) {
+            emptyBucketInfo = new InternalHistogram.EmptyBucketInfo(interval, offset, minBound, maxBound, buildEmptySubAggregations());
+        }
+        return new InternalHistogram(name, Collections.emptyList(), order, minDocCount, emptyBucketInfo, formatter, keyed, metadata());
+    }
+
+    @Override
+    public void doClose() {
+        Releasables.close(bucketOrds);
+    }
+
+    @Override
+    public void collectDebugInfo(BiConsumer<String, Object> add) {
+        add.accept("total_buckets", bucketOrds.size());
+        super.collectDebugInfo(add);
+    }
+}

server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/HistogramAggregatorFactory.java

@@ -19,8 +19,8 @@
 
 package org.elasticsearch.search.aggregations.bucket.histogram;
 
+import org.elasticsearch.common.collect.List;
 import org.elasticsearch.index.query.QueryShardContext;
-import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -36,8 +36,6 @@
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.Collections;
 import java.util.Map;
 
 /**
@@ -54,40 +52,11 @@ public final class HistogramAggregatorFactory extends ValuesSourceAggregatorFact
 
     static void registerAggregators(ValuesSourceRegistry.Builder builder) {
         builder.register(HistogramAggregationBuilder.NAME, CoreValuesSourceType.RANGE,
-            new HistogramAggregatorSupplier() {
-                @Override
-                public Aggregator build(String name, AggregatorFactories factories, double interval, double offset,
-                                        BucketOrder order, boolean keyed, long minDocCount, double minBound, double maxBound,
-                                        ValuesSource valuesSource, DocValueFormat formatter, SearchContext context,
-                                        Aggregator parent,
-                                        Map<String, Object> metadata) throws IOException {
-                    ValuesSource.Range rangeValueSource = (ValuesSource.Range) valuesSource;
-                    if (rangeValueSource.rangeType().isNumeric() == false) {
-                        throw new IllegalArgumentException("Expected numeric range type but found non-numeric range ["
-                            + rangeValueSource.rangeType().name + "]");
-                    }
-                    return new RangeHistogramAggregator(name, factories, interval, offset, order, keyed, minDocCount, minBound,
-                        maxBound, rangeValueSource, formatter, context, parent, metadata);
-                }
-            }
-        );
+            (HistogramAggregatorSupplier) RangeHistogramAggregator::new);
 
         builder.register(HistogramAggregationBuilder.NAME,
-            Collections.unmodifiableList(Arrays.asList(CoreValuesSourceType.NUMERIC,
-                                                       CoreValuesSourceType.DATE,
-                                                       CoreValuesSourceType.BOOLEAN)),
-            new HistogramAggregatorSupplier() {
-                @Override
-                public Aggregator build(String name, AggregatorFactories factories, double interval, double offset,
-                                        BucketOrder order, boolean keyed, long minDocCount, double minBound, double maxBound,
-                                        ValuesSource valuesSource, DocValueFormat formatter, SearchContext context,
-                                        Aggregator parent,
-                                        Map<String, Object> metadata) throws IOException {
-                    return new NumericHistogramAggregator(name, factories, interval, offset, order, keyed, minDocCount, minBound,
-                        maxBound, (ValuesSource.Numeric) valuesSource, formatter, context, parent, metadata);
-                }
-            }
-        );
+            List.of(CoreValuesSourceType.NUMERIC, CoreValuesSourceType.DATE, CoreValuesSourceType.BOOLEAN),
+            (HistogramAggregatorSupplier) NumericHistogramAggregator::new);
     }
 
     public HistogramAggregatorFactory(String name,
@@ -123,10 +92,6 @@ protected Aggregator doCreateInternal(ValuesSource valuesSource,
                                             Aggregator parent,
                                             boolean collectsFromSingleBucket,
                                             Map<String, Object> metadata) throws IOException {
-        if (collectsFromSingleBucket == false) {
-            return asMultiBucketAggregator(this, searchContext, parent);
-        }
-
         AggregatorSupplier aggregatorSupplier = queryShardContext.getValuesSourceRegistry().getAggregator(config.valueSourceType(),
             HistogramAggregationBuilder.NAME);
         if (aggregatorSupplier instanceof HistogramAggregatorSupplier == false) {
@@ -135,14 +100,14 @@ protected Aggregator doCreateInternal(ValuesSource valuesSource,
         }
         HistogramAggregatorSupplier histogramAggregatorSupplier = (HistogramAggregatorSupplier) aggregatorSupplier;
         return histogramAggregatorSupplier.build(name, factories, interval, offset, order, keyed, minDocCount, minBound, maxBound,
-                valuesSource, config.format(), searchContext, parent, metadata);
+                valuesSource, config.format(), searchContext, parent, collectsFromSingleBucket, metadata);
     }
 
     @Override
     protected Aggregator createUnmapped(SearchContext searchContext,
                                             Aggregator parent,
                                             Map<String, Object> metadata) throws IOException {
         return new NumericHistogramAggregator(name, factories, interval, offset, order, keyed, minDocCount, minBound, maxBound,
-            null, config.format(), searchContext, parent, metadata);
+            null, config.format(), searchContext, parent, false, metadata);
     }
 }