ESQL: Lazy collection copying during node transform (#124424) (#124528)

* 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

Author: costin

docs/changelog/124424.yaml

@@ -0,0 +1,5 @@
+pr: 124424
+summary: Lazy collection copying during node transform
+area: ES|QL
+type: bug
+issues: []

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

@@ -12,7 +12,6 @@
 import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
 
 import java.io.IOException;
-import java.util.Objects;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -34,7 +33,7 @@ public NameId() {
 
  @Override
  public int hashCode() {
- return Objects.hash(id);
+ return Long.hashCode(id);
  }
 
  @Override

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

@@ -110,7 +110,7 @@ public <E> void forEachPropertyUp(Class<E> typeToken, Consumer<? super E> rule)
  protected <E> void forEachProperty(Class<E> typeToken, Consumer<? super E> rule) {
  for (Object prop : info().properties()) {
  // skip children (only properties are interesting)
- if (prop != children && children.contains(prop) == false && typeToken.isInstance(prop)) {
+ if (prop != children && typeToken.isInstance(prop) && children.contains(prop) == false) {
  rule.accept((E) prop);
  }
  }
@@ -203,20 +203,21 @@ public <E extends T> T transformUp(Class<E> typeToken, Function<E, ? extends T>
  protected <R extends Function<? super T, ? extends T>> T transformChildren(Function<T, ? extends T> traversalOperation) {
  boolean childrenChanged = false;
 
- // stream() could be used but the code is just as complicated without any advantages
- // further more, it would include bring in all the associated stream/collector object creation even though in
- // most cases the immediate tree would be quite small (0,1,2 elements)
- List<T> transformedChildren = new ArrayList<>(children().size());
+ // Avoid creating a new array of children if no change is needed.
+ // And when it happens, look at using replacement to minimize the amount of method invocations.
+ List<T> transformedChildren = null;
 
- for (T child : children) {
+ for (int i = 0, s = children.size(); i < s; i++) {
+ T child = children.get(i);
  T next = traversalOperation.apply(child);
- if (child.equals(next)) {
- // use the initial value
- next = child;
- } else {
- childrenChanged = true;
+ if (child.equals(next) == false) {
+ // lazy copy + replacement in place
+ if (childrenChanged == false) {
+ childrenChanged = true;
+ transformedChildren = new ArrayList<>(children);
+ }
+ transformedChildren.set(i, next);
  }
- transformedChildren.add(next);
  }
 
  return (childrenChanged ? replaceChildrenSameSize(transformedChildren) : (T) this);

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

@@ -52,7 +52,7 @@ final <E> T transform(Function<? super E, ? extends E> rule, Class<E> typeToken)
  List<?> children = node.children();
 
  Function<Object, Object> realRule = p -> {
- if (p != children && false == children.contains(p) && (p == null || typeToken.isInstance(p))) {
+ if (p != children && (p == null || typeToken.isInstance(p)) && false == children.contains(p)) {
  return rule.apply(typeToken.cast(p));
  }
  return p;

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

@@ -131,8 +131,8 @@ public <E extends Expression> PlanType transformExpressionsUp(Class<E> typeToken
 
  @SuppressWarnings("unchecked")
  private static Object doTransformExpression(Object arg, Function<Expression, ? extends Expression> traversal) {
- if (arg instanceof Expression) {
- return traversal.apply((Expression) arg);
+ if (arg instanceof Expression exp) {
+ return traversal.apply(exp);
  }
 
  // WARNING: if the collection is typed, an incompatible function will be applied to it
@@ -141,17 +141,19 @@ private static Object doTransformExpression(Object arg, Function<Expression, ? e
  // has no type info so it's difficult to have automatic checking without having base classes).
 
  if (arg instanceof Collection<?> c) {
- List<Object> transformed = new ArrayList<>(c.size());
+ List<Object> transformed = null;
  boolean hasChanged = false;
+ int i = 0;
  for (Object e : c) {
  Object next = doTransformExpression(e, traversal);
- if (e.equals(next)) {
- // use the initial value
- next = e;
- } else {
- hasChanged = true;
+ if (e.equals(next) == false) {
+ if (hasChanged == false) {
+ hasChanged = true;
+ transformed = new ArrayList<>(c);
+ }
+ transformed.set(i, next);
  }
- transformed.add(next);
+ i++;
  }
 
  return hasChanged ? transformed : arg;
@@ -186,8 +188,8 @@ public <E extends Expression> void forEachExpressionUp(Class<E> typeToken, Consu
 
  @SuppressWarnings("unchecked")
  private static void doForEachExpression(Object arg, Consumer<Expression> traversal) {
- if (arg instanceof Expression) {
- traversal.accept((Expression) arg);
+ if (arg instanceof Expression exp) {
+ traversal.accept(exp);
  } else if (arg instanceof Collection<?> c) {
  for (Object o : c) {
  doForEachExpression(o, traversal);

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

@@ -6,7 +6,6 @@
  */
 package org.elasticsearch.xpack.ql.expression;
 
-import java.util.Objects;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -28,7 +27,7 @@ public NameId() {
 
  @Override
  public int hashCode() {
- return Objects.hash(id);
+ return Long.hashCode(id);
  }
 
  @Override

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

@@ -109,8 +109,8 @@ public <E extends Expression> PlanType transformExpressionsUp(Class<E> typeToken
 
  @SuppressWarnings("unchecked")
  private static Object doTransformExpression(Object arg, Function<Expression, ? extends Expression> traversal) {
- if (arg instanceof Expression) {
- return traversal.apply((Expression) arg);
+ if (arg instanceof Expression exp) {
+ return traversal.apply(exp);
  }
 
  // 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
  // has no type info so it's difficult to have automatic checking without having base classes).
 
  if (arg instanceof Collection<?> c) {
- List<Object> transformed = new ArrayList<>(c.size());
+ List<Object> transformed = null;
  boolean hasChanged = false;
+ int i = 0;
  for (Object e : c) {
  Object next = doTransformExpression(e, traversal);
- if (e.equals(next)) {
- // use the initial value
- next = e;
- } else {
- hasChanged = true;
+ if (e.equals(next) == false) {
+ if (hasChanged == false) {
+ hasChanged = true;
+ transformed = new ArrayList<>(c);
+ }
+ transformed.set(i, next);
  }
- transformed.add(next);
+ i++;
  }
 
  return hasChanged ? transformed : arg;
@@ -164,8 +166,8 @@ public <E extends Expression> void forEachExpressionUp(Class<E> typeToken, Consu
 
  @SuppressWarnings("unchecked")
  private static void doForEachExpression(Object arg, Consumer<Expression> traversal) {
- if (arg instanceof Expression) {
- traversal.accept((Expression) arg);
+ if (arg instanceof Expression exp) {
+ traversal.accept(exp);
  } else if (arg instanceof Collection<?> c) {
  for (Object o : c) {
  doForEachExpression(o, traversal);

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

@@ -109,7 +109,7 @@ public <E> void forEachPropertyUp(Class<E> typeToken, Consumer<? super E> rule)
  protected <E> void forEachProperty(Class<E> typeToken, Consumer<? super E> rule) {
  for (Object prop : info().properties()) {
  // skip children (only properties are interesting)
- if (prop != children && children.contains(prop) == false && typeToken.isInstance(prop)) {
+ if (prop != children && typeToken.isInstance(prop) && children.contains(prop) == false) {
  rule.accept((E) prop);
  }
  }
@@ -202,20 +202,21 @@ public <E extends T> T transformUp(Class<E> typeToken, Function<E, ? extends T>
  protected <R extends Function<? super T, ? extends T>> T transformChildren(Function<T, ? extends T> traversalOperation) {
  boolean childrenChanged = false;
 
- // stream() could be used but the code is just as complicated without any advantages
- // further more, it would include bring in all the associated stream/collector object creation even though in
- // most cases the immediate tree would be quite small (0,1,2 elements)
- List<T> transformedChildren = new ArrayList<>(children().size());
+ // Avoid creating a new array of children if no change is needed.
+ // And when it happens, look at using replacement to minimize the amount of method invocations.
+ List<T> transformedChildren = null;
 
- for (T child : children) {
+ for (int i = 0, s = children.size(); i < s; i++) {
+ T child = children.get(i);
  T next = traversalOperation.apply(child);
- if (child.equals(next)) {
- // use the initial value
- next = child;
- } else {
- childrenChanged = true;
+ if (child.equals(next) == false) {
+ // lazy copy + replacement in place
+ if (childrenChanged == false) {
+ childrenChanged = true;
+ transformedChildren = new ArrayList<>(children);
+ }
+ transformedChildren.set(i, next);
  }
- transformedChildren.add(next);
  }
 
  return (childrenChanged ? replaceChildrenSameSize(transformedChildren) : (T) this);

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

@@ -52,7 +52,7 @@ final <E> T transform(Function<? super E, ? extends E> rule, Class<E> typeToken)
  List<?> children = node.children();
 
  Function<Object, Object> realRule = p -> {
- if (p != children && false == children.contains(p) && (p == null || typeToken.isInstance(p))) {
+ if (p != children && (p == null || typeToken.isInstance(p)) && false == children.contains(p)) {
  return rule.apply(typeToken.cast(p));
  }
  return p;