Return slice and future timestamp in time-series source (#134767)

With this change, the time-series source operator will return two 
additional constant blocks: one for the slice index and one for the
future max timestamp.

1. For the slice index, the rate can keep the last slice index and flush 
the rate buffer when the slice index changes. This allows partitioning
the time-series source by time periods.


2. For the future max timestamp, we are currently returning a dummy 
Long.MAX_VALUE, but in a follow-up, we will return the max timestamp
from segments. This will enable rate aggregation to reduce the
outstanding buffer and avoid buffering unnecessary data points.

This change is a preparation for future work.

Relates #134324

Author: dnhatn

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/RateDoubleGroupingAggregatorFunction.java

@@ -126,6 +126,11 @@ public final class Rate$Type$GroupingAggregatorFunction implements GroupingAggre
             assert false : "expected timestamp vector in time-series aggregation";
             throw new IllegalStateException("expected timestamp vector in time-series aggregation");
         }
+        IntVector sliceIndices = ((IntBlock) page.getBlock(channels.get(2))).asVector();
+        assert sliceIndices != null : "expected slice indices vector in time-series aggregation";
+        LongVector futureMaxTimestamps = ((LongBlock) page.getBlock(channels.get(3))).asVector();
+        assert futureMaxTimestamps != null : "expected future max timestamps vector in time-series aggregation";
+
         return new AddInput() {
             @Override
             public void add(int positionOffset, IntArrayBlock groupIds) {

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/RateIntGroupingAggregatorFunction.java

@@ -65,7 +65,7 @@ public abstract class LuceneOperator extends SourceOperator {
     final Set<Query> processedQueries = new HashSet<>();
     final Set<String> processedShards = new HashSet<>();
 
-    private LuceneSlice currentSlice;
+    protected LuceneSlice currentSlice;
     private int sliceIndex;
 
     private LuceneScorer currentScorer;

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/RateLongGroupingAggregatorFunction.java

@@ -60,9 +60,9 @@ public class LuceneSourceOperator extends LuceneOperator {
     private final int minPageSize;
 
     public static class Factory extends LuceneOperator.Factory {
-        private final List<? extends RefCounted> contexts;
-        private final int maxPageSize;
-        private final Limiter limiter;
+        protected final List<? extends RefCounted> contexts;
+        protected final int maxPageSize;
+        protected final Limiter limiter;
 
         public Factory(
             List<? extends ShardContext> contexts,
@@ -327,7 +327,8 @@ public Page getCheckedOutput() throws IOException {
                 IntVector shard = null;
                 IntVector leaf = null;
                 IntVector docs = null;
-                Block[] blocks = new Block[1 + (scoreBuilder == null ? 0 : 1) + scorer.tags().size()];
+                int metadataBlocks = numMetadataBlocks();
+                Block[] blocks = new Block[1 + metadataBlocks + scorer.tags().size()];
                 currentPagePos -= discardedDocs;
                 try {
                     int shardId = scorer.shardContext().index();
@@ -340,10 +341,8 @@ public Page getCheckedOutput() throws IOException {
                     shard = null;
                     leaf = null;
                     docs = null;
-                    if (scoreBuilder != null) {
-                        blocks[b++] = buildScoresVector(currentPagePos).asBlock();
-                        scoreBuilder = blockFactory.newDoubleVectorBuilder(Math.min(remainingDocs, maxPageSize));
-                    }
+                    buildMetadataBlocks(blocks, b, currentPagePos);
+                    b += metadataBlocks;
                     for (Object e : scorer.tags()) {
                         blocks[b++] = BlockUtils.constantBlock(blockFactory, e, currentPagePos);
                     }
@@ -393,6 +392,17 @@ private DoubleVector buildScoresVector(int upToPositions) {
         }
     }
 
+    protected int numMetadataBlocks() {
+        return scoreBuilder != null ? 1 : 0;
+    }
+
+    protected void buildMetadataBlocks(Block[] blocks, int offset, int currentPagePos) {
+        if (scoreBuilder != null) {
+            blocks[offset] = buildScoresVector(currentPagePos).asBlock();
+            scoreBuilder = blockFactory.newDoubleVectorBuilder(Math.min(remainingDocs, maxPageSize));
+        }
+    }
+
     @Override
     public void additionalClose() {
         Releasables.close(docsBuilder, scoreBuilder);

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/X-RateGroupingAggregatorFunction.java.st

@@ -0,0 +1,79 @@
+/*
+ * 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.compute.lucene;
+
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.compute.operator.Limiter;
+import org.elasticsearch.compute.operator.SourceOperator;
+import org.elasticsearch.core.RefCounted;
+
+import java.util.List;
+import java.util.function.Function;
+
+/**
+ * Extension of {@link LuceneSourceOperator} for time-series aggregation that inserts metadata blocks,
+ * such as slice index and future max timestamp, to allow downstream operators to optimize processing.
+ */
+public final class TimeSeriesSourceOperator extends LuceneSourceOperator {
+
+    public static final class Factory extends LuceneSourceOperator.Factory {
+        public Factory(
+            List<? extends ShardContext> contexts,
+            Function<ShardContext, List<LuceneSliceQueue.QueryAndTags>> queryFunction,
+            int taskConcurrency,
+            int maxPageSize,
+            int limit
+        ) {
+            super(
+                contexts,
+                queryFunction,
+                DataPartitioning.SHARD,
+                query -> { throw new UnsupportedOperationException("locked to SHARD partitioning"); },
+                taskConcurrency,
+                maxPageSize,
+                limit,
+                false
+            );
+        }
+
+        @Override
+        public SourceOperator get(DriverContext driverContext) {
+            return new TimeSeriesSourceOperator(contexts, driverContext.blockFactory(), maxPageSize, sliceQueue, limit, limiter);
+        }
+
+        @Override
+        public String describe() {
+            return "TimeSeriesSourceOperator[maxPageSize = " + maxPageSize + ", limit = " + limit + "]";
+        }
+    }
+
+    public TimeSeriesSourceOperator(
+        List<? extends RefCounted> shardContextCounters,
+        BlockFactory blockFactory,
+        int maxPageSize,
+        LuceneSliceQueue sliceQueue,
+        int limit,
+        Limiter limiter
+    ) {
+        super(shardContextCounters, blockFactory, maxPageSize, sliceQueue, limit, limiter, false);
+    }
+
+    @Override
+    protected int numMetadataBlocks() {
+        // See EsQueryExec#TIME_SERIES_SOURCE_FIELDS
+        return 2;
+    }
+
+    @Override
+    protected void buildMetadataBlocks(Block[] blocks, int offset, int currentPagePos) {
+        blocks[offset] = blockFactory.newConstantIntVector(currentSlice.slicePosition(), currentPagePos).asBlock();
+        blocks[offset + 1] = blockFactory.newConstantLongVector(Long.MAX_VALUE, currentPagePos).asBlock();
+    }
+}