Optimize InternalAggregations construction a little (elastic#120868) (elastic#126516)

We can streamline and optimize this logic a little to
see less copying and more compact results.

Author: original-brownbear

benchmarks/src/main/java/org/elasticsearch/benchmark/search/aggregations/TermsReduceBenchmark.java

@@ -111,7 +111,7 @@ public void setup() {
                 dict[i] = new BytesRef(Long.toString(rand.nextLong()));
             }
             for (int i = 0; i < numShards; i++) {
-                aggsList.add(InternalAggregations.from(Collections.singletonList(newTerms(rand, dict, true))));
+                aggsList.add(InternalAggregations.from(newTerms(rand, dict, true)));
             }
         }
 
@@ -124,7 +124,7 @@ private StringTerms newTerms(Random rand, BytesRef[] dict, boolean withNested) {
             for (BytesRef term : randomTerms) {
                 InternalAggregations subAggs;
                 if (withNested) {
-                    subAggs = InternalAggregations.from(Collections.singletonList(newTerms(rand, dict, false)));
+                    subAggs = InternalAggregations.from(newTerms(rand, dict, false));
                 } else {
                     subAggs = InternalAggregations.EMPTY;
                 }

benchmarks/src/main/java/org/elasticsearch/benchmark/search/aggregations/bucket/terms/StringTermsSerializationBenchmark.java

@@ -50,13 +50,13 @@ public class StringTermsSerializationBenchmark {
 
     @Setup
     public void initResults() {
-        results = DelayableWriteable.referencing(InternalAggregations.from(List.of(newTerms(true))));
+        results = DelayableWriteable.referencing(InternalAggregations.from(newTerms(true)));
     }
 
     private StringTerms newTerms(boolean withNested) {
         List<StringTerms.Bucket> resultBuckets = new ArrayList<>(buckets);
         for (int i = 0; i < buckets; i++) {
-            InternalAggregations inner = withNested ? InternalAggregations.from(List.of(newTerms(false))) : InternalAggregations.EMPTY;
+            InternalAggregations inner = withNested ? InternalAggregations.from(newTerms(false)) : InternalAggregations.EMPTY;
             resultBuckets.add(new StringTerms.Bucket(new BytesRef("test" + i), i, inner, false, 0, DocValueFormat.RAW));
         }
         return new StringTerms(

modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/histogram/InternalAutoDateHistogram.java

@@ -414,7 +414,7 @@ private BucketReduceResult addEmptyBuckets(BucketReduceResult current, Aggregati
 
         Bucket lastBucket = null;
         ListIterator<Bucket> iter = list.listIterator();
-        InternalAggregations reducedEmptySubAggs = InternalAggregations.reduce(List.of(bucketInfo.emptySubAggregations), reduceContext);
+        InternalAggregations reducedEmptySubAggs = InternalAggregations.reduce(bucketInfo.emptySubAggregations, reduceContext);
 
         // Add the empty buckets within the data,
         // e.g. if the data series is [1,2,3,7] there're 3 empty buckets that will be created for 4,5,6

modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/timeseries/InternalTimeSeries.java

@@ -222,7 +222,7 @@ public InternalAggregation get() {
                     InternalBucket reducedBucket;
                     if (bucketsWithSameKey.size() == 1) {
                         reducedBucket = bucketsWithSameKey.get(0);
-                        reducedBucket.aggregations = InternalAggregations.reduce(List.of(reducedBucket.aggregations), reduceContext);
+                        reducedBucket.aggregations = InternalAggregations.reduce(reducedBucket.aggregations, reduceContext);
                     } else {
                         reducedBucket = reduceBucket(bucketsWithSameKey, reduceContext);
                     }

modules/aggregations/src/main/java/org/elasticsearch/aggregations/pipeline/DerivativePipelineAggregator.java

@@ -22,8 +22,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 
 import static org.elasticsearch.search.aggregations.pipeline.BucketHelpers.resolveBucketValue;
 
@@ -69,11 +67,16 @@ public InternalAggregation reduce(InternalAggregation aggregation, AggregationRe
                 if (xAxisUnits != null) {
                     xDiff = (thisBucketKey.doubleValue() - lastBucketKey.doubleValue()) / xAxisUnits;
                 }
-                final List<InternalAggregation> aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false)
-                    .collect(Collectors.toCollection(ArrayList::new));
-                aggs.add(new Derivative(name(), gradient, xDiff, formatter, metadata()));
-                Bucket newBucket = factory.createBucket(factory.getKey(bucket), bucket.getDocCount(), InternalAggregations.from(aggs));
-                newBuckets.add(newBucket);
+                newBuckets.add(
+                    factory.createBucket(
+                        factory.getKey(bucket),
+                        bucket.getDocCount(),
+                        InternalAggregations.append(
+                            bucket.getAggregations(),
+                            new Derivative(name(), gradient, xDiff, formatter, metadata())
+                        )
+                    )
+                );
             } else {
                 newBuckets.add(bucket);
             }

modules/aggregations/src/main/java/org/elasticsearch/aggregations/pipeline/MovFnPipelineAggregator.java

@@ -25,8 +25,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 
 import static org.elasticsearch.search.aggregations.pipeline.BucketHelpers.resolveBucketValue;
 
@@ -117,12 +115,11 @@ public InternalAggregation reduce(InternalAggregation aggregation, AggregationRe
                     vars,
                     values.subList(fromIndex, toIndex).stream().mapToDouble(Double::doubleValue).toArray()
                 );
-
-                List<InternalAggregation> aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false)
-                    .map(InternalAggregation.class::cast)
-                    .collect(Collectors.toCollection(ArrayList::new));
-                aggs.add(new InternalSimpleValue(name(), result, formatter, metadata()));
-                newBucket = factory.createBucket(factory.getKey(bucket), bucket.getDocCount(), InternalAggregations.from(aggs));
+                newBucket = factory.createBucket(
+                    factory.getKey(bucket),
+                    bucket.getDocCount(),
+                    InternalAggregations.append(bucket.getAggregations(), new InternalSimpleValue(name(), result, formatter, metadata()))
+                );
                 index++;
             }
             newBuckets.add(newBucket);

server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregations.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.pipeline.SiblingPipelineAggregator;
@@ -75,14 +76,15 @@ public List<InternalAggregation> asList() {
     }
 
     private Map<String, InternalAggregation> asMap() {
-        if (aggregationsAsMap == null) {
+        var res = aggregationsAsMap;
+        if (res == null) {
             Map<String, InternalAggregation> newAggregationsAsMap = Maps.newMapWithExpectedSize(aggregations.size());
             for (InternalAggregation aggregation : aggregations) {
                 newAggregationsAsMap.put(aggregation.getName(), aggregation);
             }
-            this.aggregationsAsMap = unmodifiableMap(newAggregationsAsMap);
+            res = this.aggregationsAsMap = unmodifiableMap(newAggregationsAsMap);
         }
-        return aggregationsAsMap;
+        return res;
     }
 
     /**
@@ -147,13 +149,27 @@ public static InternalAggregations fromXContent(XContentParser parser) throws IO
         return new InternalAggregations(aggregations);
     }
 
+    public static InternalAggregations from(InternalAggregation aggregation) {
+        return new InternalAggregations(List.of(aggregation));
+    }
+
     public static InternalAggregations from(List<InternalAggregation> aggregations) {
         if (aggregations.isEmpty()) {
             return EMPTY;
         }
+        if (aggregations.size() == 1) {
+            return from(aggregations.get(0));
+        }
         return new InternalAggregations(aggregations);
     }
 
+    public static InternalAggregations append(InternalAggregations aggs, InternalAggregation toAppend) {
+        if (aggs.aggregations.isEmpty()) {
+            return from(toAppend);
+        }
+        return new InternalAggregations(CollectionUtils.appendToCopyNoNullElements(aggs.aggregations, toAppend));
+    }
+
     public static InternalAggregations readFrom(StreamInput in) throws IOException {
         return from(in.readNamedWriteableCollectionAsList(InternalAggregation.class));
     }
@@ -253,19 +269,7 @@ public static InternalAggregations reduce(List<InternalAggregations> aggregation
         }
         // handle special case when there is just one aggregation
         if (aggregationsList.size() == 1) {
-            final List<InternalAggregation> internalAggregations = aggregationsList.get(0).asList();
-            final List<InternalAggregation> reduced = new ArrayList<>(internalAggregations.size());
-            for (InternalAggregation aggregation : internalAggregations) {
-                if (aggregation.mustReduceOnSingleInternalAgg()) {
-                    try (AggregatorReducer aggregatorReducer = aggregation.getReducer(context.forAgg(aggregation.getName()), 1)) {
-                        aggregatorReducer.accept(aggregation);
-                        reduced.add(aggregatorReducer.get());
-                    }
-                } else {
-                    reduced.add(aggregation);
-                }
-            }
-            return from(reduced);
+            return reduce(aggregationsList.get(0), context);
         }
         // general case
         try (AggregatorsReducer reducer = new AggregatorsReducer(aggregationsList.get(0), context, aggregationsList.size())) {
@@ -276,6 +280,29 @@ public static InternalAggregations reduce(List<InternalAggregations> aggregation
         }
     }
 
+    public static InternalAggregations reduce(InternalAggregations aggregations, AggregationReduceContext context) {
+        final List<InternalAggregation> internalAggregations = aggregations.asList();
+        int size = internalAggregations.size();
+        if (size == 0) {
+            return EMPTY;
+        }
+        boolean noneReduced = true;
+        final List<InternalAggregation> reduced = new ArrayList<>(size);
+        for (int i = 0; i < size; i++) {
+            InternalAggregation aggregation = internalAggregations.get(i);
+            if (aggregation.mustReduceOnSingleInternalAgg()) {
+                noneReduced = false;
+                try (AggregatorReducer aggregatorReducer = aggregation.getReducer(context.forAgg(aggregation.getName()), 1)) {
+                    aggregatorReducer.accept(aggregation);
+                    reduced.add(aggregatorReducer.get());
+                }
+            } else {
+                reduced.add(aggregation);
+            }
+        }
+        return noneReduced ? aggregations : from(reduced);
+    }
+
     /**
      * Finalizes the sampling for all the internal aggregations
      * @param samplingContext the sampling context

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

@@ -377,7 +377,7 @@ public void accept(long key) {
         iterateEmptyBuckets(list, list.listIterator(), counter);
         reduceContext.consumeBucketsAndMaybeBreak(counter.size);
 
-        InternalAggregations reducedEmptySubAggs = InternalAggregations.reduce(List.of(emptyBucketInfo.subAggregations), reduceContext);
+        InternalAggregations reducedEmptySubAggs = InternalAggregations.reduce(emptyBucketInfo.subAggregations, reduceContext);
         ListIterator<Bucket> iter = list.listIterator();
         iterateEmptyBuckets(list, iter, new LongConsumer() {
             private int size = 0;

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

@@ -349,10 +349,7 @@ public void accept(double key) {
         /*
          * Now that we're sure we have space we allocate all the buckets.
          */
-        InternalAggregations reducedEmptySubAggs = InternalAggregations.reduce(
-            Collections.singletonList(emptyBucketInfo.subAggregations),
-            reduceContext
-        );
+        InternalAggregations reducedEmptySubAggs = InternalAggregations.reduce(emptyBucketInfo.subAggregations, reduceContext);
         ListIterator<Bucket> iter = list.listIterator();
         iterateEmptyBuckets(list, iter, new DoubleConsumer() {
             private int size;

server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java

@@ -22,8 +22,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 
 import static org.elasticsearch.search.aggregations.pipeline.BucketHelpers.resolveBucketValue;
 
@@ -80,13 +78,11 @@ public InternalAggregation reduce(InternalAggregation aggregation, AggregationRe
                 if (returned == null) {
                     newBuckets.add(bucket);
                 } else {
-                    final List<InternalAggregation> aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false)
-                        .collect(Collectors.toCollection(ArrayList::new));
-
-                    InternalSimpleValue simpleValue = new InternalSimpleValue(name(), returned.doubleValue(), formatter, metadata());
-                    aggs.add(simpleValue);
                     InternalMultiBucketAggregation.InternalBucket newBucket = originalAgg.createBucket(
-                        InternalAggregations.from(aggs),
+                        InternalAggregations.append(
+                            bucket.getAggregations(),
+                            new InternalSimpleValue(name(), returned.doubleValue(), formatter, metadata())
+                        ),
                         bucket
                     );
                     newBuckets.add(newBucket);