Add logical/physical plans time-series aggregate (#126178)

We need to store extra information for Aggregate and AggregateExec for 
time-series aggregations. Previously, I added a type (standard or
time_series), but this was not enough. This PR removes it and replaces
it with extensions of Aggregate and AggregateExec. I considered adding
an extra map of metadata to Aggregate and AggregateExec, but this
approach seems simpler.

Author: dnhatn

Diff for: server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -211,6 +211,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion REPOSITORIES_METADATA_AS_PROJECT_CUSTOM = def(9_042_0_00);
     public static final TransportVersion BATCHED_QUERY_PHASE_VERSION = def(9_043_0_00);
     public static final TransportVersion REMOTE_EXCEPTION = def(9_044_0_00);
+    public static final TransportVersion ESQL_REMOVE_AGGREGATE_TYPE = def(9_045_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

Diff for: x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/TimeSeriesAggregationOperator.java

@@ -23,19 +23,18 @@
 public class TimeSeriesAggregationOperator extends HashAggregationOperator {
 
     public record Factory(
-        BlockHash.GroupSpec tsidGroup,
-        BlockHash.GroupSpec timestampGroup,
+        List<BlockHash.GroupSpec> groups,
         AggregatorMode aggregatorMode,
         List<GroupingAggregator.Factory> aggregators,
         int maxPageSize
     ) implements OperatorFactory {
         @Override
         public Operator get(DriverContext driverContext) {
             // TODO: use TimeSeriesBlockHash when possible
-            return new HashAggregationOperator(
+            return new TimeSeriesAggregationOperator(
                 aggregators,
                 () -> BlockHash.build(
-                    List.of(tsidGroup, timestampGroup),
+                    groups,
                     driverContext.blockFactory(),
                     maxPageSize,
                     true // we can enable optimizations as the inputs are vectors

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/CombineProjections.java

@@ -53,7 +53,7 @@ protected LogicalPlan rule(UnaryPlan plan) {
                 // project can be fully removed
                 if (newAggs != null) {
                     var newGroups = replacePrunedAliasesUsedInGroupBy(a.groupings(), aggs, newAggs);
-                    plan = new Aggregate(a.source(), a.child(), a.aggregateType(), newGroups, newAggs);
+                    plan = a.with(newGroups, newAggs);
                 }
             }
             return plan;
@@ -75,10 +75,8 @@ protected LogicalPlan rule(UnaryPlan plan) {
                         throw new EsqlIllegalArgumentException("Expected an Attribute, got {}", grouping);
                     }
                 }
-                plan = new Aggregate(
-                    a.source(),
+                plan = a.with(
                     p.child(),
-                    a.aggregateType(),
                     combineUpperGroupingsAndLowerProjections(groupingAttrs, p.projections()),
                     combineProjections(a.aggregates(), p.projections())
                 );

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/ExtractAggregateCommonFilter.java

@@ -67,12 +67,6 @@ protected LogicalPlan rule(Aggregate aggregate) {
         }
 
         // build the new agg on top of extracted filter
-        return new Aggregate(
-            aggregate.source(),
-            new Filter(aggregate.source(), aggregate.child(), common.v1()),
-            aggregate.aggregateType(),
-            aggregate.groupings(),
-            newAggs
-        );
+        return aggregate.with(new Filter(aggregate.source(), aggregate.child(), common.v1()), aggregate.groupings(), newAggs);
     }
 }

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java

@@ -73,22 +73,10 @@ public LogicalPlan apply(LogicalPlan plan) {
                             } else {
                                 // Aggs cannot produce pages with 0 columns, so retain one grouping.
                                 remaining = List.of(Expressions.attribute(aggregate.groupings().get(0)));
-                                p = new Aggregate(
-                                    aggregate.source(),
-                                    aggregate.child(),
-                                    aggregate.aggregateType(),
-                                    aggregate.groupings(),
-                                    remaining
-                                );
+                                p = aggregate.with(aggregate.groupings(), remaining);
                             }
                         } else {
-                            p = new Aggregate(
-                                aggregate.source(),
-                                aggregate.child(),
-                                aggregate.aggregateType(),
-                                aggregate.groupings(),
-                                remaining
-                            );
+                            p = aggregate.with(aggregate.groupings(), remaining);
                         }
                     }
                 } else if (p instanceof Eval eval) {

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/SubstituteSurrogates.java

@@ -107,7 +107,7 @@ protected LogicalPlan rule(Aggregate aggregate) {
         if (changed) {
             var source = aggregate.source();
             if (newAggs.isEmpty() == false) {
-                plan = new Aggregate(source, aggregate.child(), aggregate.aggregateType(), aggregate.groupings(), newAggs);
+                plan = aggregate.with(aggregate.child(), aggregate.groupings(), newAggs);
             } else {
                 // All aggs actually have been surrogates for (foldable) expressions, e.g.
                 // \_Aggregate[[],[AVG([1, 2][INTEGER]) AS s]]

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/TranslateTimeSeriesAggregate.java

@@ -26,12 +26,12 @@
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.TimeSeriesAggregate;
 
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Stream;
 
 /**
  * Rate aggregation is special because it must be computed per time series, regardless of the grouping keys.
@@ -123,10 +123,14 @@ public TranslateTimeSeriesAggregate() {
 
     @Override
     protected LogicalPlan rule(Aggregate aggregate) {
-        return translate(aggregate);
+        if (aggregate instanceof TimeSeriesAggregate ts) {
+            return translate(ts);
+        } else {
+            return aggregate;
+        }
     }
 
-    LogicalPlan translate(Aggregate aggregate) {
+    LogicalPlan translate(TimeSeriesAggregate aggregate) {
         Map<Rate, Alias> rateAggs = new HashMap<>();
         List<NamedExpression> firstPassAggs = new ArrayList<>();
         List<NamedExpression> secondPassAggs = new ArrayList<>();
@@ -153,7 +157,8 @@ LogicalPlan translate(Aggregate aggregate) {
             }
         }
         if (rateAggs.isEmpty()) {
-            return aggregate;
+            // no time-series aggregations, run a regular aggregation instead.
+            return new Aggregate(aggregate.source(), aggregate.child(), aggregate.groupings(), aggregate.aggregates());
         }
         Holder<Attribute> tsid = new Holder<>();
         Holder<Attribute> timestamp = new Holder<>();
@@ -204,7 +209,7 @@ LogicalPlan translate(Aggregate aggregate) {
             }
             secondPassGroupings.add(new Alias(g.source(), g.name(), newFinalGroup.toAttribute(), g.id()));
         }
-        LogicalPlan relation = aggregate.child().transformUp(EsRelation.class, r -> {
+        LogicalPlan newChild = aggregate.child().transformUp(EsRelation.class, r -> {
             if (r.output().contains(tsid.get()) == false) {
                 return new EsRelation(
                     r.source(),
@@ -217,26 +222,22 @@ LogicalPlan translate(Aggregate aggregate) {
                 return r;
             }
         });
-        return newAggregate(
-            newAggregate(relation, Aggregate.AggregateType.TIME_SERIES, firstPassAggs, firstPassGroupings),
-            Aggregate.AggregateType.STANDARD,
-            secondPassAggs,
-            secondPassGroupings
+        final var firstPhase = new TimeSeriesAggregate(
+            newChild.source(),
+            newChild,
+            firstPassGroupings,
+            mergeExpressions(firstPassAggs, firstPassGroupings)
         );
+        return new Aggregate(firstPhase.source(), firstPhase, secondPassGroupings, mergeExpressions(secondPassAggs, secondPassGroupings));
     }
 
-    private static Aggregate newAggregate(
-        LogicalPlan child,
-        Aggregate.AggregateType type,
+    private static List<? extends NamedExpression> mergeExpressions(
         List<? extends NamedExpression> aggregates,
         List<Expression> groupings
     ) {
-        return new Aggregate(
-            child.source(),
-            child,
-            type,
-            groupings,
-            Stream.concat(aggregates.stream(), groupings.stream().map(Expressions::attribute)).toList()
-        );
+        List<NamedExpression> merged = new ArrayList<>(aggregates.size() + groupings.size());
+        merged.addAll(aggregates);
+        groupings.forEach(g -> merged.add(Expressions.attribute(g)));
+        return merged;
     }
 }

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -66,6 +66,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Rename;
 import org.elasticsearch.xpack.esql.plan.logical.Row;
 import org.elasticsearch.xpack.esql.plan.logical.RrfScoreEval;
+import org.elasticsearch.xpack.esql.plan.logical.TimeSeriesAggregate;
 import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
 import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
 import org.elasticsearch.xpack.esql.plan.logical.show.ShowInfo;
@@ -312,7 +313,13 @@ public PlanFactory visitInsistCommand(EsqlBaseParser.InsistCommandContext ctx) {
     @Override
     public PlanFactory visitStatsCommand(EsqlBaseParser.StatsCommandContext ctx) {
         final Stats stats = stats(source(ctx), ctx.grouping, ctx.stats);
-        return input -> new Aggregate(source(ctx), input, Aggregate.AggregateType.STANDARD, stats.groupings, stats.aggregates);
+        return input -> {
+            if (input.anyMatch(p -> p instanceof UnresolvedRelation ur && ur.indexMode() == IndexMode.TIME_SERIES)) {
+                return new TimeSeriesAggregate(source(ctx), input, stats.groupings, stats.aggregates);
+            } else {
+                return new Aggregate(source(ctx), input, stats.groupings, stats.aggregates);
+            }
+        };
     }
 
     private record Stats(List<Expression> groupings, List<? extends NamedExpression> aggregates) {}
@@ -362,10 +369,7 @@ public PlanFactory visitInlinestatsCommand(EsqlBaseParser.InlinestatsCommandCont
         List<NamedExpression> groupings = visitGrouping(ctx.grouping);
         aggregates.addAll(groupings);
         // TODO: add support for filters
-        return input -> new InlineStats(
-            source(ctx),
-            new Aggregate(source(ctx), input, Aggregate.AggregateType.STANDARD, new ArrayList<>(groupings), aggregates)
-        );
+        return input -> new InlineStats(source(ctx), new Aggregate(source(ctx), input, new ArrayList<>(groupings), aggregates));
     }
 
     @Override

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
 import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
+import org.elasticsearch.xpack.esql.plan.logical.TimeSeriesAggregate;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
 import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
@@ -46,6 +47,7 @@
 import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
 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.TopNExec;
 
 import java.util.ArrayList;
@@ -79,6 +81,7 @@ public static List<NamedWriteableRegistry.Entry> logical() {
             MvExpand.ENTRY,
             OrderBy.ENTRY,
             Project.ENTRY,
+            TimeSeriesAggregate.ENTRY,
             TopN.ENTRY
         );
     }
@@ -105,6 +108,7 @@ public static List<NamedWriteableRegistry.Entry> physical() {
             ProjectExec.ENTRY,
             ShowExec.ENTRY,
             SubqueryExec.ENTRY,
+            TimeSeriesAggregateExec.ENTRY,
             TopNExec.ENTRY
         );
     }

Diff for: x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Aggregate.java

@@ -48,61 +48,35 @@ public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAwar
         Aggregate::new
     );
 
-    public enum AggregateType {
-        STANDARD,
-        TIME_SERIES;
-
-        static void writeType(StreamOutput out, AggregateType type) throws IOException {
-            if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_15_0)) {
-                out.writeString(type.name());
-            } else if (type != STANDARD) {
-                throw new IllegalStateException("cluster is not ready to support aggregate type [" + type + "]");
-            }
-        }
-
-        static AggregateType readType(StreamInput in) throws IOException {
-            if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_15_0)) {
-                return AggregateType.valueOf(in.readString());
-            } else {
-                return STANDARD;
-            }
-        }
-    }
-
-    private final AggregateType aggregateType;
-    private final List<Expression> groupings;
-    private final List<? extends NamedExpression> aggregates;
+    protected final List<Expression> groupings;
+    protected final List<? extends NamedExpression> aggregates;
 
-    private List<Attribute> lazyOutput;
+    protected List<Attribute> lazyOutput;
 
-    public Aggregate(
-        Source source,
-        LogicalPlan child,
-        AggregateType aggregateType,
-        List<Expression> groupings,
-        List<? extends NamedExpression> aggregates
-    ) {
+    public Aggregate(Source source, LogicalPlan child, List<Expression> groupings, List<? extends NamedExpression> aggregates) {
         super(source, child);
-        this.aggregateType = aggregateType;
         this.groupings = groupings;
         this.aggregates = aggregates;
     }
 
     public Aggregate(StreamInput in) throws IOException {
-        this(
-            Source.readFrom((PlanStreamInput) in),
-            in.readNamedWriteable(LogicalPlan.class),
-            AggregateType.readType(in),
-            in.readNamedWriteableCollectionAsList(Expression.class),
-            in.readNamedWriteableCollectionAsList(NamedExpression.class)
-        );
+        super(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(LogicalPlan.class));
+        if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_15_0)
+            && in.getTransportVersion().before(TransportVersions.ESQL_REMOVE_AGGREGATE_TYPE)) {
+            in.readString();
+        }
+        this.groupings = in.readNamedWriteableCollectionAsList(Expression.class);
+        this.aggregates = in.readNamedWriteableCollectionAsList(NamedExpression.class);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         Source.EMPTY.writeTo(out);
         out.writeNamedWriteable(child());
-        AggregateType.writeType(out, aggregateType());
+        if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_15_0)
+            && out.getTransportVersion().before(TransportVersions.ESQL_REMOVE_AGGREGATE_TYPE)) {
+            out.writeString("STANDARD");
+        }
         out.writeNamedWriteableCollection(groupings);
         out.writeNamedWriteableCollection(aggregates());
     }
@@ -113,25 +87,21 @@ public String getWriteableName() {
     }
 
     @Override
-    protected NodeInfo<Aggregate> info() {
-        return NodeInfo.create(this, Aggregate::new, child(), aggregateType, groupings, aggregates);
+    protected NodeInfo<? extends Aggregate> info() {
+        return NodeInfo.create(this, Aggregate::new, child(), groupings, aggregates);
     }
 
     @Override
     public Aggregate replaceChild(LogicalPlan newChild) {
-        return new Aggregate(source(), newChild, aggregateType, groupings, aggregates);
+        return new Aggregate(source(), newChild, groupings, aggregates);
     }
 
     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, aggregateType(), newGroupings, newAggregates);
-    }
-
-    public AggregateType aggregateType() {
-        return aggregateType;
+        return new Aggregate(source(), child, newGroupings, newAggregates);
     }
 
     public List<Expression> groupings() {
@@ -144,10 +114,7 @@ public List<? extends NamedExpression> aggregates() {
 
     @Override
     public String telemetryLabel() {
-        return switch (aggregateType) {
-            case STANDARD -> "STATS";
-            case TIME_SERIES -> "TIME_SERIES";
-        };
+        return "STATS";
     }
 
     @Override
@@ -184,7 +151,7 @@ public static AttributeSet computeReferences(List<? extends NamedExpression> agg
 
     @Override
     public int hashCode() {
-        return Objects.hash(aggregateType, groupings, aggregates, child());
+        return Objects.hash(groupings, aggregates, child());
     }
 
     @Override
@@ -198,8 +165,7 @@ public boolean equals(Object obj) {
         }
 
         Aggregate other = (Aggregate) obj;
-        return aggregateType == other.aggregateType
-            && Objects.equals(groupings, other.groupings)
+        return Objects.equals(groupings, other.groupings)
             && Objects.equals(aggregates, other.aggregates)
             && Objects.equals(child(), other.child());
     }