Try to save memory on aggregations

server/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java

@@ -67,8 +67,11 @@
 import java.util.Map;
 import java.util.function.Function;
 import java.util.function.IntFunction;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
+import static java.util.stream.Collectors.toList;
+
 public final class SearchPhaseController {
     private static final ScoreDoc[] EMPTY_DOCS = new ScoreDoc[0];
 
@@ -429,7 +432,7 @@ public ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResul
      * @see QuerySearchResult#consumeProfileResult()
      */
     private ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResult> queryResults,
-                                                List<InternalAggregations> bufferedAggs, List<TopDocs> bufferedTopDocs,
+                                                List<Supplier<InternalAggregations>> bufferedAggs, List<TopDocs> bufferedTopDocs,
                                                 TopDocsStats topDocsStats, int numReducePhases, boolean isScrollRequest,
                                                 InternalAggregation.ReduceContextBuilder aggReduceContextBuilder,
                                                 boolean performFinalReduce) {
@@ -453,7 +456,7 @@ private ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResu
         final boolean hasSuggest = firstResult.suggest() != null;
         final boolean hasProfileResults = firstResult.hasProfileResults();
         final boolean consumeAggs;
-        final List<InternalAggregations> aggregationsList;
+        final List<Supplier<InternalAggregations>> aggregationsList;
         if (bufferedAggs != null) {
             consumeAggs = false;
             // we already have results from intermediate reduces and just need to perform the final reduce
@@ -492,7 +495,7 @@ private ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResu
                 }
             }
             if (consumeAggs) {
-                aggregationsList.add((InternalAggregations) result.consumeAggs());
+                aggregationsList.add(result.consumeAggs());
             }
             if (hasProfileResults) {
                 String key = result.getSearchShardTarget().toString();
@@ -508,8 +511,9 @@ private ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResu
             reducedSuggest = new Suggest(Suggest.reduce(groupedSuggestions));
             reducedCompletionSuggestions = reducedSuggest.filter(CompletionSuggestion.class);
         }
-        final InternalAggregations aggregations = aggregationsList.isEmpty() ? null : InternalAggregations.topLevelReduce(aggregationsList,
-                    performFinalReduce ? aggReduceContextBuilder.forFinalReduction() : aggReduceContextBuilder.forPartialReduction());
+        final InternalAggregations aggregations = aggregationsList.isEmpty() ? null : InternalAggregations.topLevelReduce(
+                aggregationsList.stream().map(Supplier::get).collect(toList()),
+                performFinalReduce ? aggReduceContextBuilder.forFinalReduction() : aggReduceContextBuilder.forPartialReduction());
         final SearchProfileShardResults shardResults = profileResults.isEmpty() ? null : new SearchProfileShardResults(profileResults);
         final SortedTopDocs sortedTopDocs = sortDocs(isScrollRequest, queryResults, bufferedTopDocs, topDocsStats, from, size,
             reducedCompletionSuggestions);
@@ -600,7 +604,7 @@ public InternalSearchResponse buildResponse(SearchHits hits) {
      */
     static final class QueryPhaseResultConsumer extends ArraySearchPhaseResults<SearchPhaseResult> {
         private final SearchShardTarget[] processedShards;
-        private final InternalAggregations[] aggsBuffer;
+        private final Supplier<InternalAggregations>[] aggsBuffer;
         private final TopDocs[] topDocsBuffer;
         private final boolean hasAggs;
         private final boolean hasTopDocs;
@@ -642,7 +646,9 @@ private QueryPhaseResultConsumer(SearchProgressListener progressListener, Search
             this.progressListener = progressListener;
             this.processedShards = new SearchShardTarget[expectedResultSize];
             // no need to buffer anything if we have less expected results. in this case we don't consume any results ahead of time.
-            this.aggsBuffer = new InternalAggregations[hasAggs ? bufferSize : 0];
+            @SuppressWarnings("unchecked")
+            Supplier<InternalAggregations>[] aggsBuffer = new Supplier[hasAggs ? bufferSize : 0];
+            this.aggsBuffer = aggsBuffer;
             this.topDocsBuffer = new TopDocs[hasTopDocs ? bufferSize : 0];
             this.hasTopDocs = hasTopDocs;
             this.hasAggs = hasAggs;
@@ -665,10 +671,11 @@ private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
             if (querySearchResult.isNull() == false) {
                 if (index == bufferSize) {
                     if (hasAggs) {
-                        ReduceContext reduceContext = aggReduceContextBuilder.forPartialReduction();
-                        InternalAggregations reducedAggs = InternalAggregations.topLevelReduce(Arrays.asList(aggsBuffer), reduceContext);
+                        InternalAggregations reducedAggs = InternalAggregations.topLevelReduce(
+                                Arrays.stream(aggsBuffer).map(Supplier::get).collect(toList()),
+                                aggReduceContextBuilder.forPartialReduction());
                         Arrays.fill(aggsBuffer, null);
-                        aggsBuffer[0] = reducedAggs;
+                        aggsBuffer[0] = () -> reducedAggs;
                     }
                     if (hasTopDocs) {
                         TopDocs reducedTopDocs = mergeTopDocs(Arrays.asList(topDocsBuffer),
@@ -681,12 +688,12 @@ private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
                     index = 1;
                     if (hasAggs || hasTopDocs) {
                         progressListener.notifyPartialReduce(SearchProgressListener.buildSearchShards(processedShards),
-                            topDocsStats.getTotalHits(), hasAggs ? aggsBuffer[0] : null, numReducePhases);
+                            topDocsStats.getTotalHits(), hasAggs ? aggsBuffer[0].get() : null, numReducePhases);
                     }
                 }
                 final int i = index++;
                 if (hasAggs) {
-                    aggsBuffer[i] = (InternalAggregations) querySearchResult.consumeAggs();
+                    aggsBuffer[i] = querySearchResult.consumeAggs();
                 }
                 if (hasTopDocs) {
                     final TopDocsAndMaxScore topDocs = querySearchResult.consumeTopDocs(); // can't be null
@@ -698,7 +705,7 @@ private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
             processedShards[querySearchResult.getShardIndex()] = querySearchResult.getSearchShardTarget();
         }
 
-        private synchronized List<InternalAggregations> getRemainingAggs() {
+        private synchronized List<Supplier<InternalAggregations>> getRemainingAggs() {
             return hasAggs ? Arrays.asList(aggsBuffer).subList(0, index) : null;
         }
 

server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java

@@ -0,0 +1,115 @@
+/*
+ * 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.common.io.stream;
+
+import org.elasticsearch.Version;
+import org.elasticsearch.common.bytes.BytesReference;
+
+import java.io.IOException;
+import java.util.function.Supplier;
+
+/**
+ * A holder for {@link Writeable}s that can delays reading the underlying
+ * {@linkplain Writeable} when it is read from a remote node.
+ */
+public abstract class DelayableWriteable<T extends Writeable> implements Supplier<T>, Writeable {
+    /**
+     * Build a {@linkplain DelayableWriteable} that wraps an existing object
+     * but is serialized so that deserializing it can be delayed.
+     */
+    public static <T extends Writeable> DelayableWriteable<T> referencing(T reference) {
+        return new Referencing<>(reference);
+    }
+    /**
+     * Build a {@linkplain DelayableWriteable} that copies a buffer from
+     * the provided {@linkplain StreamInput} and deserializes the buffer
+     * when {@link Supplier#get()} is called.
+     */
+    public static <T extends Writeable> DelayableWriteable<T> delayed(Writeable.Reader<T> reader, StreamInput in) throws IOException {
+        return new Delayed<>(reader, in);
+    }
+
+    private DelayableWriteable() {}
+
+    public abstract boolean isDelayed();
+
+    private static class Referencing<T extends Writeable> extends DelayableWriteable<T> {
+        private T reference;
+
+        Referencing(T reference) {
+            this.reference = reference;
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            try (BytesStreamOutput buffer = new BytesStreamOutput()) {
+                reference.writeTo(buffer);
+                out.writeBytesReference(buffer.bytes());
+            }
+        }
+
+        @Override
+        public T get() {
+            return reference;
+        }
+
+        @Override
+        public boolean isDelayed() {
+            return false;
+        }
+    }
+
+    private static class Delayed<T extends Writeable> extends DelayableWriteable<T> {
+        private final Writeable.Reader<T> reader;
+        private final Version remoteVersion;
+        private final BytesReference serialized;
+        private final NamedWriteableRegistry registry;
+
+        Delayed(Writeable.Reader<T> reader, StreamInput in) throws IOException {
+            this.reader = reader;
+            remoteVersion = in.getVersion();
+            serialized = in.readBytesReference();
+            registry = in.namedWriteableRegistry();
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public T get() {
+            try {
+                try (StreamInput in = registry == null ?
+                        serialized.streamInput() : new NamedWriteableAwareStreamInput(serialized.streamInput(), registry)) {
+                    in.setVersion(remoteVersion);
+                    return reader.read(in);
+                }
+            } catch (IOException e) {
+                throw new RuntimeException("unexpected error expanding aggregations", e);
+            }
+        }
+
+        @Override
+        public boolean isDelayed() {
+            return true;
+        }
+    }
+}

server/src/main/java/org/elasticsearch/common/io/stream/FilterStreamInput.java

@@ -94,4 +94,9 @@ public void setVersion(Version version) {
     protected void ensureCanReadBytes(int length) throws EOFException {
         delegate.ensureCanReadBytes(length);
     }
+
+    @Override
+    public NamedWriteableRegistry namedWriteableRegistry() {
+        return delegate.namedWriteableRegistry();
+    }
 }

server/src/main/java/org/elasticsearch/common/io/stream/NamedWriteableAwareStreamInput.java

@@ -52,4 +52,9 @@ public <C extends NamedWriteable> C readNamedWriteable(@SuppressWarnings("unused
             + "] than it was read from [" + name + "].";
         return c;
     }
+
+    @Override
+    public NamedWriteableRegistry namedWriteableRegistry() {
+        return namedWriteableRegistry;
+    }
 }

server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java

@@ -1093,6 +1093,14 @@ public <T extends Exception> T readException() throws IOException {
         return null;
     }
 
+    /**
+     * Get the registry of named writeables is his stream has one,
+     * {@code null} otherwise.
+     */
+    public NamedWriteableRegistry namedWriteableRegistry() {
+        return null;
+    }
+
     /**
      * Reads a {@link NamedWriteable} from the current stream, by first reading its name and then looking for
      * the corresponding entry in the registry by name, so that the proper object can be read and returned.

server/src/main/java/org/elasticsearch/search/query/QuerySearchResult.java

@@ -22,13 +22,13 @@
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.Version;
+import org.elasticsearch.common.io.stream.DelayableWriteable;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchShardTarget;
-import org.elasticsearch.search.aggregations.Aggregations;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.InternalAggregations;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
@@ -40,6 +40,7 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.common.lucene.Lucene.readTopDocs;
@@ -54,7 +55,7 @@ public final class QuerySearchResult extends SearchPhaseResult {
     private TotalHits totalHits;
     private float maxScore = Float.NaN;
     private DocValueFormat[] sortValueFormats;
-    private InternalAggregations aggregations;
+    private DelayableWriteable<InternalAggregations> aggregations;
     private boolean hasAggs;
     private Suggest suggest;
     private boolean searchTimedOut;
@@ -196,21 +197,21 @@ public boolean hasAggs() {
      * Returns and nulls out the aggregation for this search results. This allows to free up memory once the aggregation is consumed.
      * @throws IllegalStateException if the aggregations have already been consumed.
      */
-    public Aggregations consumeAggs() {
+    public Supplier<InternalAggregations> consumeAggs() {
         if (aggregations == null) {
             throw new IllegalStateException("aggs already consumed");
         }
-        Aggregations aggs = aggregations;
+        Supplier<InternalAggregations> aggs = aggregations;
         aggregations = null;
         return aggs;
     }
 
     public void aggregations(InternalAggregations aggregations) {
-        this.aggregations = aggregations;
+        this.aggregations = aggregations == null ? null : DelayableWriteable.referencing(aggregations);
         hasAggs = aggregations != null;
     }
 
-    public InternalAggregations aggregations() {
+    public DelayableWriteable<InternalAggregations> aggregations() {
         return aggregations;
     }
 
@@ -314,18 +315,22 @@ public void readFromWithId(SearchContextId id, StreamInput in) throws IOExceptio
         }
         setTopDocs(readTopDocs(in));
         if (hasAggs = in.readBoolean()) {
-            aggregations = new InternalAggregations(in);
+            if (in.getVersion().before(Version.V_8_0_0)) {
+                aggregations = DelayableWriteable.referencing(new InternalAggregations(in));
+            } else {
+                aggregations = DelayableWriteable.delayed(InternalAggregations::new, in);
+            }
         }
         if (in.getVersion().before(Version.V_7_2_0)) {
             List<SiblingPipelineAggregator> pipelineAggregators = in.readNamedWriteableList(PipelineAggregator.class).stream()
                 .map(a -> (SiblingPipelineAggregator) a).collect(Collectors.toList());
             if (hasAggs && pipelineAggregators.isEmpty() == false) {
-                List<InternalAggregation> internalAggs = aggregations.asList().stream()
+                List<InternalAggregation> internalAggs = aggregations.get().asList().stream()
                     .map(agg -> (InternalAggregation) agg).collect(Collectors.toList());
                 //Earlier versions serialize sibling pipeline aggs separately as they used to be set to QuerySearchResult directly, while
                 //later versions include them in InternalAggregations. Note that despite serializing sibling pipeline aggs as part of
                 //InternalAggregations is supported since 6.7.0, the shards set sibling pipeline aggs to InternalAggregations only from 7.1.
-                this.aggregations = new InternalAggregations(internalAggs, pipelineAggregators);
+                this.aggregations = DelayableWriteable.referencing(new InternalAggregations(internalAggs, pipelineAggregators));
             }
         }
         if (in.readBoolean()) {
@@ -366,7 +371,11 @@ public void writeToNoId(StreamOutput out) throws IOException {
             out.writeBoolean(false);
         } else {
             out.writeBoolean(true);
-            aggregations.writeTo(out);
+            if (out.getVersion().before(Version.V_8_0_0)) {
+                aggregations.get().writeTo(out);
+            } else {
+                aggregations.writeTo(out);
+            }
         }
         if (out.getVersion().before(Version.V_7_2_0)) {
             //Earlier versions expect sibling pipeline aggs separately as they used to be set to QuerySearchResult directly,
@@ -375,7 +384,7 @@ public void writeToNoId(StreamOutput out) throws IOException {
             if (aggregations == null) {
                 out.writeNamedWriteableList(Collections.emptyList());
             } else {
-                out.writeNamedWriteableList(aggregations.getTopLevelPipelineAggregators());
+                out.writeNamedWriteableList(aggregations.get().getTopLevelPipelineAggregators());
             }
         }
         if (suggest == null) {