Remove deprecated function isNotNullAndFoldable

docs/changelog/130944.yaml

@@ -0,0 +1,6 @@
+pr: 130944
+summary: Remove unnecessary calls to Fold
+area: ES|QL
+type: enhancement
+issues:
+ - 119756

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/TypeResolutions.java

@@ -133,36 +133,6 @@ public static TypeResolution isFoldable(Expression e, String operationName, Para
         return TypeResolution.TYPE_RESOLVED;
     }
 
-    /**
-     * Is this {@link Expression#foldable()} and not {@code null}.
-     *
-     * @deprecated instead of calling this, check for a {@link Literal} containing
-     *             {@code null}. Foldable expressions will be folded by other rules,
-     *             eventually, to a {@link Literal}.
-     */
-    @Deprecated
-    public static TypeResolution isNotNullAndFoldable(Expression e, String operationName, ParamOrdinal paramOrd) {
-        TypeResolution resolution = isFoldable(e, operationName, paramOrd);
-
-        if (resolution.unresolved()) {
-            return resolution;
-        }
-
-        if (e.dataType() == DataType.NULL || e.fold(FoldContext.small()) == null) {
-            resolution = new TypeResolution(
-                format(
-                    null,
-                    "{}argument of [{}] cannot be null, received [{}]",
-                    paramOrd == null || paramOrd == DEFAULT ? "" : paramOrd.name().toLowerCase(Locale.ROOT) + " ",
-                    operationName,
-                    Expressions.name(e)
-                )
-            );
-        }
-
-        return resolution;
-    }
-
     public static TypeResolution isNotNull(Expression e, String operationName, ParamOrdinal paramOrd) {
         if (e.dataType() == DataType.NULL) {
             return new TypeResolution(

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

@@ -0,0 +1,151 @@
+/*
+ * 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;
+
+import org.elasticsearch.common.lucene.BytesRefs;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.common.Failure;
+import org.elasticsearch.xpack.esql.common.Failures;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+
+import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
+import static org.elasticsearch.xpack.esql.common.Failure.fail;
+
+public class FunctionUtils {
+    /**
+     * A utility class to validate the type resolution of expressions before and after logical planning.
+     * If null is passed for Failures to the constructor, it means we are only type resolution.
+     * This is usually called when doing pre-logical planning validation.
+     * If a {@link Failures} instance is passed, it means we are doing post-logical planning validation as well.
+     * This is usually called after folding is done, during
+     * {@link org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware} verification
+     */
+    public static class TypeResolutionValidator {
+
+        Expression.TypeResolution typeResolution = Expression.TypeResolution.TYPE_RESOLVED;
+        @Nullable
+        private final Failures postValidationFailures; // null means we are doing pre-folding validation only
+        private final Expression field;
+
+        public static TypeResolutionValidator forPreOptimizationValidation(Expression field) {
+            return new TypeResolutionValidator(field, null);
+        }
+
+        public static TypeResolutionValidator forPostOptimizationValidation(Expression field, Failures failures) {
+            return new TypeResolutionValidator(field, failures);
+        }
+
+        private TypeResolutionValidator(Expression field, Failures failures) {
+            this.field = field;
+            this.postValidationFailures = failures;
+        }
+
+        public void invalidIfPostValidation(Failure failure) {
+            if (postValidationFailures != null) {
+                postValidationFailures.add(failure);
+            }
+        }
+
+        public void invalid(Expression.TypeResolution message) {
+            typeResolution = message;
+            if (postValidationFailures != null) {
+                postValidationFailures.add(fail(field, message.message()));
+            }
+        }
+
+        public Expression.TypeResolution getResolvedType() {
+            return typeResolution;
+        }
+    }
+
+    public static Integer limitValue(Expression limitField, String sourceText) {
+        if (limitField instanceof Literal literal) {
+            Object value = literal.value();
+            if (value instanceof Integer intValue) {
+                return intValue;
+            }
+        }
+        throw new EsqlIllegalArgumentException(format(null, "Limit value must be an integer in [{}], found [{}]", sourceText, limitField));
+    }
+
+    /**
+     * We check that the limit is not null and that if it is a literal, it is a positive integer
+     * During postOptimizationVerification folding is already done, so we also verify that it is definitively a literal
+     */
+    public static Expression.TypeResolution resolveTypeLimit(Expression limitField, String sourceText, TypeResolutionValidator validator) {
+        if (limitField == null) {
+            validator.invalid(
+                new Expression.TypeResolution(format(null, "Limit must be a constant integer in [{}], found [{}]", sourceText, limitField))
+            );
+        } else if (limitField instanceof Literal literal) {
+            if (literal.value() == null) {
+                validator.invalid(
+                    new Expression.TypeResolution(
+                        format(null, "Limit must be a constant integer in [{}], found [{}]", sourceText, limitField)
+                    )
+                );
+            } else {
+                int value = (Integer) literal.value();
+                if (value <= 0) {
+                    validator.invalid(
+                        new Expression.TypeResolution(format(null, "Limit must be greater than 0 in [{}], found [{}]", sourceText, value))
+                    );
+                }
+            }
+        } else {
+            // it is expected that the expression is a literal after folding
+            // we fail if it is not a literal
+            validator.invalidIfPostValidation(
+                fail(limitField, "Limit must be a constant integer in [{}], found [{}]", sourceText, limitField)
+            );
+        }
+        return validator.getResolvedType();
+    }
+
+    /**
+     * We check that the query is not null and that if it is a literal, it is a string
+     * During postOptimizationVerification folding is already done, so we also verify that it is definitively a literal
+     */
+    public static Expression.TypeResolution resolveTypeQuery(Expression queryField, String sourceText, TypeResolutionValidator validator) {
+        if (queryField == null) {
+            validator.invalid(
+                new Expression.TypeResolution(format(null, "Query must be a valid string in [{}], found [{}]", sourceText, queryField))
+            );
+        } else if (queryField instanceof Literal literal) {
+            if (literal.value() == null) {
+                validator.invalid(
+                    new Expression.TypeResolution(format(null, "Query value cannot be null in [{}], but got [{}]", sourceText, queryField))
+                );
+            }
+        } else {
+            // it is expected that the expression is a literal after folding
+            // we fail if it is not a literal
+            validator.invalidIfPostValidation(fail(queryField, "Query must be a valid string in [{}], found [{}]", sourceText, queryField));
+        }
+        return validator.getResolvedType();
+    }
+
+    public static Object queryAsObject(Expression queryField, String sourceText) {
+        if (queryField instanceof Literal literal) {
+            return literal.value();
+        }
+        throw new EsqlIllegalArgumentException(
+            format(null, "Query value must be a constant string in [{}], found [{}]", sourceText, queryField)
+        );
+    }
+
+    public static String queryAsString(Expression queryField, String sourceText) {
+        if (queryField instanceof Literal literal) {
+            return BytesRefs.toString(literal.value());
+        }
+        throw new EsqlIllegalArgumentException(
+            format(null, "Query value must be a constant string in [{}], found [{}]", sourceText, queryField)
+        );
+    }
+}

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

@@ -17,29 +17,33 @@
 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.TypeResolutionValidator.forPostOptimizationValidation;
+import static org.elasticsearch.xpack.esql.expression.function.FunctionUtils.TypeResolutionValidator.forPreOptimizationValidation;
+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 +114,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(), forPreOptimizationValidation(limitField()));
+        if (result.equals(TypeResolution.TYPE_RESOLVED) == false) {
+            return result;
         }
         return TypeResolution.TYPE_RESOLVED;
     }
@@ -164,11 +168,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) {
+        FunctionUtils.resolveTypeLimit(limitField(), sourceText(), forPostOptimizationValidation(limitField(), failures));
+    }
 }