Skip to content

Commit cd631ac

Browse files
costincostin
committed
ESQL: Lazy collection copying during node transform
A set of optimization for tree traversal: 1. perform lazy copying during children transform 2. use long hashing to avoid object creation 3. perform type check first before collection checking Relates #124395
1 parent 6f0756a commit cd631ac

File tree

6 files changed

+57
-51
lines changed

6 files changed

+57
-51
lines changed

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
1313

1414
import java.io.IOException;
15-
import java.util.Objects;
1615
import java.util.concurrent.atomic.AtomicLong;
1716

1817
/**
@@ -34,7 +33,7 @@ public NameId() {
3433

3534
@Override
3635
public int hashCode() {
37-
return Objects.hash(id);
36+
return Long.hashCode(id);
3837
}
3938

4039
@Override

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ public <E> void forEachPropertyUp(Class<E> typeToken, Consumer<? super E> rule)
110110
protected <E> void forEachProperty(Class<E> typeToken, Consumer<? super E> rule) {
111111
for (Object prop : info().properties()) {
112112
// skip children (only properties are interesting)
113-
if (prop != children && children.contains(prop) == false && typeToken.isInstance(prop)) {
113+
if (prop != children && typeToken.isInstance(prop) && children.contains(prop) == false) {
114114
rule.accept((E) prop);
115115
}
116116
}
@@ -203,20 +203,21 @@ public <E extends T> T transformUp(Class<E> typeToken, Function<E, ? extends T>
203203
protected <R extends Function<? super T, ? extends T>> T transformChildren(Function<T, ? extends T> traversalOperation) {
204204
boolean childrenChanged = false;
205205

206-
// stream() could be used but the code is just as complicated without any advantages
207-
// further more, it would include bring in all the associated stream/collector object creation even though in
208-
// most cases the immediate tree would be quite small (0,1,2 elements)
209-
List<T> transformedChildren = new ArrayList<>(children().size());
206+
// Avoid creating a new array of children if no change is needed.
207+
// And when it happens, look at using replacement to minimize the amount of method invocations.
208+
List<T> transformedChildren = null;
210209

211-
for (T child : children) {
210+
for (int i = 0, s = children.size(); i < s; i++) {
211+
T child = children.get(i);
212212
T next = traversalOperation.apply(child);
213-
if (child.equals(next)) {
214-
// use the initial value
215-
next = child;
216-
} else {
217-
childrenChanged = true;
213+
if (child.equals(next) == false) {
214+
// lazy copy + replacement in place
215+
if (childrenChanged == false) {
216+
childrenChanged = true;
217+
transformedChildren = new ArrayList<>(children);
218+
}
219+
transformedChildren.set(i, next);
218220
}
219-
transformedChildren.add(next);
220221
}
221222

222223
return (childrenChanged ? replaceChildrenSameSize(transformedChildren) : (T) this);

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

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
import java.util.ArrayList;
1717
import java.util.Collection;
18+
import java.util.Iterator;
1819
import java.util.List;
1920
import java.util.function.Consumer;
2021
import java.util.function.Function;
@@ -131,8 +132,8 @@ public <E extends Expression> PlanType transformExpressionsUp(Class<E> typeToken
131132

132133
@SuppressWarnings("unchecked")
133134
private static Object doTransformExpression(Object arg, Function<Expression, ? extends Expression> traversal) {
134-
if (arg instanceof Expression) {
135-
return traversal.apply((Expression) arg);
135+
if (arg instanceof Expression exp) {
136+
return traversal.apply(exp);
136137
}
137138

138139
// WARNING: if the collection is typed, an incompatible function will be applied to it
@@ -141,17 +142,19 @@ private static Object doTransformExpression(Object arg, Function<Expression, ? e
141142
// has no type info so it's difficult to have automatic checking without having base classes).
142143

143144
if (arg instanceof Collection<?> c) {
144-
List<Object> transformed = new ArrayList<>(c.size());
145+
List<Object> transformed = null;
145146
boolean hasChanged = false;
146-
for (Object e : c) {
147+
int i = 0;
148+
for (Object e: c) {
147149
Object next = doTransformExpression(e, traversal);
148-
if (e.equals(next)) {
149-
// use the initial value
150-
next = e;
151-
} else {
152-
hasChanged = true;
150+
if (e.equals(next) == false) {
151+
if (hasChanged == false) {
152+
hasChanged = true;
153+
transformed = new ArrayList<>(c);
154+
}
155+
transformed.set(i, next);
153156
}
154-
transformed.add(next);
157+
i++;
155158
}
156159

157160
return hasChanged ? transformed : arg;
@@ -186,8 +189,8 @@ public <E extends Expression> void forEachExpressionUp(Class<E> typeToken, Consu
186189

187190
@SuppressWarnings("unchecked")
188191
private static void doForEachExpression(Object arg, Consumer<Expression> traversal) {
189-
if (arg instanceof Expression) {
190-
traversal.accept((Expression) arg);
192+
if (arg instanceof Expression exp) {
193+
traversal.accept(exp);
191194
} else if (arg instanceof Collection<?> c) {
192195
for (Object o : c) {
193196
doForEachExpression(o, traversal);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public NameId() {
2828

2929
@Override
3030
public int hashCode() {
31-
return Objects.hash(id);
31+
return Long.hashCode(id);
3232
}
3333

3434
@Override

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

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,8 @@ public <E extends Expression> PlanType transformExpressionsUp(Class<E> typeToken
109109

110110
@SuppressWarnings("unchecked")
111111
private static Object doTransformExpression(Object arg, Function<Expression, ? extends Expression> traversal) {
112-
if (arg instanceof Expression) {
113-
return traversal.apply((Expression) arg);
112+
if (arg instanceof Expression exp) {
113+
return traversal.apply(exp);
114114
}
115115

116116
// WARNING: if the collection is typed, an incompatible function will be applied to it
@@ -119,17 +119,19 @@ private static Object doTransformExpression(Object arg, Function<Expression, ? e
119119
// has no type info so it's difficult to have automatic checking without having base classes).
120120

121121
if (arg instanceof Collection<?> c) {
122-
List<Object> transformed = new ArrayList<>(c.size());
122+
List<Object> transformed = null;
123123
boolean hasChanged = false;
124-
for (Object e : c) {
124+
int i = 0;
125+
for (Object e: c) {
125126
Object next = doTransformExpression(e, traversal);
126-
if (e.equals(next)) {
127-
// use the initial value
128-
next = e;
129-
} else {
130-
hasChanged = true;
127+
if (e.equals(next) == false) {
128+
if (hasChanged == false) {
129+
hasChanged = true;
130+
transformed = new ArrayList<>(c);
131+
}
132+
transformed.set(i, next);
131133
}
132-
transformed.add(next);
134+
i++;
133135
}
134136

135137
return hasChanged ? transformed : arg;
@@ -164,8 +166,8 @@ public <E extends Expression> void forEachExpressionUp(Class<E> typeToken, Consu
164166

165167
@SuppressWarnings("unchecked")
166168
private static void doForEachExpression(Object arg, Consumer<Expression> traversal) {
167-
if (arg instanceof Expression) {
168-
traversal.accept((Expression) arg);
169+
if (arg instanceof Expression exp) {
170+
traversal.accept(exp);
169171
} else if (arg instanceof Collection<?> c) {
170172
for (Object o : c) {
171173
doForEachExpression(o, traversal);

‎x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/tree/Node.java‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ public <E> void forEachPropertyUp(Class<E> typeToken, Consumer<? super E> rule)
109109
protected <E> void forEachProperty(Class<E> typeToken, Consumer<? super E> rule) {
110110
for (Object prop : info().properties()) {
111111
// skip children (only properties are interesting)
112-
if (prop != children && children.contains(prop) == false && typeToken.isInstance(prop)) {
112+
if (prop != children && typeToken.isInstance(prop) && children.contains(prop) == false) {
113113
rule.accept((E) prop);
114114
}
115115
}
@@ -202,20 +202,21 @@ public <E extends T> T transformUp(Class<E> typeToken, Function<E, ? extends T>
202202
protected <R extends Function<? super T, ? extends T>> T transformChildren(Function<T, ? extends T> traversalOperation) {
203203
boolean childrenChanged = false;
204204

205-
// stream() could be used but the code is just as complicated without any advantages
206-
// further more, it would include bring in all the associated stream/collector object creation even though in
207-
// most cases the immediate tree would be quite small (0,1,2 elements)
208-
List<T> transformedChildren = new ArrayList<>(children().size());
205+
// Avoid creating a new array of children if no change is needed.
206+
// And when it happens, look at using replacement to minimize the amount of method invocations.
207+
List<T> transformedChildren = null;
209208

210-
for (T child : children) {
209+
for (int i = 0, s = children.size(); i < s; i++) {
210+
T child = children.get(i);
211211
T next = traversalOperation.apply(child);
212-
if (child.equals(next)) {
213-
// use the initial value
214-
next = child;
215-
} else {
216-
childrenChanged = true;
212+
if (child.equals(next) == false) {
213+
// lazy copy + replacement in place
214+
if (childrenChanged == false) {
215+
childrenChanged = true;
216+
transformedChildren = new ArrayList<>(children);
217+
}
218+
transformedChildren.set(i, next);
217219
}
218-
transformedChildren.add(next);
219220
}
220221

221222
return (childrenChanged ? replaceChildrenSameSize(transformedChildren) : (T) this);

0 commit comments

Comments
 (0)