ESQL: Reorganize optimizer rules (#112338)

- Organize the optimizer rules consistently for all 4 optimizers (logical, physical, local logical, local physical).
- Move helper methods meant for optimizer rules out of the optimizers into the relevant rules or into helper classes.
- Consolidate the 2 nearly identical logical ParameterizedRules into one.

Author: alex-spies

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

@@ -61,7 +61,7 @@
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.type.DataType.BOOLEAN;
-import static org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizer.PushFiltersToSource.canPushToSource;
+import static org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushFiltersToSource.canPushToSource;
 
 /**
  * This class is part of the planner. Responsible for failing impossible queries with a human-readable error message.  In particular, this

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

@@ -9,12 +9,12 @@
 
 import org.elasticsearch.xpack.esql.capabilities.Validatable;
 import org.elasticsearch.xpack.esql.common.Failures;
-import org.elasticsearch.xpack.esql.optimizer.OptimizerRules.DependencyConsistency;
+import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
 public final class LogicalVerifier {
 
-    private static final DependencyConsistency<LogicalPlan> DEPENDENCY_CHECK = new DependencyConsistency<>();
+    private static final PlanConsistencyChecker<LogicalPlan> DEPENDENCY_CHECK = new PlanConsistencyChecker<>();
     public static final LogicalVerifier INSTANCE = new LogicalVerifier();
 
     private LogicalVerifier() {}

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

@@ -11,7 +11,7 @@
 import org.elasticsearch.xpack.esql.core.rule.ParameterizedRule;
 import org.elasticsearch.xpack.esql.core.rule.Rule;
 import org.elasticsearch.xpack.esql.core.util.ReflectionUtils;
-import org.elasticsearch.xpack.esql.optimizer.rules.OptimizerRules.TransformDirection;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
 public class PhysicalOptimizerRules {

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

@@ -9,32 +9,16 @@
 
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.common.Failure;
-import org.elasticsearch.xpack.esql.core.expression.Alias;
-import org.elasticsearch.xpack.esql.core.expression.Attribute;
-import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
-import org.elasticsearch.xpack.esql.core.expression.Expressions;
-import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.rule.ParameterizedRuleExecutor;
-import org.elasticsearch.xpack.esql.core.rule.Rule;
 import org.elasticsearch.xpack.esql.core.rule.RuleExecutor;
-import org.elasticsearch.xpack.esql.core.tree.Source;
-import org.elasticsearch.xpack.esql.core.util.Holder;
-import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
-import org.elasticsearch.xpack.esql.plan.logical.Eval;
-import org.elasticsearch.xpack.esql.plan.logical.Project;
-import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.ProjectAwayColumns;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
-import org.elasticsearch.xpack.esql.plan.physical.UnaryExec;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 
-import static java.lang.Boolean.FALSE;
-import static java.lang.Boolean.TRUE;
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonList;
 
 /**
  * This class is part of the planner. Performs global (coordinator) optimization of the physical plan. Local (data-node) optimizations
@@ -70,64 +54,4 @@ static List<RuleExecutor.Batch<PhysicalPlan>> initializeRules(boolean isOptimize
     protected Iterable<RuleExecutor.Batch<PhysicalPlan>> batches() {
         return rules;
     }
-
-    /**
-     * Adds an explicit project to minimize the amount of attributes sent from the local plan to the coordinator.
-     * This is done here to localize the project close to the data source and simplify the upcoming field
-     * extraction.
-     */
-    static class ProjectAwayColumns extends Rule<PhysicalPlan, PhysicalPlan> {
-
-        @Override
-        public PhysicalPlan apply(PhysicalPlan plan) {
-            Holder<Boolean> keepTraversing = new Holder<>(TRUE);
-            // Invariant: if we add a projection with these attributes after the current plan node, the plan remains valid
-            // and the overall output will not change.
-            Holder<AttributeSet> requiredAttributes = new Holder<>(plan.outputSet());
-
-            // This will require updating should we choose to have non-unary execution plans in the future.
-            return plan.transformDown(UnaryExec.class, currentPlanNode -> {
-                if (keepTraversing.get() == false) {
-                    return currentPlanNode;
-                }
-                if (currentPlanNode instanceof ExchangeExec exec) {
-                    keepTraversing.set(FALSE);
-                    var child = exec.child();
-                    // otherwise expect a Fragment
-                    if (child instanceof FragmentExec fragmentExec) {
-                        var logicalFragment = fragmentExec.fragment();
-
-                        // no need for projection when dealing with aggs
-                        if (logicalFragment instanceof Aggregate == false) {
-                            List<Attribute> output = new ArrayList<>(requiredAttributes.get());
-                            // if all the fields are filtered out, it's only the count that matters
-                            // however until a proper fix (see https://github.com/elastic/elasticsearch/issues/98703)
-                            // 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);
-                                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)
-                            FragmentExec newChild = new FragmentExec(
-                                Source.EMPTY,
-                                new Project(logicalFragment.source(), logicalFragment, output),
-                                fragmentExec.esFilter(),
-                                fragmentExec.estimatedRowSize(),
-                                fragmentExec.reducer()
-                            );
-                            return new ExchangeExec(exec.source(), output, exec.inBetweenAggs(), newChild);
-                        }
-                    }
-                } else {
-                    AttributeSet childOutput = currentPlanNode.inputSet();
-                    AttributeSet addedAttributes = currentPlanNode.outputSet().subtract(childOutput);
-                    requiredAttributes.set(requiredAttributes.get().subtract(addedAttributes).combine(currentPlanNode.references()));
-                }
-                return currentPlanNode;
-            });
-        }
-    }
 }

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

@@ -10,7 +10,7 @@
 import org.elasticsearch.xpack.esql.common.Failure;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
-import org.elasticsearch.xpack.esql.optimizer.OptimizerRules.DependencyConsistency;
+import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
@@ -24,7 +24,7 @@
 public final class PhysicalVerifier {
 
     public static final PhysicalVerifier INSTANCE = new PhysicalVerifier();
-    private static final DependencyConsistency<PhysicalPlan> DEPENDENCY_CHECK = new DependencyConsistency<>();
+    private static final PlanConsistencyChecker<PhysicalPlan> DEPENDENCY_CHECK = new PlanConsistencyChecker<>();
 
     private PhysicalVerifier() {}
 

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

@@ -0,0 +1,48 @@
+/*
+ * 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.optimizer.rules;
+
+import org.elasticsearch.xpack.esql.common.Failures;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
+import org.elasticsearch.xpack.esql.core.expression.NameId;
+import org.elasticsearch.xpack.esql.plan.QueryPlan;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import static org.elasticsearch.xpack.esql.common.Failure.fail;
+
+public class PlanConsistencyChecker<P extends QueryPlan<P>> {
+
+    /**
+     * Check whether a single {@link QueryPlan} produces no duplicate attributes and its children provide all of its required
+     * {@link QueryPlan#references() references}. Otherwise, add
+     * {@link org.elasticsearch.xpack.esql.common.Failure Failure}s to the {@link Failures} object.
+     */
+    public void checkPlan(P p, Failures failures) {
+        AttributeSet refs = p.references();
+        AttributeSet input = p.inputSet();
+        AttributeSet missing = refs.subtract(input);
+        // TODO: for Joins, we should probably check if the required fields from the left child are actually in the left child, not
+        // just any child (and analogously for the right child).
+        if (missing.isEmpty() == false) {
+            failures.add(fail(p, "Plan [{}] optimized incorrectly due to missing references {}", p.nodeString(), missing));
+        }
+
+        Set<String> outputAttributeNames = new HashSet<>();
+        Set<NameId> outputAttributeIds = new HashSet<>();
+        for (Attribute outputAttr : p.output()) {
+            if (outputAttributeNames.add(outputAttr.name()) == false || outputAttributeIds.add(outputAttr.id()) == false) {
+                failures.add(
+                    fail(p, "Plan [{}] optimized incorrectly due to duplicate output attribute {}", p.nodeString(), outputAttr.toString())
+                );
+            }
+        }
+    }
+}