Late materialization of dimension fields in time-series (#135961)

This change adds an optimization rule for time-series queries that moves 
reading dimension fields from before the time-series operator to after,
reading each dimension field once per group. This is possible because
dimension field values for `_tsid` are identical across all documents in
the same time-series.

For example:

```
TS .. | STATS sum(rate(r1)), sum(rate(r2)) BY cluster, host, tbucket(1m)
```

Without this rule:

```
TS ..
| EXTRACT_FIELDS(r1, r2, cluster, host)
| STATS rate(r1), rate(r2), VALUES(cluster), VALUES(host) BY _tsid, tbucket(1m)
| ...
```

With this rule:

```
TS ..
| EXTRACT_FIELDS(r1, r2)
| STATS rate(r1), rate(r2), FIRST_DOC_ID(_doc) BY _tsid, tbucket(1m)
| EXTRACT_FIELDS(cluster, host)
| ...
```

Ideally, dimension fields should be read once per _tsid in the final 
result, similar to the fetch phase. Currently, dimension fields are read
once per group key in each pipeline; if there are multiple time buckets,
dimensions for the same _tsid are read multiple times. This can be
avoided by extending ValuesSourceReaderOperator to understand the
ordinals of _tsid. I will follow up with this improvement later to keep
this PR small.

Author: dnhatn

docs/changelog/135961.yaml

@@ -0,0 +1,5 @@
+pr: 135961
+summary: Late materialization of dimension fields in time-series
+area: TSDB
+type: enhancement
+issues: []

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

@@ -0,0 +1,285 @@
+/*
+ * 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.aggregation;
+
+// begin generated imports
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.BytesRefVector;
+import org.elasticsearch.compute.data.ElementType;
+import org.elasticsearch.compute.data.IntArrayBlock;
+import org.elasticsearch.compute.data.IntBigArrayBlock;
+import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.IntVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.DriverContext;
+
+import java.util.List;
+// end generated imports
+
+public final class DimensionValuesByteRefGroupingAggregatorFunction implements GroupingAggregatorFunction {
+
+    public static final class FunctionSupplier implements AggregatorFunctionSupplier {
+
+        @Override
+        public List<IntermediateStateDesc> nonGroupingIntermediateStateDesc() {
+            throw new UnsupportedOperationException("non-grouping aggregator is not supported");
+        }
+
+        @Override
+        public List<IntermediateStateDesc> groupingIntermediateStateDesc() {
+            return INTERMEDIATE_STATE_DESC;
+        }
+
+        @Override
+        public AggregatorFunction aggregator(DriverContext driverContext, List<Integer> channels) {
+            throw new UnsupportedOperationException("non-grouping aggregator is not supported");
+        }
+
+        @Override
+        public DimensionValuesByteRefGroupingAggregatorFunction groupingAggregator(DriverContext driverContext, List<Integer> channels) {
+            return new DimensionValuesByteRefGroupingAggregatorFunction(channels, driverContext);
+        }
+
+        @Override
+        public String describe() {
+            return "dimensions of bytes";
+        }
+    }
+
+    static final List<IntermediateStateDesc> INTERMEDIATE_STATE_DESC = List.of(new IntermediateStateDesc("values", ElementType.BYTES_REF));
+
+    private final BytesRefBlock.Builder builder;
+    private final int channel;
+    private final DriverContext driverContext;
+    private int maxGroupId = -1;
+
+    public DimensionValuesByteRefGroupingAggregatorFunction(List<Integer> channels, DriverContext driverContext) {
+        this.channel = channels.getFirst();
+        this.driverContext = driverContext;
+        this.builder = driverContext.blockFactory().newBytesRefBlockBuilder(4096);
+    }
+
+    @Override
+    public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) {
+        // manage nulls
+    }
+
+    @Override
+    public AddInput prepareProcessRawInputPage(SeenGroupIds seenGroupIds, Page page) {
+        BytesRefBlock valuesBlock = page.getBlock(0);
+        if (valuesBlock.areAllValuesNull()) {
+            return new AddInput() {
+                @Override
+                public void add(int positionOffset, IntArrayBlock groupIds) {
+
+                }
+
+                @Override
+                public void add(int positionOffset, IntBigArrayBlock groupIds) {
+
+                }
+
+                @Override
+                public void add(int positionOffset, IntVector groupIds) {
+
+                }
+
+                @Override
+                public void close() {
+
+                }
+            };
+        }
+        return new AddInput() {
+            @Override
+            public void add(int positionOffset, IntArrayBlock groupIds) {
+                addInputValuesBlock(positionOffset, groupIds, valuesBlock);
+            }
+
+            @Override
+            public void add(int positionOffset, IntBigArrayBlock groupIds) {
+                addInputValuesBlock(positionOffset, groupIds, valuesBlock);
+            }
+
+            @Override
+            public void add(int positionOffset, IntVector groupIds) {
+                var valuesVector = valuesBlock.asVector();
+                if (valuesVector != null) {
+                    addInputValuesVector(positionOffset, groupIds, valuesVector);
+                } else {
+                    addInputValuesBlock(positionOffset, groupIds, valuesBlock);
+                }
+            }
+
+            @Override
+            public void close() {
+
+            }
+        };
+    }
+
+    // Note that this path can be executed randomly in tests, not in production
+    private void addInputValuesBlock(int positionOffset, IntBlock groups, BytesRefBlock valueBlock) {
+        var scratch = new BytesRef();
+        int positionCount = groups.getPositionCount();
+        for (int p = 0; p < positionCount; p++) {
+            if (groups.isNull(p)) {
+                continue;
+            }
+            int valuePosition = p + positionOffset;
+            int groupStart = groups.getFirstValueIndex(p);
+            int groupEnd = groupStart + groups.getValueCount(p);
+            for (int g = groupStart; g < groupEnd; g++) {
+                final int groupId = groups.getInt(g);
+                if (maxGroupId < groupId) {
+                    fillNullsUpTo(groupId);
+                    builder.copyFrom(valueBlock, valuePosition, scratch);
+                    maxGroupId = groupId;
+                }
+            }
+        }
+    }
+
+    private void addInputValuesBlock(int positionOffset, IntVector groups, BytesRefBlock valueBlock) {
+        var scratch = new BytesRef();
+        int positionCount = groups.getPositionCount();
+        if (groups.isConstant()) {
+            int groupId = groups.getInt(0);
+            if (groupId > maxGroupId) {
+                fillNullsUpTo(groupId);
+                builder.copyFrom(valueBlock, positionOffset, scratch);
+                maxGroupId = groupId;
+            }
+        } else {
+            for (int p = 0; p < positionCount; p++) {
+                int groupId = groups.getInt(p);
+                if (groupId > maxGroupId) {
+                    fillNullsUpTo(groupId);
+                    builder.copyFrom(valueBlock, positionOffset + p, scratch);
+                    maxGroupId = groupId;
+                }
+            }
+        }
+    }
+
+    private void addInputValuesVector(int positionOffset, IntVector groups, BytesRefVector valueVector) {
+        var scratch = new BytesRef();
+        int positionCount = groups.getPositionCount();
+        if (groups.isConstant()) {
+            int groupId = groups.getInt(0);
+            if (groupId > maxGroupId) {
+                fillNullsUpTo(groupId);
+                builder.appendBytesRef(valueVector.getBytesRef(positionOffset, scratch));
+                maxGroupId = groupId;
+            }
+        } else {
+            for (int p = 0; p < positionCount; p++) {
+                int groupId = groups.getInt(p);
+                if (groupId > maxGroupId) {
+                    fillNullsUpTo(groupId);
+                    builder.appendBytesRef(valueVector.getBytesRef(positionOffset + p, scratch));
+                    maxGroupId = groupId;
+                }
+            }
+        }
+    }
+
+    private void fillNullsUpTo(int groupId) {
+        for (int i = maxGroupId + 1; i < groupId; i++) {
+            builder.appendNull();
+        }
+    }
+
+    @Override
+    public int intermediateBlockCount() {
+        return INTERMEDIATE_STATE_DESC.size();
+    }
+
+    @Override
+    public void addIntermediateInput(int positionOffset, IntArrayBlock groups, Page page) {
+        BytesRefBlock valuesBlock = page.getBlock(channel);
+        if (valuesBlock.areAllValuesNull()) {
+            return;
+        }
+        addInputValuesBlock(positionOffset, groups, valuesBlock);
+    }
+
+    @Override
+    public void addIntermediateInput(int positionOffset, IntBigArrayBlock groups, Page page) {
+        BytesRefBlock valuesBlock = page.getBlock(channel);
+        if (valuesBlock.areAllValuesNull()) {
+            return;
+        }
+        addInputValuesBlock(positionOffset, groups, valuesBlock);
+    }
+
+    @Override
+    public void addIntermediateInput(int positionOffset, IntVector groups, Page page) {
+        BytesRefBlock valuesBlock = page.getBlock(channel);
+        if (valuesBlock.areAllValuesNull()) {
+            return;
+        }
+        var valuesVector = valuesBlock.asVector();
+        if (valuesVector != null) {
+            addInputValuesVector(positionOffset, groups, valuesVector);
+        } else {
+            addInputValuesBlock(positionOffset, groups, valuesBlock);
+        }
+    }
+
+    @Override
+    public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) {
+        int positionCount = selected.getPositionCount();
+        boolean allSelected = positionCount == maxGroupId + 1;
+        if (allSelected) {
+            for (int i = 0; i < selected.getPositionCount(); i++) {
+                if (selected.getInt(i) == i) {
+                    allSelected = false;
+                    break;
+                }
+            }
+        }
+        if (allSelected) {
+            blocks[offset] = builder.build();
+            return;
+        }
+        BytesRef scratch = new BytesRef();
+        try (var block = builder.build(); var outputBuilder = driverContext.blockFactory().newBytesRefBlockBuilder(positionCount)) {
+            for (int p = 0; p < positionCount; p++) {
+                int groupId = selected.getInt(p);
+                if (groupId <= maxGroupId) {
+                    outputBuilder.copyFrom(block, groupId, scratch);
+                } else {
+                    outputBuilder.appendNull();
+                }
+            }
+            blocks[offset] = outputBuilder.build();
+        }
+    }
+
+    @Override
+    public void close() {
+        builder.close();
+    }
+
+    @Override
+    public void evaluateFinal(Block[] blocks, int offset, IntVector selected, GroupingAggregatorEvaluationContext evalContext) {
+        evaluateIntermediate(blocks, offset, selected);
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append(getClass().getSimpleName()).append("[");
+        sb.append("channels=").append(channel);
+        sb.append("]");
+        return sb.toString();
+    }
+}