Made TopNAggregate a child of Aggregate

Author: ivancea

x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats_top_n.csv-spec

@@ -44,3 +44,14 @@ height:double
 2.1
 2.1
 ;
+
+multipleTopN
+from employees
+| stats c = count(height) by languages.long
+| sort languages.long | limit 4
+| sort languages.long desc | limit 2;
+
+c:long | languages.long:long
+17     | 3
+18     | 4
+;

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

@@ -30,7 +30,7 @@ public LogicalPlan apply(LogicalPlan plan) {
     }
 
     private LogicalPlan applyRule(TopN topN) {
-        if (topN.child() instanceof Aggregate aggregate) {
+        if (topN.child() instanceof Aggregate aggregate && aggregate instanceof TopNAggregate == false) {
             // TimeSeriesAggregate shouldn't appear after a TopN when this rule is executed
             assert aggregate instanceof TimeSeriesAggregate == false : "TimeSeriesAggregate should not be replaced with TopNAggregate";
 

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

@@ -76,7 +76,7 @@ public PhysicalPlan apply(PhysicalPlan plan) {
                     var logicalFragment = fragmentExec.fragment();
 
                     // no need for projection when dealing with aggs
-                    if (logicalFragment instanceof Aggregate == false && logicalFragment instanceof TopNAggregate == false) {
+                    if (logicalFragment instanceof Aggregate == false) {
                         List<Attribute> output = new ArrayList<>(requiredAttrBuilder.build());
                         // 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)

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
@@ -18,6 +19,7 @@
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * An extension of {@link Aggregate} to perform time-series aggregation per time-series, such as rate or _over_time.
@@ -74,8 +76,35 @@ public TimeSeriesAggregate with(LogicalPlan child, List<Expression> newGroupings
         return new TimeSeriesAggregate(source(), child, newGroupings, newAggregates, timeBucket);
     }
 
+    @Override
+    public boolean expressionsResolved() {
+        return super.expressionsResolved() && timeBucket.resolved();
+    }
+
     @Nullable
     public Bucket timeBucket() {
         return timeBucket;
     }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(groupings, aggregates, child(), timeBucket);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+
+        if (obj == null || getClass() != obj.getClass()) {
+            return false;
+        }
+
+        TimeSeriesAggregate other = (TimeSeriesAggregate) obj;
+        return Objects.equals(groupings, other.groupings)
+            && Objects.equals(aggregates, other.aggregates)
+            && Objects.equals(child(), other.child())
+            && Objects.equals(timeBucket, other.timeBucket);
+    }
 }

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

@@ -23,17 +23,13 @@
 import java.util.List;
 import java.util.Objects;
 
-// TODO: Should this be TelemetryAware?
-public class TopNAggregate extends UnaryPlan {
+public class TopNAggregate extends Aggregate {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "TopNAggregate",
         TopNAggregate::new
     );
 
-    protected final List<Expression> groupings;
-    protected final List<? extends NamedExpression> aggregates;
-
     private final List<Order> order;
     private final Expression limit;
 
@@ -47,27 +43,20 @@ public TopNAggregate(
         List<Order> order,
         Expression limit
     ) {
-        super(source, child);
-        this.groupings = groupings;
-        this.aggregates = aggregates;
+        super(source, child, groupings, aggregates);
         this.order = order;
         this.limit = limit;
     }
 
     public TopNAggregate(StreamInput in) throws IOException {
-        super(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(LogicalPlan.class));
-        this.groupings = in.readNamedWriteableCollectionAsList(Expression.class);
-        this.aggregates = in.readNamedWriteableCollectionAsList(NamedExpression.class);
+        super(in);
         this.order = in.readCollectionAsList(Order::new);
         this.limit = in.readNamedWriteable(Expression.class);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        Source.EMPTY.writeTo(out);
-        out.writeNamedWriteable(child());
-        out.writeNamedWriteableCollection(groupings);
-        out.writeNamedWriteableCollection(aggregates());
+        super.writeTo(out);
         out.writeCollection(order);
         out.writeNamedWriteable(limit);
     }
@@ -87,14 +76,6 @@ public TopNAggregate replaceChild(LogicalPlan newChild) {
         return new TopNAggregate(source(), newChild, groupings, aggregates, order, limit);
     }
 
-    public List<Expression> groupings() {
-        return groupings;
-    }
-
-    public List<? extends NamedExpression> aggregates() {
-        return aggregates;
-    }
-
     public List<Order> order() {
         return order;
     }
@@ -105,20 +86,7 @@ public Expression limit() {
 
     @Override
     public boolean expressionsResolved() {
-        return Resolvables.resolved(groupings) && Resolvables.resolved(aggregates);
-    }
-
-    @Override
-    public List<Attribute> output() {
-        if (lazyOutput == null) {
-            lazyOutput = Aggregate.output(aggregates);
-        }
-        return lazyOutput;
-    }
-
-    @Override
-    protected AttributeSet computeReferences() {
-        return Aggregate.computeReferences(aggregates, groupings);
+        return super.expressionsResolved() && Resolvables.resolved(order) && limit.resolved();
     }
 
     @Override

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

@@ -73,16 +73,16 @@ private PhysicalPlan mapUnary(UnaryPlan unary) {
         // Pipeline breakers
         //
 
-        if (unary instanceof Aggregate aggregate) {
-            List<Attribute> intermediate = MapperUtils.intermediateAttributes(aggregate);
-            return MapperUtils.aggExec(aggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
-        }
-
         if (unary instanceof TopNAggregate topNAggregate) {
             List<Attribute> intermediate = MapperUtils.intermediateAttributes(topNAggregate);
             return MapperUtils.topNAggExec(topNAggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
         }
 
+        if (unary instanceof Aggregate aggregate) {
+            List<Attribute> intermediate = MapperUtils.intermediateAttributes(aggregate);
+            return MapperUtils.aggExec(aggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
+        }
+
         if (unary instanceof Limit limit) {
             return new LimitExec(limit.source(), mappedChild, limit.limit(), null);
         }

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

@@ -146,7 +146,7 @@ private PhysicalPlan mapUnary(UnaryPlan unary) {
         //
         // Pipeline breakers
         //
-        if (unary instanceof Aggregate aggregate) {
+        if (unary instanceof TopNAggregate aggregate) {
             List<Attribute> intermediate = MapperUtils.intermediateAttributes(aggregate);
 
             // create both sides of the aggregate (for parallelism purposes), if no fragment is present
@@ -160,14 +160,14 @@ private PhysicalPlan mapUnary(UnaryPlan unary) {
             }
             // if no exchange was added (aggregation happening on the coordinator), create the initial agg
             else {
-                mappedChild = MapperUtils.aggExec(aggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
+                mappedChild = MapperUtils.topNAggExec(aggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
             }
 
             // always add the final/reduction agg
-            return MapperUtils.aggExec(aggregate, mappedChild, AggregatorMode.FINAL, intermediate);
+            return MapperUtils.topNAggExec(aggregate, mappedChild, AggregatorMode.FINAL, intermediate);
         }
 
-        if (unary instanceof TopNAggregate aggregate) {
+        if (unary instanceof Aggregate aggregate) {
             List<Attribute> intermediate = MapperUtils.intermediateAttributes(aggregate);
 
             // create both sides of the aggregate (for parallelism purposes), if no fragment is present
@@ -181,11 +181,11 @@ private PhysicalPlan mapUnary(UnaryPlan unary) {
             }
             // if no exchange was added (aggregation happening on the coordinator), create the initial agg
             else {
-                mappedChild = MapperUtils.topNAggExec(aggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
+                mappedChild = MapperUtils.aggExec(aggregate, mappedChild, AggregatorMode.INITIAL, intermediate);
             }
 
             // always add the final/reduction agg
-            return MapperUtils.topNAggExec(aggregate, mappedChild, AggregatorMode.FINAL, intermediate);
+            return MapperUtils.aggExec(aggregate, mappedChild, AggregatorMode.FINAL, intermediate);
         }
 
         if (unary instanceof Limit limit) {

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

@@ -152,14 +152,6 @@ static List<Attribute> intermediateAttributes(Aggregate aggregate) {
         return intermediateAttributes;
     }
 
-    static List<Attribute> intermediateAttributes(TopNAggregate aggregate) {
-        List<Attribute> intermediateAttributes = AbstractPhysicalOperationProviders.intermediateAttributes(
-            aggregate.aggregates(),
-            aggregate.groupings()
-        );
-        return intermediateAttributes;
-    }
-
     static AggregateExec aggExec(Aggregate aggregate, PhysicalPlan child, AggregatorMode aggMode, List<Attribute> intermediateAttributes) {
         if (aggregate instanceof TimeSeriesAggregate ts) {
             return new TimeSeriesAggregateExec(