Remove folding from CountDistinct

Author: julian-elastic

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

@@ -19,7 +19,6 @@
 import org.elasticsearch.compute.aggregation.CountDistinctLongAggregatorFunctionSupplier;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
-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;
@@ -43,6 +42,7 @@
 import java.util.Map;
 import java.util.function.Function;
 
+import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
@@ -210,16 +210,23 @@ protected TypeResolution resolveType() {
     @Override
     public AggregatorFunctionSupplier supplier() {
         DataType type = field().dataType();
-        int precision = this.precision == null
-            ? DEFAULT_PRECISION
-            : ((Number) this.precision.fold(FoldContext.small() /* TODO remove me */)).intValue();
+        int precision = this.precision == null ? DEFAULT_PRECISION : precisionValue();
         if (SUPPLIERS.containsKey(type) == false) {
             // If the type checking did its job, this should never happen
             throw EsqlIllegalArgumentException.illegalDataType(type);
         }
         return SUPPLIERS.get(type).apply(precision);
     }
 
+    private int precisionValue() {
+        if (precision instanceof Literal literal) {
+            return ((Number) literal.value()).intValue();
+        }
+        throw new EsqlIllegalArgumentException(
+            format(null, "Precision value must be a constant integer in [{}], found [{}]", source(), precision)
+        );
+    }
+
     @Override
     public Expression surrogate() {
         var s = source();

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/230_folding.yml

@@ -581,3 +581,37 @@ Percentile with non-foldable value:
             FROM employees | STATS p = PERCENTILE(salary, salary + 1)  | LIMIT 1
   - match: { error.type: "verification_exception" }
   - contains: { error.reason: "second argument of [PERCENTILE(salary, salary + 1)] must be a constant, received [salary + 1]" }
+
+---
+COUNT_DISTINCT with foldable precision:
+  - do:
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | STATS c1 = COUNT_DISTINCT(salary, 234 + 543)  | LIMIT 1
+  - match: { columns.0.name: "c1" }
+  - length: { values: 1 }
+
+---
+COUNT_DISTINCT with no precision:
+  - do:
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | STATS c1 = COUNT_DISTINCT(salary)  | LIMIT 1
+  - match: { columns.0.name: "c1" }
+  - length: { values: 1 }
+
+---
+COUNT_DISTINCT with non-foldable precision:
+  - do:
+      catch: bad_request
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | STATS COUNT_DISTINCT(salary, salary + 1)  | LIMIT 1
+  - match: { error.type: "verification_exception" }
+  - contains: { error.reason: "second argument of [COUNT_DISTINCT(salary, salary + 1)] must be a constant, received [salary + 1]" }