QL: Make UnaryPlan.replaceChild public and use it where appropriate (elastic#76071) (elastic#76154)

Co-authored-by: Lukas Wegmann <[email protected]>

Author: elasticsearchmachine

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/optimizer/Optimizer.java

@@ -65,7 +65,6 @@
 import java.util.Objects;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 import static org.elasticsearch.xpack.ql.optimizer.OptimizerRules.PropagateNullable;
 
@@ -448,13 +447,13 @@ protected LogicalPlan rule(OrderBy orderBy) {
                         // preserve the order for the base query, everything else needs to be ascending
                         List<Order> pushedOrder = baseFilter ? orderBy.order() : ascendingOrders;
                         OrderBy order = new OrderBy(filter.source(), filter.child(), pushedOrder);
-                        orderedQueries.add((KeyedFilter) filter.replaceChildrenSameSize(singletonList(order)));
+                        orderedQueries.add(filter.replaceChild(order));
                         baseFilter = false;
                     }
 
                     KeyedFilter until = join.until();
                     OrderBy order = new OrderBy(until.source(), until.child(), ascendingOrders);
-                    until = (KeyedFilter) until.replaceChildrenSameSize(singletonList(order));
+                    until = until.replaceChild(order);
 
                     OrderDirection direction = orderBy.order().get(0).direction();
                     plan = join.with(orderedQueries, until, direction);

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/plan/logical/Head.java

@@ -25,7 +25,7 @@ protected NodeInfo<Limit> info() {
     }
 
     @Override
-    protected Head replaceChild(LogicalPlan newChild) {
+    public Head replaceChild(LogicalPlan newChild) {
         return new Head(source(), limit(), newChild);
     }
 }

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/plan/logical/KeyedFilter.java

@@ -43,7 +43,7 @@ protected NodeInfo<KeyedFilter> info() {
     }
 
     @Override
-    protected KeyedFilter replaceChild(LogicalPlan newChild) {
+    public KeyedFilter replaceChild(LogicalPlan newChild) {
         return new KeyedFilter(source(), newChild, keys, timestamp, tiebreaker);
     }
 

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/plan/logical/LimitWithOffset.java

@@ -34,7 +34,7 @@ protected NodeInfo<Limit> info() {
     }
 
     @Override
-    protected LimitWithOffset replaceChild(LogicalPlan newChild) {
+    public LimitWithOffset replaceChild(LogicalPlan newChild) {
         return new LimitWithOffset(source(), limit(), offset, newChild);
     }
 

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/plan/logical/Tail.java

@@ -33,7 +33,7 @@ protected NodeInfo<Limit> info() {
     }
 
     @Override
-    protected Tail replaceChild(LogicalPlan newChild) {
+    public Tail replaceChild(LogicalPlan newChild) {
         return new Tail(source(), newChild, limit());
     }
 }

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/optimizer/OptimizerRules.java

@@ -68,7 +68,6 @@
 
 import static java.lang.Math.signum;
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonList;
 import static org.elasticsearch.xpack.ql.expression.Literal.FALSE;
 import static org.elasticsearch.xpack.ql.expression.Literal.TRUE;
 import static org.elasticsearch.xpack.ql.expression.predicate.Predicates.combineAnd;
@@ -1207,14 +1206,12 @@ else if (child instanceof UnaryPlan) {
                     }
                     // if at least one expression can be pushed down, update the tree
                     if (conjunctions.size() > 0) {
-                        child = child.replaceChildrenSameSize(
-                            singletonList(filter.with(unary.child(), Predicates.combineAnd(conjunctions)))
-                        );
+                        child = unary.replaceChild(filter.with(unary.child(), Predicates.combineAnd(conjunctions)));
                         plan = filter.with(child, Predicates.combineAnd(inPlace));
                     }
                 } else {
                     // push down filter
-                    plan = child.replaceChildrenSameSize(singletonList(filter.with(unary.child(), condition)));
+                    plan = unary.replaceChild(filter.with(unary.child(), condition));
                 }
             }
 

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

@@ -34,7 +34,7 @@ protected NodeInfo<Aggregate> info() {
     }
 
     @Override
-    protected Aggregate replaceChild(LogicalPlan newChild) {
+    public Aggregate replaceChild(LogicalPlan newChild) {
         return new Aggregate(source(), newChild, groupings, aggregates);
     }
 

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

@@ -32,7 +32,7 @@ protected NodeInfo<Filter> info() {
     }
 
     @Override
-    protected Filter replaceChild(LogicalPlan newChild) {
+    public Filter replaceChild(LogicalPlan newChild) {
         return new Filter(source(), newChild, condition);
     }
 

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

@@ -27,7 +27,7 @@ protected NodeInfo<Limit> info() {
     }
 
     @Override
-    protected Limit replaceChild(LogicalPlan newChild) {
+    public Limit replaceChild(LogicalPlan newChild) {
         return new Limit(source(), limit, newChild);
     }
 

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

@@ -29,7 +29,7 @@ protected NodeInfo<OrderBy> info() {
     }
 
     @Override
-    protected OrderBy replaceChild(LogicalPlan newChild) {
+    public OrderBy replaceChild(LogicalPlan newChild) {
         return new OrderBy(source(), newChild, order);
     }
 

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

@@ -35,7 +35,7 @@ protected NodeInfo<Project> info() {
     }
 
     @Override
-    protected Project replaceChild(LogicalPlan newChild) {
+    public Project replaceChild(LogicalPlan newChild) {
         return new Project(source(), newChild, projections);
     }
 

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

@@ -31,7 +31,7 @@ public final UnaryPlan replaceChildren(List<LogicalPlan> newChildren) {
         return replaceChild(newChildren.get(0));
     }
 
-    protected abstract UnaryPlan replaceChild(LogicalPlan newChild);
+    public abstract UnaryPlan replaceChild(LogicalPlan newChild);
 
     public LogicalPlan child() {
         return child;

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java

@@ -784,7 +784,8 @@ private static LogicalPlan propagateMissing(LogicalPlan plan, AttributeSet missi
 
             // LeafPlans are tables and BinaryPlans are joins so pushing can only happen on unary
             if (plan instanceof UnaryPlan) {
-                return plan.replaceChildrenSameSize(singletonList(propagateMissing(((UnaryPlan) plan).child(), missing, failed)));
+                UnaryPlan unary = (UnaryPlan) plan;
+                return unary.replaceChild(propagateMissing(unary.child(), missing, failed));
             }
 
             failed.addAll(missing);

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

@@ -1240,7 +1240,7 @@ static class SkipQueryForLiteralAggregations extends OptimizerRule<Aggregate> {
         @Override
         protected LogicalPlan rule(Aggregate plan) {
             if (plan.groupings().isEmpty() && plan.child() instanceof EsRelation && plan.aggregates().stream().allMatch(this::foldable)) {
-                return plan.replaceChildrenSameSize(singletonList(new LocalRelation(plan.source(), new SingletonExecutable())));
+                return plan.replaceChild(new LocalRelation(plan.source(), new SingletonExecutable()));
             }
 
             return plan;

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Distinct.java

@@ -23,7 +23,7 @@ protected NodeInfo<Distinct> info() {
     }
 
     @Override
-    protected Distinct replaceChild(LogicalPlan newChild) {
+    public Distinct replaceChild(LogicalPlan newChild) {
         return new Distinct(source(), newChild);
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Having.java

@@ -29,7 +29,7 @@ protected NodeInfo<Filter> info() {
     }
 
     @Override
-    protected Having replaceChild(LogicalPlan newChild) {
+    public Having replaceChild(LogicalPlan newChild) {
         return new Having(source(), newChild, condition());
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Pivot.java

@@ -74,7 +74,7 @@ protected NodeInfo<Pivot> info() {
     }
 
     @Override
-    protected Pivot replaceChild(LogicalPlan newChild) {
+    public Pivot replaceChild(LogicalPlan newChild) {
         return new Pivot(source(), newChild, column, values, aggregates, grouping);
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/SubQueryAlias.java

@@ -33,7 +33,7 @@ protected NodeInfo<SubQueryAlias> info() {
     }
 
     @Override
-    protected SubQueryAlias replaceChild(LogicalPlan newChild) {
+    public SubQueryAlias replaceChild(LogicalPlan newChild) {
         return new SubQueryAlias(source(), newChild, alias);
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/With.java

@@ -28,7 +28,7 @@ protected NodeInfo<With> info() {
     }
 
     @Override
-    protected With replaceChild(LogicalPlan newChild) {
+    public With replaceChild(LogicalPlan newChild) {
         return new With(source(), newChild, subQueries);
     }