Add a Physical Plan Verifier (ESQL-1045)

Add a Physical Plan Verifier phase, that verifies the final plan.

Currently we only need this for the source attribute check of the Field
extractor, but other checks could be accommodated later as needed.

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

@@ -26,6 +26,9 @@ import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
 import org.elasticsearch.xpack.esql.plan.physical.RegexExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
 import org.elasticsearch.xpack.esql.plan.physical.UnaryExec;
+import org.elasticsearch.xpack.esql.planner.VerificationException;
+import org.elasticsearch.xpack.esql.planner.Verifier;
+import org.elasticsearch.xpack.ql.common.Failure;
 import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.AttributeSet;
@@ -48,6 +51,7 @@ import org.elasticsearch.xpack.ql.util.Holder;
 import org.elasticsearch.xpack.ql.util.ReflectionUtils;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
@@ -68,12 +72,23 @@ public class PhysicalPlanOptimizer extends ParameterizedRuleExecutor<PhysicalPla
 
     private static final Iterable<RuleExecutor.Batch<PhysicalPlan>> rules = initializeRules(true);
 
+    private final Verifier verifier;
+
     public PhysicalPlanOptimizer(PhysicalOptimizerContext context) {
         super(context);
+        this.verifier = new Verifier();
     }
 
     public PhysicalPlan optimize(PhysicalPlan plan) {
-        return execute(plan);
+        return verify(execute(plan));
+    }
+
+    PhysicalPlan verify(PhysicalPlan plan) {
+        Collection<Failure> failures = verifier.verify(plan);
+        if (failures.isEmpty() == false) {
+            throw new VerificationException(failures);
+        }
+        return plan;
     }
 
     static Iterable<RuleExecutor.Batch<PhysicalPlan>> initializeRules(boolean isOptimizedForEsSource) {

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

@@ -8,9 +8,7 @@
 package org.elasticsearch.xpack.esql.plan.physical;
 
 import org.elasticsearch.compute.ann.Experimental;
-import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
 import org.elasticsearch.xpack.ql.expression.Attribute;
-import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.NodeUtils;
 import org.elasticsearch.xpack.ql.tree.Source;
@@ -29,15 +27,6 @@ public class FieldExtractExec extends UnaryExec {
         super(source, child);
         this.attributesToExtract = attributesToExtract;
         this.sourceAttribute = extractSourceAttributesFrom(child);
-
-        // TODO: this can be moved into the physical verifier
-        if (sourceAttribute == null) {
-            throw new QlIllegalArgumentException(
-                "Need to add field extractor for [{}] but cannot detect source attributes from node [{}]",
-                Expressions.names(attributesToExtract),
-                child
-            );
-        }
     }
 
     public static Attribute extractSourceAttributesFrom(PhysicalPlan plan) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/VerificationException.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.planner;
+
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.esql.EsqlClientException;
+import org.elasticsearch.xpack.ql.common.Failure;
+
+import java.util.Collection;
+
+public class VerificationException extends EsqlClientException {
+
+    public VerificationException(Collection<Failure> sources) {
+        super(Failure.failMessage(sources));
+    }
+
+    @Override
+    public RestStatus status() {
+        return RestStatus.BAD_REQUEST;
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/Verifier.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.planner;
+
+import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
+import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
+import org.elasticsearch.xpack.ql.common.Failure;
+import org.elasticsearch.xpack.ql.expression.Attribute;
+import org.elasticsearch.xpack.ql.expression.Expressions;
+
+import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+import static org.elasticsearch.xpack.ql.common.Failure.fail;
+
+/** Physical plan verifier. */
+public final class Verifier {
+
+    /** Verifies the physical plan. */
+    public Collection<Failure> verify(PhysicalPlan plan) {
+        Set<Failure> failures = new LinkedHashSet<>();
+
+        plan.forEachDown(p -> {
+            if (p instanceof FieldExtractExec fieldExtractExec) {
+                Attribute sourceAttribute = fieldExtractExec.sourceAttribute();
+                if (sourceAttribute == null) {
+                    failures.add(
+                        fail(
+                            fieldExtractExec,
+                            "Need to add field extractor for [{}] but cannot detect source attributes from node [{}]",
+                            Expressions.names(fieldExtractExec.attributesToExtract()),
+                            fieldExtractExec.child()
+                        )
+                    );
+                }
+            }
+        });
+
+        return failures;
+    }
+}

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

@@ -39,9 +39,11 @@ import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
 import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
 import org.elasticsearch.xpack.esql.planner.Mapper;
+import org.elasticsearch.xpack.esql.planner.VerificationException;
 import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
+import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.FieldAttribute;
 import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.expression.Order;
@@ -1166,4 +1168,25 @@ public class PhysicalPlanOptimizerTests extends ESTestCase {
         assertThat(remoteSink.mode(), equalTo(ExchangeExec.Mode.REMOTE_SINK));
         return remoteSink;
     }
+
+    public void testFieldExtractWithoutSourceAttributes() {
+        PhysicalPlan verifiedPlan = optimizedPlan(physicalPlan("""
+            from test
+            | where round(emp_no) > 10
+            """));
+        // Transform the verified plan so that it is invalid (i.e. no source attributes)
+        List<Attribute> emptyAttrList = List.of();
+        var badPlan = verifiedPlan.transformDown(
+            EsQueryExec.class,
+            node -> new EsSourceExec(node.source(), node.index(), emptyAttrList, node.query())
+        );
+
+        var e = expectThrows(VerificationException.class, () -> physicalPlanOptimizer.verify(badPlan));
+        assertThat(
+            e.getMessage(),
+            containsString(
+                "Need to add field extractor for [[emp_no]] but cannot detect source attributes from node [EsSourceExec[test][]]"
+            )
+        );
+    }
 }