Fix more cases and address code review

Author: julian-elastic

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

@@ -138,7 +138,7 @@ public String nodeString() {
     /**
      * Compares the size and datatypes of two lists of attributes for equality.
      */
-    public static boolean datatypeEquals(List<Attribute> left, List<Attribute> right) {
+    public static boolean dataTypeEquals(List<Attribute> left, List<Attribute> right) {
         if (left.size() != right.size()) {
             return false;
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalVerifier.java

@@ -10,17 +10,17 @@
 import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
-import org.elasticsearch.xpack.esql.plan.QueryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
-public final class LogicalVerifier extends PostOptimizationPhasePlanVerifier {
+public final class LogicalVerifier extends PostOptimizationPhasePlanVerifier<LogicalPlan> {
 
     public static final LogicalVerifier INSTANCE = new LogicalVerifier();
 
     private LogicalVerifier() {}
 
     @Override
-    boolean skipVerification(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVerification) {
+    boolean skipVerification(LogicalPlan optimizedPlan, boolean skipRemoteEnrichVerification) {
         if (skipRemoteEnrichVerification) {
             // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
             var enriches = optimizedPlan.collectFirstChildren(Enrich.class::isInstance);
@@ -32,7 +32,7 @@ boolean skipVerification(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVer
     }
 
     @Override
-    void checkPlanConsistency(QueryPlan<?> optimizedPlan, Failures failures, Failures depFailures) {
+    void checkPlanConsistency(LogicalPlan optimizedPlan, Failures failures, Failures depFailures) {
         optimizedPlan.forEachUp(p -> {
             PlanConsistencyChecker.checkPlan(p, depFailures);
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalVerifier.java

@@ -12,22 +12,22 @@
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
-import org.elasticsearch.xpack.esql.plan.QueryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
+import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 
 /** Physical plan verifier. */
-public final class PhysicalVerifier extends PostOptimizationPhasePlanVerifier {
+public final class PhysicalVerifier extends PostOptimizationPhasePlanVerifier<PhysicalPlan> {
 
     public static final PhysicalVerifier INSTANCE = new PhysicalVerifier();
 
     private PhysicalVerifier() {}
 
     @Override
-    boolean skipVerification(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVerification) {
+    boolean skipVerification(PhysicalPlan optimizedPlan, boolean skipRemoteEnrichVerification) {
         if (skipRemoteEnrichVerification) {
             // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
             var enriches = optimizedPlan.collectFirstChildren(EnrichExec.class::isInstance);
@@ -39,7 +39,7 @@ boolean skipVerification(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVer
     }
 
     @Override
-    void checkPlanConsistency(QueryPlan<?> optimizedPlan, Failures failures, Failures depFailures) {
+    void checkPlanConsistency(PhysicalPlan optimizedPlan, Failures failures, Failures depFailures) {
         optimizedPlan.forEachDown(p -> {
             if (p instanceof FieldExtractExec fieldExtractExec) {
                 Attribute sourceAttribute = fieldExtractExec.sourceAttribute();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PostOptimizationPhasePlanVerifier.java

@@ -9,14 +9,15 @@
 
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.ProjectAwayColumns;
 import org.elasticsearch.xpack.esql.plan.QueryPlan;
 import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec;
 
 import java.util.List;
 
 import static org.elasticsearch.index.IndexMode.LOOKUP;
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
-import static org.elasticsearch.xpack.esql.core.expression.Attribute.datatypeEquals;
+import static org.elasticsearch.xpack.esql.core.expression.Attribute.dataTypeEquals;
 
 /**
  * Verifies the plan after optimization.
@@ -25,10 +26,10 @@
  * LocalLogicalPlanOptimizer, and LocalPhysicalPlanOptimizer.
  * Note: Logical and Physical optimizers may override methods in this class to perform different checks.
  */
-public abstract class PostOptimizationPhasePlanVerifier {
+public abstract class PostOptimizationPhasePlanVerifier<P extends QueryPlan<P>> {
 
     /** Verifies the optimized plan */
-    public Failures verify(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVerification, List<Attribute> expectedOutputAttributes) {
+    public Failures verify(P optimizedPlan, boolean skipRemoteEnrichVerification, List<Attribute> expectedOutputAttributes) {
         Failures failures = new Failures();
         Failures depFailures = new Failures();
         if (skipVerification(optimizedPlan, skipRemoteEnrichVerification)) {
@@ -46,16 +47,25 @@ public Failures verify(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVerif
         return failures;
     }
 
-    abstract boolean skipVerification(QueryPlan<?> optimizedPlan, boolean skipRemoteEnrichVerification);
+    abstract boolean skipVerification(P optimizedPlan, boolean skipRemoteEnrichVerification);
 
-    abstract void checkPlanConsistency(QueryPlan<?> optimizedPlan, Failures failures, Failures depFailures);
+    abstract void checkPlanConsistency(P optimizedPlan, Failures failures, Failures depFailures);
 
     private static void verifyOutputNotChanged(QueryPlan<?> optimizedPlan, List<Attribute> expectedOutputAttributes, Failures failures) {
-        if (datatypeEquals(expectedOutputAttributes, optimizedPlan.output()) == false) {
+        if (dataTypeEquals(expectedOutputAttributes, optimizedPlan.output()) == false) {
+            // If the output level is empty we add a column called ProjectAwayColumns.ALL_FIELDS_PROJECTED
+            // We will ignore such cases for output verification
+            boolean hasProjectAwayColumns = optimizedPlan.output()
+                .stream()
+                .anyMatch(x -> x.name().equals(ProjectAwayColumns.ALL_FIELDS_PROJECTED));
             // LookupJoinExec represents the lookup index with EsSourceExec and this is turned into EsQueryExec by
             // ReplaceSourceAttributes. Because InsertFieldExtractions doesn't apply to lookup indices, the
             // right hand side will only have the EsQueryExec providing the _doc attribute and nothing else.
-            if ((optimizedPlan instanceof EsQueryExec esQueryExec && esQueryExec.indexMode() == LOOKUP) == false) {
+            // We perform an optimizer run on every fragment. LookupJoinExec also contains such a fragment,
+            // and currently it only contains an EsQueryExec after optimization.
+            boolean hasLookupJoinExec = optimizedPlan instanceof EsQueryExec esQueryExec && esQueryExec.indexMode() == LOOKUP;
+            boolean ignoreError = hasProjectAwayColumns || hasLookupJoinExec;
+            if (ignoreError == false) {
                 failures.add(
                     fail(
                         optimizedPlan,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/ProjectAwayColumns.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 
 import static java.lang.Boolean.FALSE;
@@ -36,6 +37,7 @@
  * extraction.
  */
 public class ProjectAwayColumns extends Rule<PhysicalPlan, PhysicalPlan> {
+    public static String ALL_FIELDS_PROJECTED = "<all-fields-projected>";
 
     @Override
     public PhysicalPlan apply(PhysicalPlan plan) {
@@ -44,7 +46,7 @@ public PhysicalPlan apply(PhysicalPlan plan) {
         // and the overall output will not change.
         AttributeSet.Builder requiredAttrBuilder = plan.outputSet().asBuilder();
 
-        return plan.transformDown(currentPlanNode -> {
+        PhysicalPlan planAfter = plan.transformDown(currentPlanNode -> {
             if (keepTraversing.get() == false) {
                 return currentPlanNode;
             }
@@ -82,12 +84,13 @@ public PhysicalPlan apply(PhysicalPlan plan) {
                         // add a synthetic field (so it doesn't clash with the user defined one) to return a constant
                         // to avoid the block from being trimmed
                         if (output.isEmpty()) {
-                            var alias = new Alias(logicalFragment.source(), "<all-fields-projected>", Literal.NULL, null, true);
+                            var alias = new Alias(logicalFragment.source(), ALL_FIELDS_PROJECTED, Literal.NULL, null, true);
                             List<Alias> fields = singletonList(alias);
                             logicalFragment = new Eval(logicalFragment.source(), logicalFragment, fields);
                             output = Expressions.asAttributes(fields);
                         }
                         // add a logical projection (let the local replanning remove it if needed)
+                        output = reorderOutput(logicalFragment.output(), output);
                         FragmentExec newChild = new FragmentExec(
                             Source.EMPTY,
                             new Project(logicalFragment.source(), logicalFragment, output),
@@ -105,5 +108,27 @@ public PhysicalPlan apply(PhysicalPlan plan) {
             }
             return currentPlanNode;
         });
+        return planAfter;
+    }
+
+    private List<Attribute> reorderOutput(List<Attribute> expectedOrder, List<Attribute> attributesToSort) {
+        List<Attribute> output = new ArrayList<>(attributesToSort.size());
+        // Track which attributes have been added
+        var added = new HashSet<Attribute>();
+
+        // Add attributes in expectedOrder if present in attributesToSort
+        for (Attribute expected : expectedOrder) {
+            if (attributesToSort.contains(expected)) {
+                output.add(expected);
+                added.add(expected);
+            }
+        }
+        // Add remaining attributes from attributesToSort in their original order
+        for (Attribute attr : attributesToSort) {
+            if (added.contains(attr) == false) {
+                output.add(attr);
+            }
+        }
+        return output;
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java

@@ -131,6 +131,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.local.EmptyLocalSupplier;
 import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
+import org.elasticsearch.xpack.esql.rule.RuleExecutor;
 
 import java.time.Duration;
 import java.util.ArrayList;
@@ -178,6 +179,8 @@
 import static org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison.BinaryComparisonOperation.GTE;
 import static org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison.BinaryComparisonOperation.LT;
 import static org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison.BinaryComparisonOperation.LTE;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.DOWN;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.UP;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.contains;
@@ -7855,4 +7858,101 @@ public void testSampleNoPushDownChangePoint() {
         var topN = as(changePoint.child(), TopN.class);
         var source = as(topN.child(), EsRelation.class);
     }
+
+    private LogicalPlanOptimizer getCustomRulesLogicalPlanOptimizer(List<RuleExecutor.Batch<LogicalPlan>> batches) {
+        LogicalOptimizerContext context = new LogicalOptimizerContext(EsqlTestUtils.TEST_CFG, FoldContext.small());
+        LogicalPlanOptimizer customOptimizer = new LogicalPlanOptimizer(context) {
+            @Override
+            protected List<Batch<LogicalPlan>> batches() {
+                return batches;
+            }
+        };
+        return customOptimizer;
+    }
+
+    public void testVerifierOnAdditionalAttributeAdded() throws Exception {
+        var plan = optimizedPlan("""
+            from test
+            | stats a = min(salary) by emp_no
+            """);
+
+        var limit = as(plan, Limit.class);
+        var aggregate = as(limit.child(), Aggregate.class);
+        var min = as(Alias.unwrap(aggregate.aggregates().get(0)), Min.class);
+        var salary = as(min.field(), NamedExpression.class);
+        assertThat(salary.name(), is("salary"));
+
+        // use a custom rule that adds another output attribute
+        var customRuleBatch = new RuleExecutor.Batch<>(
+            "CustomRuleBatch",
+            RuleExecutor.Limiter.ONCE,
+            new OptimizerRules.ParameterizedOptimizerRule<Aggregate, LogicalOptimizerContext>(UP) {
+                static Integer appliedCount = 0;
+
+                @Override
+                protected LogicalPlan rule(Aggregate plan, LogicalOptimizerContext context) {
+                    // This rule adds a missing attribute to the plan output
+                    // We only want to apply it once, so we use a static counter
+                    if (appliedCount == 0) {
+                        appliedCount++;
+                        Literal additionalLiteral = new Literal(Source.EMPTY, "additional literal", INTEGER);
+                        return new Eval(plan.source(), plan, List.of(new Alias(Source.EMPTY, "additionalAttribute", additionalLiteral)));
+                    }
+                    return plan;
+                }
+
+            }
+        );
+        LogicalPlanOptimizer customRulesLogicalPlanOptimizer = getCustomRulesLogicalPlanOptimizer(List.of(customRuleBatch));
+        Exception e = expectThrows(VerificationException.class, () -> customRulesLogicalPlanOptimizer.optimize(plan));
+        assertThat(e.getMessage(), containsString("Output has changed from"));
+        assertThat(e.getMessage(), containsString("additionalAttribute"));
+    }
+
+    public void testVerifierOnAttributeDatatypeChanged() {
+        var plan = optimizedPlan("""
+            from test
+            | stats a = min(salary) by emp_no
+            """);
+
+        var limit = as(plan, Limit.class);
+        var aggregate = as(limit.child(), Aggregate.class);
+        var min = as(Alias.unwrap(aggregate.aggregates().get(0)), Min.class);
+        var salary = as(min.field(), NamedExpression.class);
+        assertThat(salary.name(), is("salary"));
+
+        // use a custom rule that changes the datatype of an output attribute
+        var customRuleBatch = new RuleExecutor.Batch<>(
+            "CustomRuleBatch",
+            RuleExecutor.Limiter.ONCE,
+            new OptimizerRules.ParameterizedOptimizerRule<LogicalPlan, LogicalOptimizerContext>(DOWN) {
+                static Integer appliedCount = 0;
+
+                @Override
+                protected LogicalPlan rule(LogicalPlan plan, LogicalOptimizerContext context) {
+                    // We only want to apply it once, so we use a static counter
+                    if (appliedCount == 0) {
+                        appliedCount++;
+                        Limit limit = as(plan, Limit.class);
+                        Limit newLimit = new Limit(plan.source(), limit.limit(), limit.child()) {
+                            @Override
+                            public List<Attribute> output() {
+                                List<Attribute> oldOutput = super.output();
+                                List<Attribute> newOutput = new ArrayList<>(oldOutput);
+                                newOutput.set(0, oldOutput.get(0).withDataType(DataType.DATETIME));
+                                return newOutput;
+                            }
+                        };
+                        return newLimit;
+                    }
+                    return plan;
+                }
+
+            }
+        );
+        LogicalPlanOptimizer customRulesLogicalPlanOptimizer = getCustomRulesLogicalPlanOptimizer(List.of(customRuleBatch));
+        Exception e = expectThrows(VerificationException.class, () -> customRulesLogicalPlanOptimizer.optimize(plan));
+        assertThat(e.getMessage(), containsString("Output has changed from"));
+    }
+
 }