pre-size attribute set/map

Author: idegtiarenko

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/AttributeMap.java

@@ -6,6 +6,7 @@
  */
 package org.elasticsearch.xpack.esql.core.expression;
 
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
 
@@ -169,21 +170,25 @@ public AttributeMap() {
         delegate = new LinkedHashMap<>();
     }
 
+    public AttributeMap(int expectedSize) {
+        delegate = Maps.newLinkedHashMapWithExpectedSize(expectedSize);
+    }
+
     public AttributeMap(Attribute key, E value) {
         delegate = new LinkedHashMap<>();
         add(key, value);
     }
 
     public AttributeMap<E> combine(AttributeMap<E> other) {
-        AttributeMap<E> combine = new AttributeMap<>();
+        AttributeMap<E> combine = new AttributeMap<>(this.size() + other.size());
         combine.addAll(this);
         combine.addAll(other);
 
         return combine;
     }
 
     public AttributeMap<E> subtract(AttributeMap<E> other) {
-        AttributeMap<E> diff = new AttributeMap<>();
+        AttributeMap<E> diff = new AttributeMap<>(this.size());
         for (Entry<AttributeWrapper, E> entry : this.delegate.entrySet()) {
             if (other.delegate.containsKey(entry.getKey()) == false) {
                 diff.delegate.put(entry.getKey(), entry.getValue());
@@ -197,7 +202,7 @@ public AttributeMap<E> intersect(AttributeMap<E> other) {
         AttributeMap<E> smaller = (other.size() > size() ? this : other);
         AttributeMap<E> larger = (smaller == this ? other : this);
 
-        AttributeMap<E> intersect = new AttributeMap<>();
+        AttributeMap<E> intersect = new AttributeMap<>(smaller.size());
         for (Entry<AttributeWrapper, E> entry : smaller.delegate.entrySet()) {
             if (larger.delegate.containsKey(entry.getKey())) {
                 intersect.delegate.put(entry.getKey(), entry.getValue());
@@ -385,7 +390,7 @@ public static <E> Builder<E> builder(AttributeMap<E> map) {
     }
 
     public static class Builder<E> {
-        private AttributeMap<E> map = new AttributeMap<>();
+        private final AttributeMap<E> map = new AttributeMap<>();
 
         private Builder() {}
 

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/AttributeSet.java

@@ -30,13 +30,16 @@ public AttributeSet() {
         delegate = new AttributeMap<>();
     }
 
+    public AttributeSet(int expectedSize) {
+        delegate = new AttributeMap<>(expectedSize);
+    }
+
     public AttributeSet(Attribute attr) {
         delegate = new AttributeMap<>(attr, PRESENT);
     }
 
     public AttributeSet(Collection<? extends Attribute> attr) {
-        delegate = new AttributeMap<>();
-
+        delegate = new AttributeMap<>(attr.size());
         for (Attribute a : attr) {
             delegate.add(a, PRESENT);
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QueryPlan.java

@@ -48,11 +48,15 @@ public AttributeSet outputSet() {
 
     public AttributeSet inputSet() {
         if (lazyInputSet == null) {
-            List<Attribute> attrs = new ArrayList<>();
+            int size = 0;
             for (PlanType child : children()) {
-                attrs.addAll(child.output());
+                size += child.output().size();
             }
-            lazyInputSet = new AttributeSet(attrs);
+            var attributes = new AttributeSet(size);
+            for (PlanType child : children()) {
+                attributes.addAll(child.output());
+            }
+            lazyInputSet = attributes;
         }
         return lazyInputSet;
     }