Revert "WIP: TopNAggregate nodes, exec not fully implemented yet"

This reverts commit ee00361.

Author: ivancea

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

@@ -14,7 +14,6 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.InferNonNullAggConstraint;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.LocalPropagateEmptyRelation;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceFieldWithConstantOrNull;
-import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceTopNAggregateWithTopNAndAggregate;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceTopNWithLimitAndSort;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRuleExecutor;
@@ -39,7 +38,6 @@ public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<Logical
         new Batch<>(
             "Local rewrite",
             Limiter.ONCE,
-            new ReplaceTopNAggregateWithTopNAndAggregate(),
             new ReplaceTopNWithLimitAndSort(),
             new ReplaceFieldWithConstantOrNull(),
             new InferIsNotNull(),

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

@@ -28,7 +28,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateEvalFoldables;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateInlineEvals;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateNullable;
-import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceTopNAndAggregateWithTopNAggregate;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateTopNToAggregates;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropgateUnmappedFields;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyPlans;
@@ -209,6 +209,6 @@ protected static Batch<LogicalPlan> operators() {
     }
 
     protected static Batch<LogicalPlan> cleanup() {
-        return new Batch<>("Clean Up", new ReplaceLimitAndSortAsTopN(), new ReplaceTopNAndAggregateWithTopNAggregate(), new ReplaceRowAsLocalRelation(), new PropgateUnmappedFields());
+        return new Batch<>("Clean Up", new ReplaceLimitAndSortAsTopN(), new PropagateTopNToAggregates(), new ReplaceRowAsLocalRelation(), new PropgateUnmappedFields());
     }
 }

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.logical.PropagateTopNToAggregates;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRuleExecutor;

…ceTopNAndAggregateWithTopNAggregate.java …s/logical/PropagateTopNToAggregates.javax-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/ReplaceTopNAndAggregateWithTopNAggregate.java renamed to x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PropagateTopNToAggregates.java x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/ReplaceTopNAndAggregateWithTopNAggregate.java renamed to x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PropagateTopNToAggregates.java

@@ -10,18 +10,18 @@
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
-import org.elasticsearch.xpack.esql.plan.logical.TopNAggregate;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 /**
  * Looks for the structure:
  * <pre>
- * {@link TopN}
- * \_{@link Aggregate}
+ * TopN
+ * \_Aggregate
  * </pre>
- * And replaces it with {@link TopNAggregate}.
+ * And replaces the Aggregate with an Aggregate with the TopN data.
+ * (TODO: Create a new TopNAggregate node instead)
  */
-public class ReplaceTopNAndAggregateWithTopNAggregate extends Rule<TopN, LogicalPlan> {
+public class PropagateTopNToAggregates extends Rule<TopN, LogicalPlan> {
 
     @Override
     public LogicalPlan apply(LogicalPlan plan) {
@@ -31,16 +31,18 @@ public LogicalPlan apply(LogicalPlan plan) {
         );
     }
 
-    private LogicalPlan applyRule(TopN topN) {
+    private TopN applyRule(TopN topN) {
         // TODO: Handle TimeSeriesAggregate
         if (topN.child() instanceof Aggregate aggregate) {
-            return new TopNAggregate(
-                aggregate.source(),
-                aggregate.child(),
-                aggregate.groupings(),
-                aggregate.aggregates(),
-                topN.order(),
-                topN.limit()
+            return topN.replaceChild(
+                new Aggregate(
+                    aggregate.source(),
+                    aggregate.child(),
+                    aggregate.groupings(),
+                    aggregate.aggregates(),
+                    topN.order(),
+                    topN.limit()
+                )
             );
         }
         return topN;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceTopNAggregateWithTopNAndAggregate.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.xpack.esql.plan.logical.Sample;
 import org.elasticsearch.xpack.esql.plan.logical.TimeSeriesAggregate;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
-import org.elasticsearch.xpack.esql.plan.logical.TopNAggregate;
 import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
 import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
 import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
@@ -53,7 +52,6 @@
 import org.elasticsearch.xpack.esql.plan.physical.ShowExec;
 import org.elasticsearch.xpack.esql.plan.physical.SubqueryExec;
 import org.elasticsearch.xpack.esql.plan.physical.TimeSeriesAggregateExec;
-import org.elasticsearch.xpack.esql.plan.physical.TopNAggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
 import org.elasticsearch.xpack.esql.plan.physical.inference.CompletionExec;
 import org.elasticsearch.xpack.esql.plan.physical.inference.RerankExec;
@@ -93,8 +91,7 @@ public static List<NamedWriteableRegistry.Entry> logical() {
             Rerank.ENTRY,
             Sample.ENTRY,
             TimeSeriesAggregate.ENTRY,
-            TopN.ENTRY,
-            TopNAggregate.ENTRY
+            TopN.ENTRY
         );
     }
 
@@ -124,8 +121,7 @@ public static List<NamedWriteableRegistry.Entry> physical() {
             ShowExec.ENTRY,
             SubqueryExec.ENTRY,
             TimeSeriesAggregateExec.ENTRY,
-            TopNExec.ENTRY,
-            TopNAggregateExec.ENTRY
+            TopNExec.ENTRY
         );
     }
 }

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

@@ -89,7 +89,6 @@ public AttributeSet references() {
     /**
      * This very likely needs to be overridden for {@link QueryPlan#references} to be correct when inheriting.
      * This can be called on unresolved plans and therefore must not rely on calls to {@link QueryPlan#output()}.
-     * @see #references()
      */
     protected AttributeSet computeReferences() {
         return Expressions.references(expressions());

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
@@ -27,6 +28,7 @@
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.util.Holder;
+import org.elasticsearch.xpack.esql.expression.Order;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.FilteredExpression;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Rate;
@@ -55,12 +57,32 @@ public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAwar
     protected final List<Expression> groupings;
     protected final List<? extends NamedExpression> aggregates;
 
+    private final List<Order> order;
+    @Nullable
+    private final Expression limit;
+
     protected List<Attribute> lazyOutput;
 
     public Aggregate(Source source, LogicalPlan child, List<Expression> groupings, List<? extends NamedExpression> aggregates) {
         super(source, child);
         this.groupings = groupings;
         this.aggregates = aggregates;
+        this.order = List.of();
+        this.limit = null;
+    }
+
+    public Aggregate(
+        Source source,
+        LogicalPlan child,
+        List<Expression> groupings, List<? extends NamedExpression> aggregates,
+        List<Order> order,
+        @Nullable Expression limit
+    ) {
+        super(source, child);
+        this.groupings = groupings;
+        this.aggregates = aggregates;
+        this.order = order;
+        this.limit = limit;
     }
 
     public Aggregate(StreamInput in) throws IOException {
@@ -71,6 +93,13 @@ public Aggregate(StreamInput in) throws IOException {
         }
         this.groupings = in.readNamedWriteableCollectionAsList(Expression.class);
         this.aggregates = in.readNamedWriteableCollectionAsList(NamedExpression.class);
+        if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_TOP_N_AGGREGATES)) {
+            this.order = in.readCollectionAsList(Order::new);
+            this.limit = in.readOptionalNamedWriteable(Expression.class);
+        } else {
+            this.order = emptyList();
+            this.limit = null;
+        }
     }
 
     @Override
@@ -83,6 +112,10 @@ public void writeTo(StreamOutput out) throws IOException {
         }
         out.writeNamedWriteableCollection(groupings);
         out.writeNamedWriteableCollection(aggregates());
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_TOP_N_AGGREGATES)) {
+            out.writeCollection(order);
+            out.writeOptionalNamedWriteable(limit);
+        }
     }
 
     @Override
@@ -97,15 +130,15 @@ protected NodeInfo<? extends Aggregate> info() {
 
     @Override
     public Aggregate replaceChild(LogicalPlan newChild) {
-        return new Aggregate(source(), newChild, groupings, aggregates);
+        return new Aggregate(source(), newChild, groupings, aggregates, order, limit);
     }
 
     public Aggregate with(List<Expression> newGroupings, List<? extends NamedExpression> newAggregates) {
         return with(child(), newGroupings, newAggregates);
     }
 
     public Aggregate with(LogicalPlan child, List<Expression> newGroupings, List<? extends NamedExpression> newAggregates) {
-        return new Aggregate(source(), child, newGroupings, newAggregates);
+        return new Aggregate(source(), child, newGroupings, newAggregates, order, limit);
     }
 
     public List<Expression> groupings() {
@@ -116,6 +149,14 @@ public List<? extends NamedExpression> aggregates() {
         return aggregates;
     }
 
+    public List<Order> order() {
+        return order;
+    }
+
+    public Expression limit() {
+        return limit;
+    }
+
     @Override
     public String telemetryLabel() {
         return "STATS";