ESQL: AVG aggregation tests and ignore complex surrogates

docs/reference/esql/functions/aggregation-functions.asciidoc

@@ -8,7 +8,7 @@
 The <<esql-stats-by>> command supports these aggregate functions:
 
 // tag::agg_list[]
-* <<esql-agg-avg>>
+* <<esql-avg>>
 * <<esql-agg-count>>
 * <<esql-agg-count-distinct>>
 * <<esql-agg-max>>
@@ -23,7 +23,6 @@ The <<esql-stats-by>> command supports these aggregate functions:
 * experimental:[] <<esql-agg-weighted-avg>>
 // end::agg_list[]
 
-include::avg.asciidoc[]
 include::count.asciidoc[]
 include::count-distinct.asciidoc[]
 include::max.asciidoc[]
@@ -33,6 +32,7 @@ include::min.asciidoc[]
 include::percentile.asciidoc[]
 include::st_centroid_agg.asciidoc[]
 include::sum.asciidoc[]
+include::layout/avg.asciidoc[]
 include::layout/top.asciidoc[]
 include::values.asciidoc[]
 include::weighted-avg.asciidoc[]

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Avg.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.SurrogateExpression;
+import org.elasticsearch.xpack.esql.expression.function.Example;
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvAvg;
@@ -28,7 +29,20 @@
 public class Avg extends AggregateFunction implements SurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Avg", Avg::new);
 
-    @FunctionInfo(returnType = "double", description = "The average of a numeric field.", isAggregation = true)
+    @FunctionInfo(
+        returnType = "double",
+        description = "The average of a numeric field.",
+        isAggregation = true,
+        examples = {
+            @Example(file = "stats", tag = "avg"),
+            @Example(
+                description = "The expression can use inline functions. For example, to calculate the average "
+                    + "over a multivalued column, first use `MV_AVG` to average the multiple values per row, "
+                    + "and use the result with the `AVG` function",
+                file = "stats",
+                tag = "docsStatsAvgNestedExpression"
+            ) }
+    )
     public Avg(Source source, @Param(name = "number", type = { "double", "integer", "long" }) Expression field) {
         super(source, field);
     }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractAggregationTestCase.java

@@ -15,6 +15,8 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.util.NumericUtils;
 import org.elasticsearch.xpack.esql.expression.SurrogateExpression;
@@ -251,6 +253,13 @@ private void resolveExpression(Expression expression, Consumer<Expression> onAgg
         expression = new FoldNull().rule(expression);
         assertThat(expression.dataType(), equalTo(testCase.expectedType()));
 
+        assumeTrue(
+            "Surrogate expression with non-trivial children cannot be evaluated",
+            expression.children()
+                .stream()
+                .allMatch(child -> child instanceof FieldAttribute || child instanceof DeepCopy || child instanceof Literal)
+        );
+
         if (expression instanceof AggregateFunction == false) {
             onEvaluableExpression.accept(expression);
             return;

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AvgTests.java

@@ -0,0 +1,95 @@
+/*
+ * 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.xpack.esql.expression.function.aggregate;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase;
+import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier;
+import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class AvgTests extends AbstractAggregationTestCase {
+    public AvgTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
+        this.testCase = testCaseSupplier.get();
+    }
+
+    @ParametersFactory
+    public static Iterable<Object[]> parameters() {
+        var suppliers = new ArrayList<TestCaseSupplier>();
+
+        Stream.of(
+            MultiRowTestCaseSupplier.intCases(1, 1000, Integer.MIN_VALUE, Integer.MAX_VALUE, true),
+            MultiRowTestCaseSupplier.longCases(1, 1000, Long.MIN_VALUE, Long.MAX_VALUE, true),
+            MultiRowTestCaseSupplier.doubleCases(1, 1000, -Double.MAX_VALUE, Double.MAX_VALUE, true)
+        ).flatMap(List::stream).map(AvgTests::makeSupplier).collect(Collectors.toCollection(() -> suppliers));
+
+        suppliers.add(
+            // Folding
+            new TestCaseSupplier(
+                List.of(DataType.INTEGER),
+                () -> new TestCaseSupplier.TestCase(
+                    List.of(TestCaseSupplier.TypedData.multiRow(List.of(200), DataType.INTEGER, "field")),
+                    "Avg[field=Attribute[channel=0]]",
+                    DataType.DOUBLE,
+                    equalTo(200.)
+                )
+            )
+        );
+
+        return parameterSuppliersFromTypedDataWithDefaultChecks(suppliers);
+    }
+
+    @Override
+    protected Expression build(Source source, List<Expression> args) {
+        return new Avg(source, args.get(0));
+    }
+
+    private static TestCaseSupplier makeSupplier(TestCaseSupplier.TypedDataSupplier fieldSupplier) {
+        return new TestCaseSupplier(List.of(fieldSupplier.type()), () -> {
+            var fieldTypedData = fieldSupplier.get();
+
+            Object expected = switch (fieldTypedData.type().widenSmallNumeric()) {
+                case INTEGER -> fieldTypedData.multiRowData()
+                    .stream()
+                    .map(v -> (Integer) v)
+                    .collect(Collectors.summarizingInt(Integer::intValue))
+                    .getAverage();
+                case LONG -> fieldTypedData.multiRowData()
+                    .stream()
+                    .map(v -> (Long) v)
+                    .collect(Collectors.summarizingLong(Long::longValue))
+                    .getAverage();
+                case DOUBLE -> fieldTypedData.multiRowData()
+                    .stream()
+                    .map(v -> (Double) v)
+                    .collect(Collectors.summarizingDouble(Double::doubleValue))
+                    .getAverage();
+                default -> throw new IllegalStateException("Unexpected value: " + fieldTypedData.type());
+            };
+
+            return new TestCaseSupplier.TestCase(
+                List.of(fieldTypedData),
+                "Avg[field=Attribute[channel=0]]",
+                DataType.DOUBLE,
+                equalTo(expected)
+            );
+        });
+    }
+}

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/TopTests.java

@@ -22,6 +22,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -37,14 +38,15 @@ public static Iterable<Object[]> parameters() {
 
         for (var limitCaseSupplier : TestCaseSupplier.intCases(1, 1000, false)) {
             for (String order : List.of("asc", "desc")) {
-                for (var fieldCaseSupplier : Stream.of(
+                Stream.of(
                     MultiRowTestCaseSupplier.intCases(1, 1000, Integer.MIN_VALUE, Integer.MAX_VALUE, true),
                     MultiRowTestCaseSupplier.longCases(1, 1000, Long.MIN_VALUE, Long.MAX_VALUE, true),
                     MultiRowTestCaseSupplier.doubleCases(1, 1000, -Double.MAX_VALUE, Double.MAX_VALUE, true),
                     MultiRowTestCaseSupplier.dateCases(1, 1000)
-                ).flatMap(List::stream).toList()) {
-                    suppliers.add(TopTests.makeSupplier(fieldCaseSupplier, limitCaseSupplier, order));
-                }
+                )
+                    .flatMap(List::stream)
+                    .map(fieldCaseSupplier -> TopTests.makeSupplier(fieldCaseSupplier, limitCaseSupplier, order))
+                    .collect(Collectors.toCollection(() -> suppliers));
             }
         }