Pre-size collections (#126382)

Author: idegtiarenko

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/CombineProjections.java

@@ -117,7 +117,7 @@ private static List<NamedExpression> combineProjections(List<? extends NamedExpr
         // collect named expressions declaration in the lower list
         AttributeMap.Builder<NamedExpression> namedExpressionsBuilder = AttributeMap.builder();
         // while also collecting the alias map for resolving the source (f1 = 1, f2 = f1, etc..)
-        AttributeMap.Builder<Expression> aliasesBuilder = AttributeMap.builder();
+        AttributeMap.Builder<Expression> aliasesBuilder = AttributeMap.builder(lower.size());
         for (NamedExpression ne : lower) {
             // record the alias
             aliasesBuilder.put(ne.toAttribute(), Alias.unwrap(ne));
@@ -128,7 +128,7 @@ private static List<NamedExpression> combineProjections(List<? extends NamedExpr
                 namedExpressionsBuilder.put(ne.toAttribute(), as.replaceChild(aliasesBuilder.build().resolve(child, child)));
             }
         }
-        List<NamedExpression> replaced = new ArrayList<>();
+        List<NamedExpression> replaced = new ArrayList<>(upper.size());
         var namedExpressions = namedExpressionsBuilder.build();
 
         // replace any matching attribute with a lower alias (if there's a match)

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/Layout.java

@@ -7,14 +7,14 @@
 
 package org.elasticsearch.xpack.esql.planner;
 
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.xpack.esql.core.expression.NameId;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -65,7 +65,7 @@ class Builder {
         private final List<ChannelSet> channels;
 
         public Builder() {
-            channels = new ArrayList<>();
+            this(new ArrayList<>());
         }
 
         Builder(List<ChannelSet> channels) {
@@ -76,7 +76,7 @@ public Builder() {
          * Appends a new channel to the layout. The channel is mapped to one or more attribute ids.
          */
         public Builder append(ChannelSet set) {
-            if (set.nameIds.size() < 1) {
+            if (set.nameIds.isEmpty()) {
                 throw new IllegalArgumentException("Channel must be mapped to at least one id.");
             }
             channels.add(set);
@@ -104,7 +104,11 @@ public Builder append(Collection<? extends NamedExpression> attributes) {
          * Build a new {@link Layout}.
          */
         public Layout build() {
-            Map<NameId, ChannelAndType> layout = new HashMap<>();
+            int size = 0;
+            for (ChannelSet channel : channels) {
+                size += channel.nameIds.size();
+            }
+            Map<NameId, ChannelAndType> layout = Maps.newHashMapWithExpectedSize(size);
             int numberOfChannels = 0;
             for (ChannelSet set : channels) {
                 int channel = numberOfChannels++;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.compute.Describable;
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.data.Block;
@@ -117,9 +118,6 @@
 import org.elasticsearch.xpack.esql.session.Configuration;
 
 import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -501,8 +499,8 @@ private PhysicalOperation planGrok(GrokExec grok, LocalExecutionPlannerContext c
         Layout.Builder layoutBuilder = source.layout.builder();
         List<Attribute> extractedFields = grok.extractedFields();
         layoutBuilder.append(extractedFields);
-        Map<String, Integer> fieldToPos = new HashMap<>(extractedFields.size());
-        Map<String, ElementType> fieldToType = new HashMap<>(extractedFields.size());
+        Map<String, Integer> fieldToPos = Maps.newHashMapWithExpectedSize(extractedFields.size());
+        Map<String, ElementType> fieldToType = Maps.newHashMapWithExpectedSize(extractedFields.size());
         ElementType[] types = new ElementType[extractedFields.size()];
         List<Attribute> extractedFieldsFromPattern = grok.pattern().extractedFields();
         for (int i = 0; i < extractedFields.size(); i++) {
@@ -559,7 +557,8 @@ private PhysicalOperation planEnrich(EnrichExec enrich, LocalExecutionPlannerCon
     private PhysicalOperation planRerank(RerankExec rerank, LocalExecutionPlannerContext context) {
         PhysicalOperation source = plan(rerank.child(), context);
 
-        Map<ColumnInfoImpl, EvalOperator.ExpressionEvaluator.Factory> rerankFieldsEvaluatorSuppliers = new LinkedHashMap<>();
+        Map<ColumnInfoImpl, EvalOperator.ExpressionEvaluator.Factory> rerankFieldsEvaluatorSuppliers = Maps
+            .newLinkedHashMapWithExpectedSize(rerank.rerankFields().size());
 
         for (var rerankField : rerank.rerankFields()) {
             rerankFieldsEvaluatorSuppliers.put(
@@ -734,31 +733,13 @@ private PhysicalOperation planProject(ProjectExec project, LocalExecutionPlanner
         List<Integer> projectionList = new ArrayList<>(projections.size());
 
         Layout.Builder layout = new Layout.Builder();
-        Map<Integer, Layout.ChannelSet> inputChannelToOutputIds = new HashMap<>();
-        for (int index = 0, size = projections.size(); index < size; index++) {
-            NamedExpression ne = projections.get(index);
-
-            NameId inputId = null;
-            if (ne instanceof Alias a) {
-                inputId = ((NamedExpression) a.child()).id();
-            } else {
-                inputId = ne.id();
-            }
+        for (NamedExpression ne : projections) {
+            NameId inputId = ne instanceof Alias a ? ((NamedExpression) a.child()).id() : ne.id();
             Layout.ChannelAndType input = source.layout.get(inputId);
             if (input == null) {
                 throw new IllegalStateException("can't find input for [" + ne + "]");
             }
-            Layout.ChannelSet channelSet = inputChannelToOutputIds.get(input.channel());
-            if (channelSet == null) {
-                channelSet = new Layout.ChannelSet(new HashSet<>(), input.type());
-                channelSet.nameIds().add(ne.id());
-                layout.append(channelSet);
-            } else {
-                channelSet.nameIds().add(ne.id());
-            }
-            if (channelSet.type() != input.type()) {
-                throw new IllegalArgumentException("type mismatch for aliases");
-            }
+            layout.append(ne);
             projectionList.add(input.channel());
         }