Make OptimizerExpressionRule conditional

Author: idegtiarenko

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

@@ -184,7 +184,6 @@ protected void doCollectFirst(Predicate<? super T> predicate, List<T> matches) {
     public T transformDown(Function<? super T, ? extends T> rule) {
         T root = rule.apply((T) this);
         Node<T> node = this.equals(root) ? this : root;
-
         return node.transformChildren(child -> child.transformDown(rule));
     }
 
@@ -194,6 +193,12 @@ public <E extends T> T transformDown(Class<E> typeToken, Function<E, ? extends T
         return transformDown((t) -> (typeToken.isInstance(t) ? rule.apply((E) t) : t));
     }
 
+    @SuppressWarnings("unchecked")
+    public <E extends T> T transformDown(Predicate<Node<?>> tokenPredicate, Function<E, ? extends T> rule) {
+        // type filtering function
+        return transformDown((t) -> (tokenPredicate.test(t) ? rule.apply((E) t) : t));
+    }
+
     @SuppressWarnings("unchecked")
     public T transformUp(Function<? super T, ? extends T> rule) {
         T transformed = transformChildren(child -> child.transformUp(rule));
@@ -207,6 +212,12 @@ public <E extends T> T transformUp(Class<E> typeToken, Function<E, ? extends T>
         return transformUp((t) -> (typeToken.isInstance(t) ? rule.apply((E) t) : t));
     }
 
+    @SuppressWarnings("unchecked")
+    public <E extends T> T transformUp(Predicate<Node<?>> tokenPredicate, Function<E, ? extends T> rule) {
+        // type filtering function
+        return transformUp((t) -> (tokenPredicate.test(t) ? rule.apply((E) t) : t));
+    }
+
     @SuppressWarnings("unchecked")
     protected <R extends Function<? super T, ? extends T>> T transformChildren(Function<T, ? extends T> traversalOperation) {
         boolean childrenChanged = false;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/OptimizerRules.java

@@ -7,9 +7,12 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.tree.Node;
 import org.elasticsearch.xpack.esql.core.util.ReflectionUtils;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
+import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRule;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
@@ -55,12 +58,20 @@ public OptimizerExpressionRule(TransformDirection direction) {
         @Override
         public final LogicalPlan apply(LogicalPlan plan, LogicalOptimizerContext ctx) {
             return direction == TransformDirection.DOWN
-                ? plan.transformExpressionsDown(expressionTypeToken, e -> rule(e, ctx))
-                : plan.transformExpressionsUp(expressionTypeToken, e -> rule(e, ctx));
+                ? plan.transformExpressionsDown(this::shouldVisit, expressionTypeToken, e -> rule(e, ctx))
+                : plan.transformExpressionsUp(this::shouldVisit, expressionTypeToken, e -> rule(e, ctx));
         }
 
         protected abstract Expression rule(E e, LogicalOptimizerContext ctx);
 
+        protected boolean shouldVisit(Node<?> node) {
+            return switch (node) {
+                case EsRelation esr -> false;
+                case Project p -> false;// this covers both keep and project
+                default -> true;
+            };
+        }
+
         public Class<E> expressionToken() {
             return expressionTypeToken;
         }

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

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 /**
  * There are two main types of plans, {@code LogicalPlan} and {@code PhysicalPlan}
@@ -109,22 +110,36 @@ public <E extends Expression> PlanType transformExpressionsOnlyUp(Class<E> typeT
         return transformPropertiesOnly(Object.class, e -> doTransformExpression(e, exp -> exp.transformUp(typeToken, rule)));
     }
 
-    public PlanType transformExpressionsDown(Function<Expression, ? extends Expression> rule) {
-        return transformExpressionsDown(Expression.class, rule);
-    }
-
     public <E extends Expression> PlanType transformExpressionsDown(Class<E> typeToken, Function<E, ? extends Expression> rule) {
         return transformPropertiesDown(Object.class, e -> doTransformExpression(e, exp -> exp.transformDown(typeToken, rule)));
     }
 
-    public PlanType transformExpressionsUp(Function<Expression, ? extends Expression> rule) {
-        return transformExpressionsUp(Expression.class, rule);
+    public <E extends Expression> PlanType transformExpressionsDown(
+        Predicate<Node<?>> shouldVisit,
+        Class<E> typeToken,
+        Function<E, ? extends Expression> rule
+    ) {
+        return transformDown(
+            shouldVisit,
+            t -> t.transformNodeProps(Object.class, e -> doTransformExpression(e, exp -> exp.transformDown(typeToken, rule)))
+        );
     }
 
     public <E extends Expression> PlanType transformExpressionsUp(Class<E> typeToken, Function<E, ? extends Expression> rule) {
         return transformPropertiesUp(Object.class, e -> doTransformExpression(e, exp -> exp.transformUp(typeToken, rule)));
     }
 
+    public <E extends Expression> PlanType transformExpressionsUp(
+        Predicate<Node<?>> shouldVisit,
+        Class<E> typeToken,
+        Function<E, ? extends Expression> rule
+    ) {
+        return transformUp(
+            shouldVisit,
+            t -> t.transformNodeProps(Object.class, e -> doTransformExpression(e, exp -> exp.transformUp(typeToken, rule)))
+        );
+    }
+
     @SuppressWarnings("unchecked")
     private static Object doTransformExpression(Object arg, Function<Expression, ? extends Expression> traversal) {
         if (arg instanceof Expression exp) {
@@ -184,18 +199,10 @@ public <E extends Expression> void forEachExpression(Class<E> typeToken, Consume
         forEachPropertyOnly(Object.class, e -> doForEachExpression(e, exp -> exp.forEachDown(typeToken, rule)));
     }
 
-    public void forEachExpressionDown(Consumer<? super Expression> rule) {
-        forEachExpressionDown(Expression.class, rule);
-    }
-
     public <E extends Expression> void forEachExpressionDown(Class<? extends E> typeToken, Consumer<? super E> rule) {
         forEachPropertyDown(Object.class, e -> doForEachExpression(e, exp -> exp.forEachDown(typeToken, rule)));
     }
 
-    public void forEachExpressionUp(Consumer<? super Expression> rule) {
-        forEachExpressionUp(Expression.class, rule);
-    }
-
     public <E extends Expression> void forEachExpressionUp(Class<E> typeToken, Consumer<? super E> rule) {
         forEachPropertyUp(Object.class, e -> doForEachExpression(e, exp -> exp.forEachUp(typeToken, rule)));
     }