Run single phase aggregation when possible

Author: dnhatn

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.ProjectAwayColumns;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.SinglePhaseAggregate;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRuleExecutor;
@@ -24,7 +25,7 @@
 public class PhysicalPlanOptimizer extends ParameterizedRuleExecutor<PhysicalPlan, PhysicalOptimizerContext> {
 
     private static final List<RuleExecutor.Batch<PhysicalPlan>> RULES = List.of(
-        new Batch<>("Plan Boundary", Limiter.ONCE, new ProjectAwayColumns())
+        new Batch<>("Plan Boundary", Limiter.ONCE, new ProjectAwayColumns(), new SinglePhaseAggregate())
     );
 
     private final PhysicalVerifier verifier = PhysicalVerifier.INSTANCE;

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

@@ -0,0 +1,31 @@
+/*
+ * 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.physical;
+
+import org.elasticsearch.compute.aggregation.AggregatorMode;
+import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
+import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
+import org.elasticsearch.xpack.esql.rule.Rule;
+
+/**
+ * Collapses two-phase aggregation into a single phase when possible.
+ * For example, in FROM .. | STATS first | STATS second, the STATS second aggregation
+ * can be executed in a single phase on the coordinator instead of two phases.
+ */
+public class SinglePhaseAggregate extends Rule<PhysicalPlan, PhysicalPlan> {
+    @Override
+    public PhysicalPlan apply(PhysicalPlan plan) {
+        if (plan instanceof AggregateExec parent
+            && parent.getMode() == AggregatorMode.INITIAL
+            && parent.child() instanceof AggregateExec child
+            && child.getMode() == AggregatorMode.FINAL) {
+            return parent.withMode(AggregatorMode.SINGLE).replaceChild(child.child());
+        }
+        return plan;
+    }
+}