QL: Make AttributeMap/Set mutable (#98641)

Currently both AttributeSet/Map are immutable requiring building to
 happen (in case of Map) through a dedicated builder.
We have a number of rules where this forces creation of a lot of clones
 due to constant addition, across nodes (and thus not in one place) of
 more attributes.

The combine/intersect/disjunct methods remain as is and return new
 instances however add/remove/removeIf methods should be properly
 implemented.
This will keep the existing code working while allowing a different
 pattern - that in which the backing collection is modified up to date.

Fix #98482

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -613,7 +613,7 @@ public class LogicalPlanOptimizer extends RuleExecutor<LogicalPlan> {
 
         @Override
         public LogicalPlan apply(LogicalPlan plan) {
-            var used = new Holder<>(new AttributeSet());
+            var used = new AttributeSet();
             // don't remove Evals without any Project/Aggregate (which might not occur as the last node in the plan)
             var seenProjection = new Holder<>(Boolean.FALSE);
 
@@ -626,14 +626,13 @@ public class LogicalPlanOptimizer extends RuleExecutor<LogicalPlan> {
                 }
 
                 // remember used
-                var usedSet = used.get();
                 boolean recheck;
                 // analyze the unused items against dedicated 'producer' nodes such as Eval and Aggregate
                 // perform a loop to retry checking if the current node is completely eliminated
                 do {
                     recheck = false;
                     if (p instanceof Aggregate aggregate) {
-                        var remaining = seenProjection.get() ? removeUnused(aggregate.aggregates(), usedSet) : null;
+                        var remaining = seenProjection.get() ? removeUnused(aggregate.aggregates(), used) : null;
                         // no aggregates, no need
                         if (remaining != null) {
                             if (remaining.isEmpty()) {
@@ -646,7 +645,7 @@ public class LogicalPlanOptimizer extends RuleExecutor<LogicalPlan> {
 
                         seenProjection.set(Boolean.TRUE);
                     } else if (p instanceof Eval eval) {
-                        var remaining = seenProjection.get() ? removeUnused(eval.fields(), usedSet) : null;
+                        var remaining = seenProjection.get() ? removeUnused(eval.fields(), used) : null;
                         // no fields, no eval
                         if (remaining != null) {
                             if (remaining.isEmpty()) {
@@ -661,8 +660,7 @@ public class LogicalPlanOptimizer extends RuleExecutor<LogicalPlan> {
                     }
                 } while (recheck);
 
-                var inUse = usedSet.combine(references(p));
-                used.set(inUse);
+                used.addAll(references(p));
 
                 // preserve the state before going to the next node
                 return p;
@@ -685,7 +683,7 @@ public class LogicalPlanOptimizer extends RuleExecutor<LogicalPlan> {
                 if (used.contains(prev.toAttribute()) == false) {
                     it.remove();
                 } else {
-                    used = used.combine(prev.references());
+                    used.addAll(prev.references());
                 }
             }
             return clone.size() != named.size() ? clone : null;

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

@@ -18,10 +18,20 @@ import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.stream.Stream;
 
-import static java.util.Collections.singletonMap;
-import static java.util.Collections.unmodifiableCollection;
-import static java.util.Collections.unmodifiableSet;
-
+import static java.util.Collections.emptyMap;
+
+/**
+ * Dedicated map for checking {@link Attribute} equality.
+ * This is typically the case when comparing the initial declaration of an Attribute, such as {@link FieldAttribute} with
+ * references to it, namely {@link ReferenceAttribute}.
+ * Using plain object equality, the two references are difference due to their type however semantically, they are the same.
+ * Expressions support semantic equality through {@link Expression#semanticEquals(Expression)} - this map is dedicated solution
+ * for attributes as its common case picked up by the plan rules.
+ * <p>
+ * The map implementation is mutable thus consumers need to be careful NOT to modify the content unless they have ownership.
+ * Worth noting the {@link #combine(AttributeMap)}, {@link #intersect(AttributeMap)} and {@link #subtract(AttributeMap)} methods which
+ * return copies, decoupled from the input maps. In other words the returned maps can be modified without affecting the input or vice-versa.
+ */
 public class AttributeMap<E> implements Map<Attribute, E> {
 
     static class AttributeWrapper {
@@ -39,11 +49,7 @@ public class AttributeMap<E> implements Map<Attribute, E> {
 
         @Override
         public boolean equals(Object obj) {
-            if (obj instanceof AttributeWrapper aw) {
-                return attr.semanticEquals(aw.attr);
-            }
-
-            return false;
+            return obj instanceof AttributeWrapper aw ? attr.semanticEquals(aw.attr) : false;
         }
 
         @Override
@@ -59,7 +65,7 @@ public class AttributeMap<E> implements Map<Attribute, E> {
         private final Set<W> set;
 
         UnwrappingSet(Set<W> originalSet) {
-            set = unmodifiableSet(originalSet);
+            set = originalSet;
         }
 
         @Override
@@ -76,6 +82,11 @@ public class AttributeMap<E> implements Map<Attribute, E> {
                 public U next() {
                     return unwrap(i.next());
                 }
+
+                @Override
+                public void remove() {
+                    i.remove();
+                }
             };
         }
 
@@ -141,35 +152,26 @@ public class AttributeMap<E> implements Map<Attribute, E> {
     }
 
     @SuppressWarnings("rawtypes")
-    private static final AttributeMap EMPTY = new AttributeMap<>();
+    private static final AttributeMap EMPTY = new AttributeMap<>(emptyMap());
 
     @SuppressWarnings("unchecked")
     public static final <E> AttributeMap<E> emptyAttributeMap() {
         return EMPTY;
     }
 
-    private static final Object NOT_FOUND = new Object();
-
     private final Map<AttributeWrapper, E> delegate;
-    private Set<Attribute> keySet = null;
-    private Collection<E> values = null;
-    private Set<Entry<Attribute, E>> entrySet = null;
+
+    private AttributeMap(Map<AttributeWrapper, E> other) {
+        delegate = other;
+    }
 
     public AttributeMap() {
         delegate = new LinkedHashMap<>();
     }
 
     public AttributeMap(Attribute key, E value) {
-        delegate = singletonMap(new AttributeWrapper(key), value);
-    }
-
-    void add(Attribute key, E value) {
-        delegate.put(new AttributeWrapper(key), value);
-    }
-
-    // a set from a collection of sets without (too much) copying
-    void addAll(AttributeMap<E> other) {
-        delegate.putAll(other.delegate);
+        delegate = new LinkedHashMap<>();
+        add(key, value);
     }
 
     public AttributeMap<E> combine(AttributeMap<E> other) {
@@ -218,6 +220,14 @@ public class AttributeMap<E> implements Map<Attribute, E> {
         return true;
     }
 
+    public void add(Attribute key, E value) {
+        put(key, value);
+    }
+
+    public void addAll(AttributeMap<E> other) {
+        putAll(other);
+    }
+
     public Set<String> attributeNames() {
         Set<String> s = Sets.newLinkedHashSetWithExpectedSize(size());
 
@@ -239,10 +249,7 @@ public class AttributeMap<E> implements Map<Attribute, E> {
 
     @Override
     public boolean containsKey(Object key) {
-        if (key instanceof NamedExpression) {
-            return delegate.containsKey(new AttributeWrapper(((NamedExpression) key).toAttribute()));
-        }
-        return false;
+        return key instanceof NamedExpression ne ? delegate.containsKey(new AttributeWrapper(ne.toAttribute())) : false;
     }
 
     @Override
@@ -252,18 +259,14 @@ public class AttributeMap<E> implements Map<Attribute, E> {
 
     @Override
     public E get(Object key) {
-        if (key instanceof NamedExpression) {
-            return delegate.get(new AttributeWrapper(((NamedExpression) key).toAttribute()));
-        }
-        return null;
+        return key instanceof NamedExpression ne ? delegate.get(new AttributeWrapper(ne.toAttribute())) : null;
     }
 
     @Override
     public E getOrDefault(Object key, E defaultValue) {
-        if (key instanceof NamedExpression) {
-            return delegate.getOrDefault(new AttributeWrapper(((NamedExpression) key).toAttribute()), defaultValue);
-        }
-        return defaultValue;
+        return key instanceof NamedExpression ne
+            ? delegate.getOrDefault(new AttributeWrapper(ne.toAttribute()), defaultValue)
+            : defaultValue;
     }
 
     public E resolve(Object key) {
@@ -290,71 +293,64 @@ public class AttributeMap<E> implements Map<Attribute, E> {
 
     @Override
     public E put(Attribute key, E value) {
-        throw new UnsupportedOperationException();
+        return delegate.put(new AttributeWrapper(key), value);
     }
 
     @Override
-    public E remove(Object key) {
-        throw new UnsupportedOperationException();
+    public void putAll(Map<? extends Attribute, ? extends E> m) {
+        for (Entry<? extends Attribute, ? extends E> entry : m.entrySet()) {
+            put(entry.getKey(), entry.getValue());
+        }
     }
 
     @Override
-    public void putAll(Map<? extends Attribute, ? extends E> m) {
-        throw new UnsupportedOperationException();
+    public E remove(Object key) {
+        return key instanceof NamedExpression ne ? delegate.remove(new AttributeWrapper(ne.toAttribute())) : null;
     }
 
     @Override
     public void clear() {
-        throw new UnsupportedOperationException();
+        delegate.clear();
     }
 
     @Override
     public Set<Attribute> keySet() {
-        if (keySet == null) {
-            keySet = new UnwrappingSet<>(delegate.keySet()) {
-                @Override
-                protected Attribute unwrap(AttributeWrapper next) {
-                    return next.attr;
-                }
-            };
-        }
-        return keySet;
+        return new UnwrappingSet<>(delegate.keySet()) {
+            @Override
+            protected Attribute unwrap(AttributeWrapper next) {
+                return next.attr;
+            }
+        };
     }
 
     @Override
     public Collection<E> values() {
-        if (values == null) {
-            values = unmodifiableCollection(delegate.values());
-        }
-        return values;
+        return delegate.values();
     }
 
     @Override
     public Set<Entry<Attribute, E>> entrySet() {
-        if (entrySet == null) {
-            entrySet = new UnwrappingSet<>(delegate.entrySet()) {
-                @Override
-                protected Entry<Attribute, E> unwrap(final Entry<AttributeWrapper, E> next) {
-                    return new Entry<>() {
-                        @Override
-                        public Attribute getKey() {
-                            return next.getKey().attr;
-                        }
-
-                        @Override
-                        public E getValue() {
-                            return next.getValue();
-                        }
-
-                        @Override
-                        public E setValue(E value) {
-                            throw new UnsupportedOperationException();
-                        }
-                    };
-                }
-            };
-        }
-        return entrySet;
+        return new UnwrappingSet<>(delegate.entrySet()) {
+            @Override
+            protected Entry<Attribute, E> unwrap(final Entry<AttributeWrapper, E> next) {
+                return new Entry<>() {
+                    @Override
+                    public Attribute getKey() {
+                        return next.getKey().attr;
+                    }
+
+                    @Override
+                    public E getValue() {
+                        return next.getValue();
+                    }
+
+                    @Override
+                    public E setValue(E value) {
+                        return next.setValue(value);
+                    }
+                };
+            }
+        };
     }
 
     @Override
@@ -369,8 +365,8 @@ public class AttributeMap<E> implements Map<Attribute, E> {
 
     @Override
     public boolean equals(Object obj) {
-        if (obj instanceof AttributeMap<?>) {
-            obj = ((AttributeMap<?>) obj).delegate;
+        if (obj instanceof AttributeMap<?> am) {
+            obj = am.delegate;
         }
         return delegate.equals(obj);
     }

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

@@ -14,11 +14,12 @@ import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
+/**
+ * Set variant of {@link AttributeMap} - please see that class Javadoc.
+ */
 public class AttributeSet implements Set<Attribute> {
 
-    private static final AttributeMap<Object> EMPTY_DELEGATE = AttributeMap.emptyAttributeMap();
-
-    public static final AttributeSet EMPTY = new AttributeSet(EMPTY_DELEGATE);
+    public static final AttributeSet EMPTY = new AttributeSet(AttributeMap.emptyAttributeMap());
 
     // use the same name as in HashSet
     private static final Object PRESENT = new Object();
@@ -34,14 +35,10 @@ public class AttributeSet implements Set<Attribute> {
     }
 
     public AttributeSet(Collection<? extends Attribute> attr) {
-        if (attr.isEmpty()) {
-            delegate = EMPTY_DELEGATE;
-        } else {
-            delegate = new AttributeMap<>();
+        delegate = new AttributeMap<>();
 
-            for (Attribute a : attr) {
-                delegate.add(a, PRESENT);
-            }
+        for (Attribute a : attr) {
+            delegate.add(a, PRESENT);
         }
     }
 
@@ -49,12 +46,6 @@ public class AttributeSet implements Set<Attribute> {
         this.delegate = delegate;
     }
 
-    // package protected - should be called through Expressions to cheaply create
-    // a set from a collection of sets without too much copying
-    void addAll(AttributeSet other) {
-        delegate.addAll(other.delegate);
-    }
-
     public AttributeSet combine(AttributeSet other) {
         return new AttributeSet(delegate.combine(other.delegate));
     }
@@ -122,42 +113,54 @@ public class AttributeSet implements Set<Attribute> {
 
     @Override
     public boolean add(Attribute e) {
-        throw new UnsupportedOperationException();
+        return delegate.put(e, PRESENT) != null;
     }
 
     @Override
     public boolean remove(Object o) {
-        throw new UnsupportedOperationException();
+        return delegate.remove(o) != null;
+    }
+
+    public void addAll(AttributeSet other) {
+        delegate.addAll(other.delegate);
     }
 
     @Override
     public boolean addAll(Collection<? extends Attribute> c) {
-        throw new UnsupportedOperationException();
+        int size = delegate.size();
+        for (var e : c) {
+            delegate.put(e, PRESENT);
+        }
+        return delegate.size() != size;
     }
 
     @Override
     public boolean retainAll(Collection<?> c) {
-        throw new UnsupportedOperationException();
+        return delegate.keySet().removeIf(e -> c.contains(e) == false);
     }
 
     @Override
     public boolean removeAll(Collection<?> c) {
-        throw new UnsupportedOperationException();
+        int size = delegate.size();
+        for (var e : c) {
+            delegate.remove(e);
+        }
+        return delegate.size() != size;
     }
 
     @Override
     public void clear() {
-        throw new UnsupportedOperationException();
+        delegate.clear();
     }
 
     @Override
     public Spliterator<Attribute> spliterator() {
-        throw new UnsupportedOperationException();
+        return delegate.keySet().spliterator();
     }
 
     @Override
     public boolean removeIf(Predicate<? super Attribute> filter) {
-        throw new UnsupportedOperationException();
+        return delegate.keySet().removeIf(filter);
     }
 
     @Override

x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java

@@ -17,6 +17,7 @@ import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
 
+import static java.util.Arrays.asList;
 import static java.util.stream.Collectors.toList;
 import static org.elasticsearch.xpack.ql.TestUtils.fieldAttribute;
 import static org.elasticsearch.xpack.ql.TestUtils.of;
@@ -25,6 +26,7 @@ import static org.hamcrest.Matchers.arrayWithSize;
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.sameInstance;
 
 public class AttributeMapTests extends ESTestCase {
 
@@ -249,4 +251,83 @@ public class AttributeMapTests extends ESTestCase {
 
         assertThat(m, is(copy));
     }
+
+    public void testEmptyMapIsImmutable() {
+        var empty = AttributeMap.emptyAttributeMap();
+        var ex = expectThrows(UnsupportedOperationException.class, () -> empty.add(a("one"), new Object()));
+    }
+
+    public void testAddPutEntriesIntoMap() {
+        var map = new AttributeMap<String>();
+        var one = a("one");
+        var two = a("two");
+        var three = a("three");
+
+        for (var i : asList(one, two, three)) {
+            map.add(i, i.name());
+        }
+
+        assertThat(map.size(), is(3));
+
+        assertThat(map.remove(one), is("one"));
+        assertThat(map.remove(two), is("two"));
+
+        assertThat(map.size(), is(1));
+    }
+
+    public void testKeyIteratorRemoval() {
+        var map = new AttributeMap<String>();
+        var one = a("one");
+        var two = a("two");
+        var three = a("three");
+
+        for (var i : asList(one, two, three)) {
+            map.add(i, i.name());
+        }
+
+        assertThat(map.attributeNames(), contains("one", "two", "three"));
+        assertThat(map.size(), is(3));
+
+        var it = map.keySet().iterator();
+        var next = it.next();
+        assertThat(next, sameInstance(one));
+        it.remove();
+        assertThat(map.size(), is(2));
+        next = it.next();
+        assertThat(next, sameInstance(two));
+        next = it.next();
+
+        assertThat(next, sameInstance(three));
+        it.remove();
+        assertThat(map.size(), is(1));
+
+        assertThat(it.hasNext(), is(false));
+    }
+
+    public void testValuesIteratorRemoval() {
+        var map = new AttributeMap<String>();
+        var one = a("one");
+        var two = a("two");
+        var three = a("three");
+
+        for (var i : asList(one, two, three)) {
+            map.add(i, i.name());
+        }
+
+        assertThat(map.values(), contains("one", "two", "three"));
+
+        map.values().removeIf(v -> v.contains("o"));
+        assertThat(map.size(), is(1));
+        assertThat(map.containsKey(three), is(true));
+        assertThat(map.containsValue("three"), is(true));
+
+        assertThat(map.containsKey("two"), is(false));
+        assertThat(map.containsKey(one), is(false));
+
+        var it = map.values().iterator();
+        assertThat(it.hasNext(), is(true));
+        assertThat(it.next(), is("three"));
+        it.remove();
+        assertThat(it.hasNext(), is(false));
+    }
 }