Fix for sample function

Author: julian-elastic

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

@@ -17,29 +17,32 @@
 import org.elasticsearch.compute.aggregation.SampleIntAggregatorFunctionSupplier;
 import org.elasticsearch.compute.aggregation.SampleLongAggregatorFunctionSupplier;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
+import org.elasticsearch.xpack.esql.common.Failures;
 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;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.Example;
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.FunctionType;
+import org.elasticsearch.xpack.esql.expression.function.FunctionUtils;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 import org.elasticsearch.xpack.esql.planner.ToAggregator;
 
 import java.io.IOException;
 import java.util.List;
 
-import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
-import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isNotNullAndFoldable;
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isNotNull;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
+import static org.elasticsearch.xpack.esql.expression.function.FunctionUtils.postOptimizationVerificationLimit;
+import static org.elasticsearch.xpack.esql.expression.function.FunctionUtils.resolveTypeLimit;
 
-public class Sample extends AggregateFunction implements ToAggregator {
+public class Sample extends AggregateFunction implements ToAggregator, PostOptimizationVerificationAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Sample", Sample::new);
 
     @FunctionInfo(
@@ -110,14 +113,14 @@ protected TypeResolution resolveType() {
             return new TypeResolution("Unresolved children");
         }
         var typeResolution = isType(field(), dt -> dt != DataType.UNSIGNED_LONG, sourceText(), FIRST, "any type except unsigned_long").and(
-            isNotNullAndFoldable(limitField(), sourceText(), SECOND)
+            isNotNull(limitField(), sourceText(), SECOND)
         ).and(isType(limitField(), dt -> dt == DataType.INTEGER, sourceText(), SECOND, "integer"));
         if (typeResolution.unresolved()) {
             return typeResolution;
         }
-        int limit = limitValue();
-        if (limit <= 0) {
-            return new TypeResolution(format(null, "Limit must be greater than 0 in [{}], found [{}]", sourceText(), limit));
+        TypeResolution result = resolveTypeLimit(limitField(), sourceText());
+        if (result.equals(TypeResolution.TYPE_RESOLVED) == false) {
+            return result;
         }
         return TypeResolution.TYPE_RESOLVED;
     }
@@ -164,11 +167,15 @@ Expression limitField() {
     }
 
     private int limitValue() {
-        return (int) limitField().fold(FoldContext.small() /* TODO remove me */);
+        return FunctionUtils.limitValue(limitField(), sourceText());
     }
 
     Expression uuid() {
         return parameters().get(1);
     }
 
+    @Override
+    public void postOptimizationVerification(Failures failures) {
+        postOptimizationVerificationLimit(failures, limitField(), sourceText());
+    }
 }

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

@@ -95,3 +95,33 @@ Top function with invalid sort order:
             | LIMIT 5
   - match: { error.type: "verification_exception" }
   - match: { error.reason: "Found 1 problem\nline 3:30: Invalid order value in [TOP(hire_date, 2, REVERSE(\"csed123\"))], expected [ASC, DESC] but got [321desc]" }
+
+---
+
+Sample function with constant folding:
+  - do:
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | STATS
+                sample_salary = SAMPLE(salary, 1+2)
+            | LIMIT 5
+  - match: { columns.0.name: "sample_salary" }
+  - length: { values: 1 }
+  - length: { values.0: 1 }
+
+---
+
+Sample function with negative limit value after folding:
+  - do:
+      catch: bad_request
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | STATS
+                sample_salary = SAMPLE(salary, 2-5)
+            | LIMIT 5
+  - match: { error.type: "verification_exception" }
+  - match: { error.reason: "Found 1 problem\nline 3:36: Limit must be greater than 0 in [SAMPLE(salary, 2-5)], found [-3]" }