Verify the layout is not changed during optimizations

julian-elastic · julian-elastic · commit 7ec9efc03c5e · 2025-07-09T13:37:48.000-04:00
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/Attribute.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/Attribute.java
@@ -134,4 +134,19 @@ public String nodeString() {
     }
 
     protected abstract String label();
+
+    /**
+     * Compares the size and datatypes of two lists of attributes for equality.
+     */
+    public static boolean datatypeEquals(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).dataType().equals(right.get(i).dataType()) == false) {
+                return false;
+            }
+        }
+        return true;
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalLogicalPlanOptimizer.java b/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);
         }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java b/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
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java b/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);
         }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalVerifier.java b/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.datatypeEquals;
+
 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 (datatypeEquals(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());
         }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizer.java b/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
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalVerifier.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalVerifier.java
@@ -14,10 +14,13 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
+import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec;
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
+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;
 
 /** Physical plan verifier. */
 public final class PhysicalVerifier {
@@ -27,19 +30,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 +70,14 @@ public Failures verify(PhysicalPlan plan, boolean skipRemoteEnrichVerification)
             }
         });
 
+        if (datatypeEquals(planBefore.output(), planAfter.output()) == false) {
+            if ((planAfter instanceof EsQueryExec esQueryExec && esQueryExec.indexMode() == LOOKUP) == false) {
+                failures.add(
+                    fail(planAfter, "Layout has changed from [{}] to [{}]. ", planBefore.output().toString(), planAfter.output().toString())
+                );
+            }
+        }
+
         if (depFailures.hasFailures()) {
             throw new IllegalStateException(depFailures.toString());
         }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/rule/RuleExecutor.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/rule/RuleExecutor.java
@@ -175,7 +175,7 @@ protected final ExecutionInfo executeWithInfo(TreeType plan) {
                     if (tf.hasChanged()) {
                         hasChanged = true;
                         if (log.isTraceEnabled()) {
-                            log.trace("Rule {} applied\n{}", rule, NodeUtils.diffString(tf.before, tf.after));
+                            log.trace("Rule {} applied with change\n{}", rule, NodeUtils.diffString(tf.before, tf.after));
                         }
                     } else {
                         if (log.isTraceEnabled()) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/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);
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java b/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}#")

Original file line number	Diff line number	Diff line change
`@@ -86,7 +86,7 @@ private static Batch<LogicalPlan> localOperators() {`
`86`	`86`
`87`	`87`	`public LogicalPlan localOptimize(LogicalPlan plan) {`
`88`	`88`	`LogicalPlan optimized = execute(plan);`
`89`		`- Failures failures = verifier.verify(optimized, true);`
	`89`	`+ Failures failures = verifier.verify(optimized, true, plan);`
`90`	`90`	`if (failures.hasFailures()) {`
`91`	`91`	`throw new VerificationException(failures);`
`92`	`92`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,15 +42,15 @@ public LocalPhysicalPlanOptimizer(LocalPhysicalOptimizerContext context) {`
`42`	`42`	`}`
`43`	`43`
`44`	`44`	`public PhysicalPlan localOptimize(PhysicalPlan plan) {`
`45`		`- return verify(execute(plan));`
	`45`	`+ return verify(execute(plan), plan);`
`46`	`46`	`}`
`47`	`47`
`48`		`- PhysicalPlan verify(PhysicalPlan plan) {`
`49`		`- Failures failures = verifier.verify(plan, true);`
	`48`	`+ PhysicalPlan verify(PhysicalPlan planAfter, PhysicalPlan planBefore) {`
	`49`	`+ Failures failures = verifier.verify(planAfter, true, planBefore);`
`50`	`50`	`if (failures.hasFailures()) {`
`51`	`51`	`throw new VerificationException(failures);`
`52`	`52`	`}`
`53`		`- return plan;`
	`53`	`+ return planAfter;`
`54`	`54`	`}`
`55`	`55`
`56`	`56`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -112,7 +112,7 @@ public LogicalPlanOptimizer(LogicalOptimizerContext optimizerContext) {`
`112`	`112`	`public LogicalPlan optimize(LogicalPlan verified) {`
`113`	`113`	`var optimized = execute(verified);`
`114`	`114`
`115`		`- Failures failures = verifier.verify(optimized, false);`
	`115`	`+ Failures failures = verifier.verify(optimized, false, verified);`
`116`	`116`	`if (failures.hasFailures()) {`
`117`	`117`	`throw new VerificationException(failures);`
`118`	`118`	`}`
Original file line number	Diff line number	Diff line change
`@@ -34,15 +34,15 @@ public PhysicalPlanOptimizer(PhysicalOptimizerContext context) {`
`34`	`34`	`}`
`35`	`35`
`36`	`36`	`public PhysicalPlan optimize(PhysicalPlan plan) {`
`37`		`- return verify(execute(plan));`
	`37`	`+ return verify(execute(plan), plan);`
`38`	`38`	`}`
`39`	`39`
`40`		`- PhysicalPlan verify(PhysicalPlan plan) {`
`41`		`- Failures failures = verifier.verify(plan, false);`
	`40`	`+ PhysicalPlan verify(PhysicalPlan planAfter, PhysicalPlan planBefore) {`
	`41`	`+ Failures failures = verifier.verify(planAfter, false, planBefore);`
`42`	`42`	`if (failures.hasFailures()) {`
`43`	`43`	`throw new VerificationException(failures);`
`44`	`44`	`}`
`45`		`- return plan;`
	`45`	`+ return planAfter;`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -175,7 +175,7 @@ protected final ExecutionInfo executeWithInfo(TreeType plan) {`
`175`	`175`	`if (tf.hasChanged()) {`
`176`	`176`	`hasChanged = true;`
`177`	`177`	`if (log.isTraceEnabled()) {`
`178`		`- log.trace("Rule {} applied\n{}", rule, NodeUtils.diffString(tf.before, tf.after));`
	`178`	`+ log.trace("Rule {} applied with change\n{}", rule, NodeUtils.diffString(tf.before, tf.after));`
`179`	`179`	`}`
`180`	`180`	`} else {`
`181`	`181`	`if (log.isTraceEnabled()) {`