Skip to content

Commit af471fe

Browse files
bpinteabpintea
authored
[9.0] ESQL: make AttributeMap and AttributeSet immutable (#125938) (#126207)
* ESQL: make `AttributeMap` and `AttributeSet` immutable (#125938) This will allow reusing them in the plan analysis and skip recreating them in UnaryPlan/UnaryExec when not needed. Introduce/adjust builders for them, which are now the only way to use a modifiable map/set. Related #124395 (cherry picked from commit 2b512bc) * Update 9.0-specific AttributeSet usage
1 parent e5bd814 commit af471fe

32 files changed

+293
-175
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ public Nullability nullable() {
6666

6767
@Override
6868
public AttributeSet references() {
69-
return new AttributeSet(this);
69+
return AttributeSet.of(this);
7070
}
7171

7272
public Attribute withLocation(Source source) {

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

Lines changed: 51 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import java.util.Map;
1717
import java.util.Set;
1818
import java.util.function.BiConsumer;
19+
import java.util.function.Function;
1920
import java.util.stream.Stream;
2021

2122
import static java.util.Collections.emptyMap;
@@ -34,7 +35,7 @@
3435
*/
3536
public final class AttributeMap<E> implements Map<Attribute, E> {
3637

37-
static class AttributeWrapper {
38+
private static class AttributeWrapper {
3839

3940
private final Attribute attr;
4041

@@ -165,7 +166,7 @@ private AttributeMap(Map<AttributeWrapper, E> other) {
165166
delegate = other;
166167
}
167168

168-
public AttributeMap() {
169+
private AttributeMap() {
169170
delegate = new LinkedHashMap<>();
170171
}
171172

@@ -220,12 +221,22 @@ public boolean subsetOf(AttributeMap<E> other) {
220221
return true;
221222
}
222223

223-
public void add(Attribute key, E value) {
224-
put(key, value);
224+
private E add(Attribute key, E value) {
225+
return delegate.put(new AttributeWrapper(key), value);
226+
}
227+
228+
private void addAll(AttributeMap<E> other) {
229+
for (Entry<? extends Attribute, ? extends E> entry : other.entrySet()) {
230+
add(entry.getKey(), entry.getValue());
231+
}
225232
}
226233

227-
public void addAll(AttributeMap<E> other) {
228-
putAll(other);
234+
private E addIfAbsent(Attribute key, Function<? super Attribute, ? extends E> mappingFunction) {
235+
return delegate.computeIfAbsent(new AttributeWrapper(key), k -> mappingFunction.apply(k.attr));
236+
}
237+
238+
private E delete(Object key) {
239+
return key instanceof NamedExpression ne ? delegate.remove(new AttributeWrapper(ne.toAttribute())) : null;
229240
}
230241

231242
public Set<String> attributeNames() {
@@ -249,7 +260,7 @@ public boolean isEmpty() {
249260

250261
@Override
251262
public boolean containsKey(Object key) {
252-
return key instanceof NamedExpression ne ? delegate.containsKey(new AttributeWrapper(ne.toAttribute())) : false;
263+
return key instanceof NamedExpression ne && delegate.containsKey(new AttributeWrapper(ne.toAttribute()));
253264
}
254265

255266
@Override
@@ -293,24 +304,22 @@ public E resolve(Object key, E defaultValue) {
293304

294305
@Override
295306
public E put(Attribute key, E value) {
296-
return delegate.put(new AttributeWrapper(key), value);
307+
throw new UnsupportedOperationException();
297308
}
298309

299310
@Override
300311
public void putAll(Map<? extends Attribute, ? extends E> m) {
301-
for (Entry<? extends Attribute, ? extends E> entry : m.entrySet()) {
302-
put(entry.getKey(), entry.getValue());
303-
}
312+
throw new UnsupportedOperationException();
304313
}
305314

306315
@Override
307316
public E remove(Object key) {
308-
return key instanceof NamedExpression ne ? delegate.remove(new AttributeWrapper(ne.toAttribute())) : null;
317+
throw new UnsupportedOperationException();
309318
}
310319

311320
@Override
312321
public void clear() {
313-
delegate.clear();
322+
throw new UnsupportedOperationException();
314323
}
315324

316325
@Override
@@ -376,29 +385,50 @@ public String toString() {
376385
return delegate.toString();
377386
}
378387

379-
public static <E> Builder<E> builder() {
380-
return new Builder<>();
388+
public static <E> AttributeMap<E> of(Attribute key, E value) {
389+
final AttributeMap<E> map = new AttributeMap<>();
390+
map.add(key, value);
391+
return map;
381392
}
382393

383-
public static <E> Builder<E> builder(AttributeMap<E> map) {
384-
return new Builder<E>().putAll(map);
394+
public static <E> Builder<E> builder() {
395+
return new Builder<>();
385396
}
386397

387398
public static class Builder<E> {
388-
private AttributeMap<E> map = new AttributeMap<>();
399+
private final AttributeMap<E> map = new AttributeMap<>();
389400

390401
private Builder() {}
391402

392-
public Builder<E> put(Attribute attr, E value) {
393-
map.add(attr, value);
394-
return this;
403+
public E put(Attribute attr, E value) {
404+
return map.add(attr, value);
395405
}
396406

397407
public Builder<E> putAll(AttributeMap<E> m) {
398408
map.addAll(m);
399409
return this;
400410
}
401411

412+
public E computeIfAbsent(Attribute key, Function<? super Attribute, ? extends E> mappingFunction) {
413+
return map.addIfAbsent(key, mappingFunction);
414+
}
415+
416+
public E remove(Object o) {
417+
return map.delete(o);
418+
}
419+
420+
public Set<Attribute> keySet() {
421+
return map.keySet();
422+
}
423+
424+
public boolean containsKey(Object key) {
425+
return map.containsKey(key);
426+
}
427+
428+
public boolean isEmpty() {
429+
return map.isEmpty();
430+
}
431+
402432
public AttributeMap<E> build() {
403433
return map;
404434
}

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

Lines changed: 79 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -26,22 +26,6 @@ public class AttributeSet implements Set<Attribute> {
2626

2727
private final AttributeMap<Object> delegate;
2828

29-
public AttributeSet() {
30-
delegate = new AttributeMap<>();
31-
}
32-
33-
public AttributeSet(Attribute attr) {
34-
delegate = new AttributeMap<>(attr, PRESENT);
35-
}
36-
37-
public AttributeSet(Collection<? extends Attribute> attr) {
38-
delegate = new AttributeMap<>();
39-
40-
for (Attribute a : attr) {
41-
delegate.add(a, PRESENT);
42-
}
43-
}
44-
4529
private AttributeSet(AttributeMap<Object> delegate) {
4630
this.delegate = delegate;
4731
}
@@ -113,44 +97,32 @@ public <T> T[] toArray(T[] a) {
11397

11498
@Override
11599
public boolean add(Attribute e) {
116-
return delegate.put(e, PRESENT) == null;
100+
throw new UnsupportedOperationException();
117101
}
118102

119103
@Override
120104
public boolean remove(Object o) {
121-
return delegate.remove(o) != null;
122-
}
123-
124-
public void addAll(AttributeSet other) {
125-
delegate.addAll(other.delegate);
105+
throw new UnsupportedOperationException();
126106
}
127107

128108
@Override
129109
public boolean addAll(Collection<? extends Attribute> c) {
130-
int size = delegate.size();
131-
for (var e : c) {
132-
delegate.put(e, PRESENT);
133-
}
134-
return delegate.size() != size;
110+
throw new UnsupportedOperationException();
135111
}
136112

137113
@Override
138114
public boolean retainAll(Collection<?> c) {
139-
return delegate.keySet().removeIf(e -> c.contains(e) == false);
115+
throw new UnsupportedOperationException();
140116
}
141117

142118
@Override
143119
public boolean removeAll(Collection<?> c) {
144-
int size = delegate.size();
145-
for (var e : c) {
146-
delegate.remove(e);
147-
}
148-
return delegate.size() != size;
120+
throw new UnsupportedOperationException();
149121
}
150122

151123
@Override
152124
public void clear() {
153-
delegate.clear();
125+
throw new UnsupportedOperationException();
154126
}
155127

156128
@Override
@@ -160,7 +132,7 @@ public Spliterator<Attribute> spliterator() {
160132

161133
@Override
162134
public boolean removeIf(Predicate<? super Attribute> filter) {
163-
return delegate.keySet().removeIf(filter);
135+
throw new UnsupportedOperationException();
164136
}
165137

166138
@Override
@@ -191,4 +163,76 @@ public int hashCode() {
191163
public String toString() {
192164
return delegate.keySet().toString();
193165
}
166+
167+
public Builder asBuilder() {
168+
return new Builder().addAll(this);
169+
}
170+
171+
public static AttributeSet of(Attribute... attrs) {
172+
final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
173+
for (var a : attrs) {
174+
mapBuilder.put(a, PRESENT);
175+
}
176+
return new AttributeSet(mapBuilder.build());
177+
}
178+
179+
public static AttributeSet of(Collection<? extends Attribute> c) {
180+
final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
181+
for (var a : c) {
182+
mapBuilder.put(a, PRESENT);
183+
}
184+
return new AttributeSet(mapBuilder.build());
185+
}
186+
187+
public static Builder builder() {
188+
return new Builder();
189+
}
190+
191+
public static class Builder {
192+
private final AttributeMap.Builder<Object> mapBuilder = AttributeMap.builder();
193+
194+
private Builder() {}
195+
196+
public Builder add(Attribute attr) {
197+
mapBuilder.put(attr, PRESENT);
198+
return this;
199+
}
200+
201+
public boolean remove(Object o) {
202+
return mapBuilder.remove(o) != null;
203+
}
204+
205+
public Builder addAll(AttributeSet other) {
206+
mapBuilder.putAll(other.delegate);
207+
return this;
208+
}
209+
210+
public Builder addAll(Builder other) {
211+
mapBuilder.putAll(other.mapBuilder.build());
212+
return this;
213+
}
214+
215+
public Builder addAll(Collection<? extends Attribute> c) {
216+
for (var e : c) {
217+
mapBuilder.put(e, PRESENT);
218+
}
219+
return this;
220+
}
221+
222+
public boolean removeIf(Predicate<? super Attribute> filter) {
223+
return mapBuilder.keySet().removeIf(filter);
224+
}
225+
226+
public boolean contains(Object o) {
227+
return mapBuilder.containsKey(o);
228+
}
229+
230+
public boolean isEmpty() {
231+
return mapBuilder.isEmpty();
232+
}
233+
234+
public AttributeSet build() {
235+
return new AttributeSet(mapBuilder.build());
236+
}
237+
}
194238
}

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,11 +40,11 @@ public static AttributeMap<Expression> asAttributeMap(List<? extends NamedExpres
4040
return AttributeMap.emptyAttributeMap();
4141
}
4242

43-
AttributeMap<Expression> map = new AttributeMap<>();
43+
AttributeMap.Builder<Expression> mapBuilder = AttributeMap.builder();
4444
for (NamedExpression exp : named) {
45-
map.add(exp.toAttribute(), exp);
45+
mapBuilder.put(exp.toAttribute(), exp);
4646
}
47-
return map;
47+
return mapBuilder.build();
4848
}
4949

5050
public static boolean anyMatch(List<? extends Expression> exps, Predicate<? super Expression> predicate) {
@@ -121,11 +121,11 @@ public static AttributeSet references(List<? extends Expression> exps) {
121121
return AttributeSet.EMPTY;
122122
}
123123

124-
AttributeSet set = new AttributeSet();
124+
var setBuilder = AttributeSet.builder();
125125
for (Expression exp : exps) {
126-
set.addAll(exp.references());
126+
setBuilder.addAll(exp.references());
127127
}
128-
return set;
128+
return setBuilder.build();
129129
}
130130

131131
public static String name(Expression e) {

0 commit comments

Comments
 (0)