pass the correct timestamp ref into translate rule

Author: not-napoleon

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

@@ -560,21 +560,31 @@ protected LogicalPlan rule(LogicalPlan plan, AnalyzerContext context) {
          * This function is meant to deal with the implicit timestamp fields that some TS functions use.
          */
         private TimeSeriesAggregate resolveTimeSeriesAggregate(TimeSeriesAggregate timeSeriesAggregate, List<Attribute> childrenOutput) {
-            if (childrenOutput.stream().noneMatch(attr -> attr.name().equals(MetadataAttribute.TIMESTAMP_FIELD))) {
-                // We only need to do something if there isn't a timestamp field in our output
+            Attribute tsAtter = null;
+            for (int i = 0; i < childrenOutput.size(); i++) {
+                if (childrenOutput.get(i).name().equals(MetadataAttribute.TIMESTAMP_FIELD)) {
+                    tsAtter = childrenOutput.get(i);
+                    break;
+                }
+            }
+            if (tsAtter == null) {
+                // if we didn't find a timestamp in the children output, time to do more work
                 Holder<String> tsAttributeName = new Holder<>(MetadataAttribute.TIMESTAMP_FIELD);
+                Holder<Attribute> tsAttribute = new Holder<>();
                 timeSeriesAggregate.forEachExpressionUp(Alias.class, a -> {
                     if (a.child() instanceof Attribute c) {
                         // will this ever not be true?
                         if (c.name().equals(tsAttributeName.get())) {
                             tsAttributeName.set(a.name());
+                            tsAttribute.set(a.toAttribute());
                         }
                     }
                 });
 
                 // Now we know what timestamp is going to be called, replace our UnresolvedAttributes referencing timestamp with that name
                 List<Expression> newGroupings = new ArrayList<>(timeSeriesAggregate.groupings().size());
                 List<NamedExpression> newAggregates = new ArrayList<>(timeSeriesAggregate.aggregates().size());
+                // TODO: Can we just resolve these here? we have the attribute
                 for (int i = 0; i < timeSeriesAggregate.groupings().size(); i++) {
                     newGroupings.add(timeSeriesAggregate.groupings().get(i).transformUp(UnresolvedAttribute.class, ua -> {
                         if (ua.name().equals(MetadataAttribute.TIMESTAMP_FIELD)) {
@@ -594,7 +604,9 @@ private TimeSeriesAggregate resolveTimeSeriesAggregate(TimeSeriesAggregate timeS
                         })
                     );
                 }
-                timeSeriesAggregate = (TimeSeriesAggregate) timeSeriesAggregate.with(newGroupings, newAggregates);
+                timeSeriesAggregate = timeSeriesAggregate.with(newGroupings, newAggregates, tsAttribute.get());
+            } else {
+                timeSeriesAggregate = timeSeriesAggregate.with(timeSeriesAggregate.groupings(), timeSeriesAggregate.aggregates(), tsAtter);
             }
 
             // After correcting the timestamps, we still need to resolve the node as normal, so delegate to resolveAggregate

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

@@ -165,15 +165,12 @@ LogicalPlan translate(TimeSeriesAggregate aggregate) {
                 if (attr.name().equals(MetadataAttribute.TSID_FIELD)) {
                     tsid.set(attr);
                 }
-                if (attr.name().equals(MetadataAttribute.TIMESTAMP_FIELD)) {
-                    timestamp.set(attr);
-                }
             }
         });
         if (tsid.get() == null) {
             tsid.set(new MetadataAttribute(aggregate.source(), MetadataAttribute.TSID_FIELD, DataType.KEYWORD, false));
         }
-        if (timestamp.get() == null) {
+        if (aggregate.timestamp() == null) {
             throw new IllegalArgumentException("_tsid or @timestamp field are missing from the time-series source");
         }
         Map<AggregateFunction, Alias> timeSeriesAggs = new HashMap<>();
@@ -216,7 +213,7 @@ LogicalPlan translate(TimeSeriesAggregate aggregate) {
                 } else {
                     // TODO: reject over_time_aggregation only
                     final Expression aggField = af.field();
-                    var tsAgg = new LastOverTime(af.source(), aggField, timestamp.get());
+                    var tsAgg = new LastOverTime(af.source(), aggField, aggregate.timestamp());
                     final AggregateFunction firstStageFn;
                     if (inlineFilter != null) {
                         firstStageFn = tsAgg.perTimeSeriesAggregation().withFilter(inlineFilter);
@@ -247,12 +244,12 @@ LogicalPlan translate(TimeSeriesAggregate aggregate) {
         Holder<NamedExpression> timeBucketRef = new Holder<>();
         aggregate.child().forEachExpressionUp(NamedExpression.class, e -> {
             for (Expression child : e.children()) {
-                if (child instanceof Bucket bucket && bucket.field().equals(timestamp.get())) {
+                if (child instanceof Bucket bucket && bucket.field().equals(aggregate.timestamp())) {
                     if (timeBucketRef.get() != null) {
                         throw new IllegalArgumentException("expected at most one time bucket");
                     }
                     timeBucketRef.set(e);
-                } else if (child instanceof TBucket tbucket && tbucket.field().equals(timestamp.get())) {
+                } else if (child instanceof TBucket tbucket && tbucket.field().equals(aggregate.timestamp())) {
                     if (timeBucketRef.get() != null) {
                         throw new IllegalArgumentException("expected at most one time tbucket");
                     }
@@ -296,7 +293,8 @@ LogicalPlan translate(TimeSeriesAggregate aggregate) {
             newChild,
             firstPassGroupings,
             mergeExpressions(firstPassAggs, firstPassGroupings),
-            (Bucket) Alias.unwrap(timeBucket)
+            (Bucket) Alias.unwrap(timeBucket),
+            aggregate.timestamp()
         );
         return new Aggregate(firstPhase.source(), firstPhase, secondPassGroupings, mergeExpressions(secondPassAggs, secondPassGroupings));
     }

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

@@ -362,7 +362,7 @@ public PlanFactory visitStatsCommand(EsqlBaseParser.StatsCommandContext ctx) {
         return input -> {
             if (input.anyMatch(p -> p instanceof Aggregate) == false
                 && input.anyMatch(p -> p instanceof UnresolvedRelation ur && ur.indexMode() == IndexMode.TIME_SERIES)) {
-                return new TimeSeriesAggregate(source(ctx), input, stats.groupings, stats.aggregates, null);
+                return new TimeSeriesAggregate(source(ctx), input, stats.groupings, stats.aggregates, null, null);
             } else {
                 return new Aggregate(source(ctx), input, stats.groupings, stats.aggregates);
             }

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
@@ -41,21 +42,30 @@ public class TimeSeriesAggregate extends Aggregate {
     );
 
     private final Bucket timeBucket;
+    /**
+     * this field is used by {@link org.elasticsearch.xpack.esql.optimizer.rules.logical.TranslateTimeSeriesAggregate} to help with
+     * resolving the timestamp field, but should not be needed after the initial logical planning on the coordinating node.  As such,
+     * it is not serialized.
+     */
+    private final Attribute timestamp;
 
     public TimeSeriesAggregate(
         Source source,
         LogicalPlan child,
         List<Expression> groupings,
         List<? extends NamedExpression> aggregates,
-        Bucket timeBucket
+        Bucket timeBucket,
+        Attribute timestamp
     ) {
         super(source, child, groupings, aggregates);
         this.timeBucket = timeBucket;
+        this.timestamp = timestamp;
     }
 
     public TimeSeriesAggregate(StreamInput in) throws IOException {
         super(in);
         this.timeBucket = in.readOptionalWriteable(inp -> (Bucket) Bucket.ENTRY.reader.read(inp));
+        this.timestamp = null;
     }
 
     @Override
@@ -71,17 +81,21 @@ public String getWriteableName() {
 
     @Override
     protected NodeInfo<Aggregate> info() {
-        return NodeInfo.create(this, TimeSeriesAggregate::new, child(), groupings, aggregates, timeBucket);
+        return NodeInfo.create(this, TimeSeriesAggregate::new, child(), groupings, aggregates, timeBucket, timestamp);
     }
 
     @Override
     public TimeSeriesAggregate replaceChild(LogicalPlan newChild) {
-        return new TimeSeriesAggregate(source(), newChild, groupings, aggregates, timeBucket);
+        return new TimeSeriesAggregate(source(), newChild, groupings, aggregates, timeBucket, timestamp);
     }
 
     @Override
     public TimeSeriesAggregate with(LogicalPlan child, List<Expression> newGroupings, List<? extends NamedExpression> newAggregates) {
-        return new TimeSeriesAggregate(source(), child, newGroupings, newAggregates, timeBucket);
+        return new TimeSeriesAggregate(source(), child, newGroupings, newAggregates, timeBucket, timestamp);
+    }
+
+    public TimeSeriesAggregate with(List<Expression> newGroupings, List<? extends NamedExpression> newAggregates, Attribute newTimestamp) {
+        return new TimeSeriesAggregate(source(), child(), newGroupings, newAggregates, timeBucket, newTimestamp);
     }
 
     @Override
@@ -94,6 +108,10 @@ public Bucket timeBucket() {
         return timeBucket;
     }
 
+    public Attribute timestamp() {
+        return timestamp;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(groupings, aggregates, child(), timeBucket);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/TranslateTimeSeriesAggregateTests.java

@@ -187,9 +187,9 @@ public void testRateWithRename() {
 
         TimeSeriesAggregate innerStats = as(outerStats.child(), TimeSeriesAggregate.class);
         assertEquals(2, innerStats.groupings().size());
-        assertEquals(timeBucketGroup, innerStats.groupings().get(1));
         Attribute tsidGroup = as(innerStats.groupings().get(0), Attribute.class);
         assertEquals("_tsid", tsidGroup.name());
+        assertEquals(timeBucketGroup, innerStats.groupings().get(1));
 
         assertEquals(2, innerStats.aggregates().size());
         Alias innerAggFunction = as(innerStats.aggregates().get(0), Alias.class);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/StatementParserTests.java

@@ -2496,6 +2496,7 @@ public void testSimpleMetricsWithStats() {
                     new Alias(EMPTY, "load", new UnresolvedFunction(EMPTY, "avg", DEFAULT, List.of(attribute("cpu")))),
                     attribute("ts")
                 ),
+                null,
                 null
             )
         );
@@ -2509,6 +2510,7 @@ public void testSimpleMetricsWithStats() {
                     new Alias(EMPTY, "load", new UnresolvedFunction(EMPTY, "avg", DEFAULT, List.of(attribute("cpu")))),
                     attribute("ts")
                 ),
+                null,
                 null
             )
         );
@@ -2532,6 +2534,7 @@ public void testSimpleMetricsWithStats() {
                     ),
                     attribute("ts")
                 ),
+                null,
                 null
             )
         );
@@ -2542,6 +2545,7 @@ public void testSimpleMetricsWithStats() {
                 unresolvedTSRelation("foo*"),
                 List.of(),
                 List.of(new Alias(EMPTY, "count(errors)", new UnresolvedFunction(EMPTY, "count", DEFAULT, List.of(attribute("errors"))))),
+                null,
                 null
             )
         );
@@ -2552,6 +2556,7 @@ public void testSimpleMetricsWithStats() {
                 unresolvedTSRelation("foo*"),
                 List.of(),
                 List.of(new Alias(EMPTY, "a(b)", new UnresolvedFunction(EMPTY, "a", DEFAULT, List.of(attribute("b"))))),
+                null,
                 null
             )
         );
@@ -2562,6 +2567,7 @@ public void testSimpleMetricsWithStats() {
                 unresolvedTSRelation("foo*"),
                 List.of(),
                 List.of(new Alias(EMPTY, "a(b)", new UnresolvedFunction(EMPTY, "a", DEFAULT, List.of(attribute("b"))))),
+                null,
                 null
             )
         );
@@ -2572,6 +2578,7 @@ public void testSimpleMetricsWithStats() {
                 unresolvedTSRelation("foo*"),
                 List.of(),
                 List.of(new Alias(EMPTY, "a1(b2)", new UnresolvedFunction(EMPTY, "a1", DEFAULT, List.of(attribute("b2"))))),
+                null,
                 null
             )
         );
@@ -2586,6 +2593,7 @@ public void testSimpleMetricsWithStats() {
                     attribute("c"),
                     attribute("d.e")
                 ),
+                null,
                 null
             )
         );

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/TimeSeriesAggregateSerializationTests.java

@@ -24,7 +24,7 @@ protected TimeSeriesAggregate createTestInstance() {
         List<Expression> groupings = randomFieldAttributes(0, 5, false).stream().map(a -> (Expression) a).toList();
         List<? extends NamedExpression> aggregates = AggregateSerializationTests.randomAggregates();
         Bucket timeBucket = BucketSerializationTests.createRandomBucket();
-        return new TimeSeriesAggregate(source, child, groupings, aggregates, timeBucket);
+        return new TimeSeriesAggregate(source, child, groupings, aggregates, timeBucket, null);
     }
 
     @Override
@@ -42,7 +42,7 @@ protected TimeSeriesAggregate mutateInstance(TimeSeriesAggregate instance) throw
             case 2 -> aggregates = randomValueOtherThan(aggregates, AggregateSerializationTests::randomAggregates);
             case 3 -> timeBucket = randomValueOtherThan(timeBucket, BucketSerializationTests::createRandomBucket);
         }
-        return new TimeSeriesAggregate(instance.source(), child, groupings, aggregates, timeBucket);
+        return new TimeSeriesAggregate(instance.source(), child, groupings, aggregates, timeBucket, null);
     }
 
     @Override