upd

Author: idegtiarenko

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

@@ -171,7 +171,7 @@ private AttributeMap() {
         delegate = new LinkedHashMap<>();
     }
 
-    public AttributeMap(int expectedSize) {
+    private AttributeMap(int expectedSize) {
         delegate = Maps.newLinkedHashMapWithExpectedSize(expectedSize);
     }
 
@@ -184,7 +184,6 @@ public AttributeMap<E> combine(AttributeMap<E> other) {
         AttributeMap<E> combine = new AttributeMap<>(this.size() + other.size());
         combine.addAll(this);
         combine.addAll(other);
-
         return combine;
     }
 
@@ -195,7 +194,6 @@ public AttributeMap<E> subtract(AttributeMap<E> other) {
                 diff.delegate.put(entry.getKey(), entry.getValue());
             }
         }
-
         return diff;
     }
 
@@ -397,13 +395,20 @@ public static <E> AttributeMap<E> of(Attribute key, E value) {
     }
 
     public static <E> Builder<E> builder() {
-        return new Builder<>();
+        return new Builder<>(new AttributeMap<>());
+    }
+
+    public static <E> Builder<E> builder(int expectedSize) {
+        return new Builder<>(new AttributeMap<>(expectedSize));
     }
 
     public static class Builder<E> {
-        private final AttributeMap<E> map = new AttributeMap<>();
 
-        private Builder() {}
+        private final AttributeMap<E> map;
+
+        public Builder(AttributeMap<E> map) {
+            this.map = map;
+        }
 
         public E put(Attribute attr, E value) {
             return map.add(attr, value);

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

@@ -165,33 +165,39 @@ public String toString() {
     }
 
     public Builder asBuilder() {
-        return new Builder().addAll(this);
+        return builder(size()).addAll(this);
     }
 
     public static AttributeSet of(Attribute... attrs) {
-        final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
+        final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder(attrs.length);
         for (var a : attrs) {
             mapBuilder.put(a, PRESENT);
         }
         return new AttributeSet(mapBuilder.build());
     }
 
     public static AttributeSet of(Collection<? extends Attribute> c) {
-        final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
+        final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder(c.size());
         for (var a : c) {
             mapBuilder.put(a, PRESENT);
         }
         return new AttributeSet(mapBuilder.build());
     }
 
     public static Builder builder() {
-        return new Builder();
+        return new Builder(AttributeMap.builder());
+    }
+
+    public static Builder builder(int expectedSize) {
+        return new Builder(AttributeMap.builder(expectedSize));
     }
 
     public static class Builder {
-        private final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
+        private final AttributeMap.Builder<Object> mapBuilder;
 
-        private Builder() {}
+        private Builder(AttributeMap.Builder<Object> mapBuilder) {
+            this.mapBuilder = mapBuilder;
+        }
 
         public Builder add(Attribute attr) {
             mapBuilder.put(attr, PRESENT);

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

@@ -35,18 +35,6 @@ public static List<Attribute> asAttributes(List<? extends NamedExpression> named
         return list;
     }
 
-    public static AttributeMap<Expression> asAttributeMap(List<? extends NamedExpression> named) {
-        if (named.isEmpty()) {
-            return AttributeMap.emptyAttributeMap();
-        }
-
-        AttributeMap.Builder<Expression> mapBuilder = AttributeMap.builder();
-        for (NamedExpression exp : named) {
-            mapBuilder.put(exp.toAttribute(), exp);
-        }
-        return mapBuilder.build();
-    }
-
     public static boolean anyMatch(List<? extends Expression> exps, Predicate<? super Expression> predicate) {
         for (Expression exp : exps) {
             if (exp.anyMatch(predicate)) {
@@ -121,11 +109,19 @@ public static AttributeSet references(List<? extends Expression> exps) {
             return AttributeSet.EMPTY;
         }
 
-        var setBuilder = AttributeSet.builder();
+        var size = 0;
+        var references = new ArrayList<AttributeSet>(exps.size());
         for (Expression exp : exps) {
-            setBuilder.addAll(exp.references());
+            var refs = exp.references();
+            size += refs.size();
+            references.add(refs);
+        }
+
+        var builder = AttributeSet.builder(size);
+        for (AttributeSet refs : references) {
+            builder.addAll(refs);
         }
-        return setBuilder.build();
+        return builder.build();
     }
 
     public static String name(Expression e) {