[9.0] ESQL: make AttributeMap and AttributeSet immutable (#125938)

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

@@ -66,7 +66,7 @@ public Nullability nullable() {
 
     @Override
     public AttributeSet references() {
-        return new AttributeSet(this);
+        return AttributeSet.of(this);
     }
 
     public Attribute withLocation(Source source) {

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

@@ -16,6 +16,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import static java.util.Collections.emptyMap;
@@ -34,7 +35,7 @@
  */
 public final class AttributeMap<E> implements Map<Attribute, E> {
 
-    static class AttributeWrapper {
+    private static class AttributeWrapper {
 
         private final Attribute attr;
 
@@ -165,7 +166,7 @@ private AttributeMap(Map<AttributeWrapper, E> other) {
         delegate = other;
     }
 
-    public AttributeMap() {
+    private AttributeMap() {
         delegate = new LinkedHashMap<>();
     }
 
@@ -220,12 +221,22 @@ public boolean subsetOf(AttributeMap<E> other) {
         return true;
     }
 
-    public void add(Attribute key, E value) {
-        put(key, value);
+    private E add(Attribute key, E value) {
+        return delegate.put(new AttributeWrapper(key), value);
+    }
+
+    private void addAll(AttributeMap<E> other) {
+        for (Entry<? extends Attribute, ? extends E> entry : other.entrySet()) {
+            add(entry.getKey(), entry.getValue());
+        }
     }
 
-    public void addAll(AttributeMap<E> other) {
-        putAll(other);
+    private E addIfAbsent(Attribute key, Function<? super Attribute, ? extends E> mappingFunction) {
+        return delegate.computeIfAbsent(new AttributeWrapper(key), k -> mappingFunction.apply(k.attr));
+    }
+
+    private E delete(Object key) {
+        return key instanceof NamedExpression ne ? delegate.remove(new AttributeWrapper(ne.toAttribute())) : null;
     }
 
     public Set<String> attributeNames() {
@@ -249,7 +260,7 @@ public boolean isEmpty() {
 
     @Override
     public boolean containsKey(Object key) {
-        return key instanceof NamedExpression ne ? delegate.containsKey(new AttributeWrapper(ne.toAttribute())) : false;
+        return key instanceof NamedExpression ne && delegate.containsKey(new AttributeWrapper(ne.toAttribute()));
     }
 
     @Override
@@ -293,24 +304,22 @@ public E resolve(Object key, E defaultValue) {
 
     @Override
     public E put(Attribute key, E value) {
-        return delegate.put(new AttributeWrapper(key), value);
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public void putAll(Map<? extends Attribute, ? extends E> m) {
-        for (Entry<? extends Attribute, ? extends E> entry : m.entrySet()) {
-            put(entry.getKey(), entry.getValue());
-        }
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public E remove(Object key) {
-        return key instanceof NamedExpression ne ? delegate.remove(new AttributeWrapper(ne.toAttribute())) : null;
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public void clear() {
-        delegate.clear();
+        throw new UnsupportedOperationException();
     }
 
     @Override
@@ -376,29 +385,50 @@ public String toString() {
         return delegate.toString();
     }
 
-    public static <E> Builder<E> builder() {
-        return new Builder<>();
+    public static <E> AttributeMap<E> of(Attribute key, E value) {
+        final AttributeMap<E> map = new AttributeMap<>();
+        map.add(key, value);
+        return map;
     }
 
-    public static <E> Builder<E> builder(AttributeMap<E> map) {
-        return new Builder<E>().putAll(map);
+    public static <E> Builder<E> builder() {
+        return new Builder<>();
     }
 
     public static class Builder<E> {
-        private AttributeMap<E> map = new AttributeMap<>();
+        private final AttributeMap<E> map = new AttributeMap<>();
 
         private Builder() {}
 
-        public Builder<E> put(Attribute attr, E value) {
-            map.add(attr, value);
-            return this;
+        public E put(Attribute attr, E value) {
+            return map.add(attr, value);
         }
 
         public Builder<E> putAll(AttributeMap<E> m) {
             map.addAll(m);
             return this;
         }
 
+        public E computeIfAbsent(Attribute key, Function<? super Attribute, ? extends E> mappingFunction) {
+            return map.addIfAbsent(key, mappingFunction);
+        }
+
+        public E remove(Object o) {
+            return map.delete(o);
+        }
+
+        public Set<Attribute> keySet() {
+            return map.keySet();
+        }
+
+        public boolean containsKey(Object key) {
+            return map.containsKey(key);
+        }
+
+        public boolean isEmpty() {
+            return map.isEmpty();
+        }
+
         public AttributeMap<E> build() {
             return map;
         }

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

@@ -26,22 +26,6 @@ public class AttributeSet implements Set<Attribute> {
 
     private final AttributeMap<Object> delegate;
 
-    public AttributeSet() {
-        delegate = new AttributeMap<>();
-    }
-
-    public AttributeSet(Attribute attr) {
-        delegate = new AttributeMap<>(attr, PRESENT);
-    }
-
-    public AttributeSet(Collection<? extends Attribute> attr) {
-        delegate = new AttributeMap<>();
-
-        for (Attribute a : attr) {
-            delegate.add(a, PRESENT);
-        }
-    }
-
     private AttributeSet(AttributeMap<Object> delegate) {
         this.delegate = delegate;
     }
@@ -113,44 +97,32 @@ public <T> T[] toArray(T[] a) {
 
     @Override
     public boolean add(Attribute e) {
-        return delegate.put(e, PRESENT) == null;
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public boolean remove(Object o) {
-        return delegate.remove(o) != null;
-    }
-
-    public void addAll(AttributeSet other) {
-        delegate.addAll(other.delegate);
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public boolean addAll(Collection<? extends Attribute> c) {
-        int size = delegate.size();
-        for (var e : c) {
-            delegate.put(e, PRESENT);
-        }
-        return delegate.size() != size;
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public boolean retainAll(Collection<?> c) {
-        return delegate.keySet().removeIf(e -> c.contains(e) == false);
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public boolean removeAll(Collection<?> c) {
-        int size = delegate.size();
-        for (var e : c) {
-            delegate.remove(e);
-        }
-        return delegate.size() != size;
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public void clear() {
-        delegate.clear();
+        throw new UnsupportedOperationException();
     }
 
     @Override
@@ -160,7 +132,7 @@ public Spliterator<Attribute> spliterator() {
 
     @Override
     public boolean removeIf(Predicate<? super Attribute> filter) {
-        return delegate.keySet().removeIf(filter);
+        throw new UnsupportedOperationException();
     }
 
     @Override
@@ -191,4 +163,76 @@ public int hashCode() {
     public String toString() {
         return delegate.keySet().toString();
     }
+
+    public Builder asBuilder() {
+        return new Builder().addAll(this);
+    }
+
+    public static AttributeSet of(Attribute... attrs) {
+        final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
+        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();
+        for (var a : c) {
+            mapBuilder.put(a, PRESENT);
+        }
+        return new AttributeSet(mapBuilder.build());
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public static class Builder {
+        private final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
+
+        private Builder() {}
+
+        public Builder add(Attribute attr) {
+            mapBuilder.put(attr, PRESENT);
+            return this;
+        }
+
+        public boolean remove(Object o) {
+            return mapBuilder.remove(o) != null;
+        }
+
+        public Builder addAll(AttributeSet other) {
+            mapBuilder.putAll(other.delegate);
+            return this;
+        }
+
+        public Builder addAll(Builder other) {
+            mapBuilder.putAll(other.mapBuilder.build());
+            return this;
+        }
+
+        public Builder addAll(Collection<? extends Attribute> c) {
+            for (var e : c) {
+                mapBuilder.put(e, PRESENT);
+            }
+            return this;
+        }
+
+        public boolean removeIf(Predicate<? super Attribute> filter) {
+            return mapBuilder.keySet().removeIf(filter);
+        }
+
+        public boolean contains(Object o) {
+            return mapBuilder.containsKey(o);
+        }
+
+        public boolean isEmpty() {
+            return mapBuilder.isEmpty();
+        }
+
+        public AttributeSet build() {
+            return new AttributeSet(mapBuilder.build());
+        }
+    }
 }

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

@@ -40,11 +40,11 @@ public static AttributeMap<Expression> asAttributeMap(List<? extends NamedExpres
             return AttributeMap.emptyAttributeMap();
         }
 
-        AttributeMap<Expression> map = new AttributeMap<>();
+        AttributeMap.Builder<Expression> mapBuilder = AttributeMap.builder();
         for (NamedExpression exp : named) {
-            map.add(exp.toAttribute(), exp);
+            mapBuilder.put(exp.toAttribute(), exp);
         }
-        return map;
+        return mapBuilder.build();
     }
 
     public static boolean anyMatch(List<? extends Expression> exps, Predicate<? super Expression> predicate) {
@@ -121,11 +121,11 @@ public static AttributeSet references(List<? extends Expression> exps) {
             return AttributeSet.EMPTY;
         }
 
-        AttributeSet set = new AttributeSet();
+        var setBuilder = AttributeSet.builder();
         for (Expression exp : exps) {
-            set.addAll(exp.references());
+            setBuilder.addAll(exp.references());
         }
-        return set;
+        return setBuilder.build();
     }
 
     public static String name(Expression e) {