Add checks that optimizers do not modify the layout

Author: julian-elastic

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

@@ -134,4 +134,16 @@ public String nodeString() {
     }
 
     protected abstract String label();
+
+    public static boolean semanticEquals(List<Attribute> left, List<Attribute> right) {
+        if (left.size() != right.size()) {
+            return false;
+        }
+        for (int i = 0; i < left.size(); i++) {
+            if (left.get(i).semanticEquals(right.get(i)) == false) {
+                return false;
+            }
+        }
+        return true;
+    }
 }

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

@@ -86,7 +86,7 @@ private static Batch<LogicalPlan> localOperators() {
 
     public LogicalPlan localOptimize(LogicalPlan plan) {
         LogicalPlan optimized = execute(plan);
-        Failures failures = verifier.verify(optimized, true);
+        Failures failures = verifier.verify(optimized, true, plan);
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }

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

@@ -42,15 +42,15 @@ public LocalPhysicalPlanOptimizer(LocalPhysicalOptimizerContext context) {
     }
 
     public PhysicalPlan localOptimize(PhysicalPlan plan) {
-        return verify(execute(plan));
+        return verify(execute(plan), plan);
     }
 
-    PhysicalPlan verify(PhysicalPlan plan) {
-        Failures failures = verifier.verify(plan, true);
+    PhysicalPlan verify(PhysicalPlan planAfter, PhysicalPlan planBefore) {
+        Failures failures = verifier.verify(planAfter, true, planBefore);
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }
-        return plan;
+        return planAfter;
     }
 
     @Override

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

@@ -112,7 +112,7 @@ public LogicalPlanOptimizer(LogicalOptimizerContext optimizerContext) {
     public LogicalPlan optimize(LogicalPlan verified) {
         var optimized = execute(verified);
 
-        Failures failures = verifier.verify(optimized, false);
+        Failures failures = verifier.verify(optimized, false, verified);
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }

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

@@ -13,26 +13,29 @@
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
+import static org.elasticsearch.xpack.esql.common.Failure.fail;
+import static org.elasticsearch.xpack.esql.core.expression.Attribute.semanticEquals;
+
 public final class LogicalVerifier {
 
     public static final LogicalVerifier INSTANCE = new LogicalVerifier();
 
     private LogicalVerifier() {}
 
     /** Verifies the optimized logical plan. */
-    public Failures verify(LogicalPlan plan, boolean skipRemoteEnrichVerification) {
+    public Failures verify(LogicalPlan planAfter, boolean skipRemoteEnrichVerification, LogicalPlan planBefore) {
         Failures failures = new Failures();
         Failures dependencyFailures = new Failures();
 
         if (skipRemoteEnrichVerification) {
             // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
-            var enriches = plan.collectFirstChildren(Enrich.class::isInstance);
+            var enriches = planAfter.collectFirstChildren(Enrich.class::isInstance);
             if (enriches.isEmpty() == false && ((Enrich) enriches.get(0)).mode() == Enrich.Mode.REMOTE) {
                 return failures;
             }
         }
 
-        plan.forEachUp(p -> {
+        planAfter.forEachUp(p -> {
             PlanConsistencyChecker.checkPlan(p, dependencyFailures);
 
             if (failures.hasFailures() == false) {
@@ -47,6 +50,12 @@ public Failures verify(LogicalPlan plan, boolean skipRemoteEnrichVerification) {
             }
         });
 
+        if (semanticEquals(planBefore.output(), planAfter.output()) == false) {
+            failures.add(
+                fail(planAfter, "Layout has changed from [{}] to [{}]. ", planBefore.output().toString(), planAfter.output().toString())
+            );
+        }
+
         if (dependencyFailures.hasFailures()) {
             throw new IllegalStateException(dependencyFailures.toString());
         }

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

@@ -34,15 +34,15 @@ public PhysicalPlanOptimizer(PhysicalOptimizerContext context) {
     }
 
     public PhysicalPlan optimize(PhysicalPlan plan) {
-        return verify(execute(plan));
+        return verify(execute(plan), plan);
     }
 
-    PhysicalPlan verify(PhysicalPlan plan) {
-        Failures failures = verifier.verify(plan, false);
+    PhysicalPlan verify(PhysicalPlan planAfter, PhysicalPlan planBefore) {
+        Failures failures = verifier.verify(planAfter, false, planBefore);
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }
-        return plan;
+        return planAfter;
     }
 
     @Override

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

@@ -27,19 +27,19 @@ public final class PhysicalVerifier {
     private PhysicalVerifier() {}
 
     /** Verifies the physical plan. */
-    public Failures verify(PhysicalPlan plan, boolean skipRemoteEnrichVerification) {
+    public Failures verify(PhysicalPlan planAfter, boolean skipRemoteEnrichVerification, PhysicalPlan planBefore) {
         Failures failures = new Failures();
         Failures depFailures = new Failures();
 
         if (skipRemoteEnrichVerification) {
             // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
-            var enriches = plan.collectFirstChildren(EnrichExec.class::isInstance);
+            var enriches = planAfter.collectFirstChildren(EnrichExec.class::isInstance);
             if (enriches.isEmpty() == false && ((EnrichExec) enriches.get(0)).mode() == Enrich.Mode.REMOTE) {
                 return failures;
             }
         }
 
-        plan.forEachDown(p -> {
+        planAfter.forEachDown(p -> {
             if (p instanceof FieldExtractExec fieldExtractExec) {
                 Attribute sourceAttribute = fieldExtractExec.sourceAttribute();
                 if (sourceAttribute == null) {
@@ -67,6 +67,12 @@ public Failures verify(PhysicalPlan plan, boolean skipRemoteEnrichVerification)
             }
         });
 
+        if (planBefore.output().equals(planAfter.output()) == false) {
+            failures.add(
+                fail(planAfter, "Layout has changed from [{}] to [{}]. ", planBefore.output().toString(), planAfter.output().toString())
+            );
+        }
+
         if (depFailures.hasFailures()) {
             throw new IllegalStateException(depFailures.toString());
         }

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

@@ -5554,7 +5554,7 @@ public void testPushShadowingGeneratingPlanPastProject() {
             List<Attribute> initialGeneratedExprs = ((GeneratingPlan) initialPlan).generatedAttributes();
             LogicalPlan optimizedPlan = testCase.rule.apply(initialPlan);
 
-            Failures inconsistencies = LogicalVerifier.INSTANCE.verify(optimizedPlan, false);
+            Failures inconsistencies = LogicalVerifier.INSTANCE.verify(optimizedPlan, false, initialPlan);
             assertFalse(inconsistencies.hasFailures());
 
             Project project = as(optimizedPlan, Project.class);
@@ -5605,7 +5605,7 @@ public void testPushShadowingGeneratingPlanPastRenamingProject() {
             List<Attribute> initialGeneratedExprs = ((GeneratingPlan) initialPlan).generatedAttributes();
             LogicalPlan optimizedPlan = testCase.rule.apply(initialPlan);
 
-            Failures inconsistencies = LogicalVerifier.INSTANCE.verify(optimizedPlan, false);
+            Failures inconsistencies = LogicalVerifier.INSTANCE.verify(optimizedPlan, false, initialPlan);
             assertFalse(inconsistencies.hasFailures());
 
             Project project = as(optimizedPlan, Project.class);
@@ -5661,7 +5661,7 @@ public void testPushShadowingGeneratingPlanPastRenamingProjectWithResolution() {
 
             // This ensures that our generating plan doesn't use invalid references, resp. that any rename from the Project has
             // been propagated into the generating plan.
-            Failures inconsistencies = LogicalVerifier.INSTANCE.verify(optimizedPlan, false);
+            Failures inconsistencies = LogicalVerifier.INSTANCE.verify(optimizedPlan, false, initialPlan);
             assertFalse(inconsistencies.hasFailures());
 
             Project project = as(optimizedPlan, Project.class);

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

@@ -2865,7 +2865,7 @@ public void testFieldExtractWithoutSourceAttributes() {
             )
         );
 
-        var e = expectThrows(VerificationException.class, () -> physicalPlanOptimizer.verify(badPlan));
+        var e = expectThrows(VerificationException.class, () -> physicalPlanOptimizer.verify(badPlan, verifiedPlan));
         assertThat(
             e.getMessage(),
             containsString(
@@ -2880,7 +2880,7 @@ public void testVerifierOnMissingReferences() {
             | stats s = sum(salary) by emp_no
             | where emp_no > 10
             """);
-
+        final var planBeforeModification = plan;
         plan = plan.transformUp(
             AggregateExec.class,
             a -> new AggregateExec(
@@ -2894,7 +2894,7 @@ public void testVerifierOnMissingReferences() {
             )
         );
         final var finalPlan = plan;
-        var e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(finalPlan));
+        var e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(finalPlan, planBeforeModification));
         assertThat(e.getMessage(), containsString(" > 10[INTEGER]]] optimized incorrectly due to missing references [emp_no{f}#"));
     }
 
@@ -2912,7 +2912,7 @@ public void testVerifierOnMissingReferencesWithBinaryPlans() throws Exception {
 
         var planWithInvalidJoinLeftSide = plan.transformUp(LookupJoinExec.class, join -> join.replaceChildren(join.right(), join.right()));
 
-        var e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(planWithInvalidJoinLeftSide));
+        var e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(planWithInvalidJoinLeftSide, plan));
         assertThat(e.getMessage(), containsString(" optimized incorrectly due to missing references from left hand side [languages"));
 
         var planWithInvalidJoinRightSide = plan.transformUp(
@@ -2929,7 +2929,7 @@ public void testVerifierOnMissingReferencesWithBinaryPlans() throws Exception {
             )
         );
 
-        e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(planWithInvalidJoinRightSide));
+        e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(planWithInvalidJoinRightSide, plan));
         assertThat(e.getMessage(), containsString(" optimized incorrectly due to missing references from right hand side [language_code"));
     }
 
@@ -2939,7 +2939,7 @@ public void testVerifierOnDuplicateOutputAttributes() {
             | stats s = sum(salary) by emp_no
             | where emp_no > 10
             """);
-
+        final var planBeforeModification = plan;
         plan = plan.transformUp(AggregateExec.class, a -> {
             List<Attribute> intermediates = new ArrayList<>(a.intermediateAttributes());
             intermediates.add(intermediates.get(0));
@@ -2954,7 +2954,7 @@ public void testVerifierOnDuplicateOutputAttributes() {
             );
         });
         final var finalPlan = plan;
-        var e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(finalPlan));
+        var e = expectThrows(IllegalStateException.class, () -> physicalPlanOptimizer.verify(finalPlan, planBeforeModification));
         assertThat(
             e.getMessage(),
             containsString("Plan [LimitExec[1000[INTEGER],null]] optimized incorrectly due to duplicate output attribute emp_no{f}#")