[ES|QL] Fuse From_Aggregate_Metric_Double to BlockLoader (elastic#138392)

This commit fuses the FROM_AGGREGATE_METRIC_DOUBLE function into the blockloaders for the aggregate_metric_double field type.
In particular, it works by directly loading the individual sub-blocks, rather than the entire AggregateMetricDoubleBlock, when called in a function that uses one of the sub-blocks via FROM_AGGREGATE_METRIC_DOUBLE

Author: limotova

muted-tests.yml

@@ -436,6 +436,12 @@ tests:
 - class: org.elasticsearch.xpack.unsignedlong.UnsignedLongSyntheticSourceNativeArrayIntegrationTests
   method: testSynthesizeArrayRandom
   issue: https://github.com/elastic/elasticsearch/issues/138684
+- class: org.elasticsearch.xpack.esql.optimizer.LocalLogicalPlanOptimizerTests
+  method: testAggregateMetricDoubleInlineStats
+  issue: https://github.com/elastic/elasticsearch/issues/138620
+- class: org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceDateTruncBucketWithRoundToTests
+  method: testAggregateMetricDoubleInlineStats
+  issue: https://github.com/elastic/elasticsearch/issues/138620
 - class: org.elasticsearch.index.mapper.HalfFloatSyntheticSourceNativeArrayIntegrationTests
   method: testSynthesizeArrayRandom
   issue: https://github.com/elastic/elasticsearch/issues/138730

server/src/main/java/org/elasticsearch/index/mapper/blockloader/BlockLoaderFunctionConfig.java

@@ -45,6 +45,10 @@ public int hashCode() {
     }
 
     enum Function {
+        AMD_COUNT,
+        AMD_MAX,
+        AMD_MIN,
+        AMD_SUM,
         MV_MAX,
         MV_MIN,
         LENGTH,

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/AggregateMetricDoubleBlockBuilder.java

@@ -203,6 +203,16 @@ public int getIndex() {
         public String getLabel() {
             return label;
         }
+
+        public static Metric indexToMetric(int i) {
+            return switch (i) {
+                case 0 -> MIN;
+                case 1 -> MAX;
+                case 2 -> SUM;
+                case 3 -> COUNT;
+                default -> null;
+            };
+        }
     }
 
     /**

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/FromAggregateMetricDouble.java

@@ -19,18 +19,22 @@
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.core.Releasables;
+import org.elasticsearch.index.mapper.blockloader.BlockLoaderFunctionConfig;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
+import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.Param;
+import org.elasticsearch.xpack.esql.expression.function.blockloader.BlockLoaderExpression;
 import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 import java.io.IOException;
 import java.util.List;
@@ -43,7 +47,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
 import static org.elasticsearch.xpack.esql.core.type.DataType.NULL;
 
-public class FromAggregateMetricDouble extends EsqlScalarFunction implements ConvertFunction {
+public class FromAggregateMetricDouble extends EsqlScalarFunction implements ConvertFunction, BlockLoaderExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Expression.class,
         "FromAggregateMetricDouble",
@@ -119,7 +123,7 @@ protected TypeResolution resolveType() {
         if (childrenResolved() == false) {
             return new TypeResolution("Unresolved children");
         }
-        return isType(field, dt -> dt == DataType.AGGREGATE_METRIC_DOUBLE, sourceText(), DEFAULT, "aggregate_metric_double only");
+        return isType(field, dt -> dt == AGGREGATE_METRIC_DOUBLE, sourceText(), DEFAULT, "aggregate_metric_double only");
     }
 
     @Override
@@ -192,4 +196,24 @@ public Expression field() {
     public Set<DataType> supportedTypes() {
         return Set.of(AGGREGATE_METRIC_DOUBLE);
     }
+
+    @Override
+    public PushedBlockLoaderExpression tryPushToFieldLoading(SearchStats stats) {
+        if (field() instanceof FieldAttribute f && f.dataType() == AGGREGATE_METRIC_DOUBLE) {
+            var folded = subfieldIndex.fold(FoldContext.small());
+            if (folded == null) {
+                throw new IllegalArgumentException("Subfield Index was null");
+            }
+            var subfield = ((Number) folded).intValue();
+            var functionConfig = switch (AggregateMetricDoubleBlockBuilder.Metric.indexToMetric(subfield)) {
+                case MIN -> BlockLoaderFunctionConfig.Function.AMD_MIN;
+                case MAX -> BlockLoaderFunctionConfig.Function.AMD_MAX;
+                case SUM -> BlockLoaderFunctionConfig.Function.AMD_SUM;
+                case COUNT -> BlockLoaderFunctionConfig.Function.AMD_COUNT;
+                case null -> throw new IllegalArgumentException("Received invalid subfield index [" + subfield + "].");
+            };
+            return new PushedBlockLoaderExpression(f, new BlockLoaderFunctionConfig.JustFunction(functionConfig));
+        }
+        return null;
+    }
 }