Add support for more functions

Author: julian-elastic

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

@@ -26,7 +26,6 @@
 import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
-import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.TypedAttribute;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -213,7 +212,7 @@ private Query translate(TranslatorHandler handler, Expression spatialExpression,
         String name = handler.nameOf(attribute);
 
         try {
-            Geometry shape = SpatialRelatesUtils.makeGeometryFromLiteral(FoldContext.small() /* TODO remove me */, constantExpression);
+            Geometry shape = SpatialRelatesUtils.makeGeometryFromLiteral(constantExpression);
             return new SpatialRelatesQuery(source(), name, queryRelation(), shape, attribute.dataType());
         } catch (IllegalArgumentException e) {
             throw new QlIllegalArgumentException(e.getMessage(), e);

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

@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.function.Function;
 
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
 import static org.elasticsearch.xpack.esql.core.expression.Foldables.valueOf;
 
 public class SpatialRelatesUtils {
@@ -168,13 +169,20 @@ private static Geometry asGeometry(BytesRefBlock valueBlock, int position) {
     }
 
     /**
-     * This function is used in two places, when evaluating a spatial constant in the SpatialRelatesFunction, as well as when
-     * we do lucene-pushdown of spatial functions.
+     * This function is used when evaluating a spatial constant in the SpatialRelatesFunction
      */
     public static Geometry makeGeometryFromLiteral(FoldContext ctx, Expression expr) {
         return makeGeometryFromLiteralValue(valueOf(ctx, expr), expr.dataType());
     }
 
+    /**
+     * This function is used when we do lucene-pushdown of spatial functions.
+     * The expression is expected to be folded already and a literal
+     */
+    public static Geometry makeGeometryFromLiteral(Expression expr) {
+        return makeGeometryFromLiteralValue(literalValueOf(expr), expr.dataType());
+    }
+
     private static Geometry makeGeometryFromLiteralValue(Object value, DataType dataType) {
         if (value instanceof BytesRef bytesRef) {
             // Single value expression

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/Range.java

@@ -33,7 +33,7 @@
 import java.util.Objects;
 
 import static java.util.Arrays.asList;
-import static org.elasticsearch.xpack.esql.core.expression.Foldables.valueOf;
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
 import static org.elasticsearch.xpack.esql.core.type.DataType.IP;
 import static org.elasticsearch.xpack.esql.core.type.DataType.UNSIGNED_LONG;
@@ -232,8 +232,8 @@ public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHand
     }
 
     private RangeQuery translate(TranslatorHandler handler) {
-        Object l = valueOf(FoldContext.small() /* TODO remove me */, lower);
-        Object u = valueOf(FoldContext.small() /* TODO remove me */, upper);
+        Object l = literalValueOf(lower);
+        Object u = literalValueOf(upper);
         String format = null;
 
         DataType dataType = value.dataType();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/EsqlBinaryComparison.java

@@ -51,7 +51,8 @@
 import java.util.Map;
 
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
-import static org.elasticsearch.xpack.esql.core.expression.Foldables.valueOf;
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.extractLiteralOrReturnSelf;
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
 import static org.elasticsearch.xpack.esql.core.type.DataType.IP;
@@ -378,7 +379,7 @@ public Expression singleValueField() {
     private Query translate(TranslatorHandler handler) {
         TypedAttribute attribute = LucenePushdownPredicates.checkIsPushableAttribute(left());
         String name = handler.nameOf(attribute);
-        Object value = valueOf(FoldContext.small() /* TODO remove me */, right());
+        Object value = extractLiteralOrReturnSelf(right());
         String format = null;
         boolean isDateLiteralComparison = false;
 
@@ -478,7 +479,7 @@ private Query translateOutOfRangeComparisons() {
         if ((left() instanceof FieldAttribute) == false || left().dataType().isNumeric() == false) {
             return null;
         }
-        Object value = valueOf(FoldContext.small() /* TODO remove me */, right());
+        Object value = literalValueOf(right());
 
         // Comparisons with multi-values always return null in ESQL.
         if (value instanceof List<?>) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/InsensitiveEquals.java

@@ -31,7 +31,7 @@
 
 import java.io.IOException;
 
-import static org.elasticsearch.xpack.esql.core.expression.Foldables.valueOf;
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
 
 public class InsensitiveEquals extends InsensitiveBinaryComparison {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
@@ -125,7 +125,7 @@ private void checkInsensitiveComparison() {
 
     private Query translate() {
         TypedAttribute attribute = LucenePushdownPredicates.checkIsPushableAttribute(left());
-        BytesRef value = BytesRefs.toBytesRef(valueOf(FoldContext.small() /* TODO remove me */, right()));
+        BytesRef value = BytesRefs.toBytesRef(literalValueOf(right()));
         String name = LucenePushdownPredicates.pushableAttributeName(attribute);
         return new TermQuery(source(), name, value.utf8ToString(), true);
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/InsensitiveEqualsMapper.java

@@ -23,6 +23,7 @@
 
 import java.util.List;
 
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
 import static org.elasticsearch.xpack.esql.evaluator.EvalMapper.toEvaluator;
 
 public class InsensitiveEqualsMapper extends ExpressionMapper<InsensitiveEquals> {
@@ -44,7 +45,7 @@ public final ExpressionEvaluator.Factory map(
         var rightEval = toEvaluator(foldCtx, bc.right(), layout, shardContexts);
         if (DataType.isString(leftType)) {
             if (bc.right().foldable() && DataType.isString(rightType)) {
-                BytesRef rightVal = BytesRefs.toBytesRef(bc.right().fold(FoldContext.small() /* TODO remove me */));
+                BytesRef rightVal = BytesRefs.toBytesRef(literalValueOf(bc.right()));
                 Automaton automaton = InsensitiveEquals.automaton(rightVal);
                 return dvrCtx -> new InsensitiveEqualsConstantEvaluator(
                     bc.source(),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Enrich.java

@@ -26,7 +26,6 @@
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.EmptyAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
-import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.NameId;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
@@ -48,6 +47,7 @@
 
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
 public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAnalysisPlanVerificationAware, TelemetryAware, SortAgnostic {
@@ -152,7 +152,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeNamedWriteable(policyName());
         out.writeNamedWriteable(matchField());
         if (out.getTransportVersion().before(TransportVersions.V_8_13_0)) {
-            out.writeString(BytesRefs.toString(policyName().fold(FoldContext.small() /* TODO remove me */))); // old policy name
+            out.writeString(BytesRefs.toString(literalValueOf(policyName()))); // old policy name
         }
         policy().writeTo(out);
         if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_13_0)) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java

@@ -11,7 +11,6 @@
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
-import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.ChangePoint;
@@ -51,6 +50,8 @@
 
 import java.util.List;
 
+import static org.elasticsearch.xpack.esql.core.expression.Foldables.literalValueOf;
+
 /**
  * Class for sharing code across Mappers.
  */
@@ -113,7 +114,7 @@ static PhysicalPlan mapUnary(UnaryPlan p, PhysicalPlan child) {
                 enrich.mode(),
                 enrich.policy().getType(),
                 enrich.matchField(),
-                BytesRefs.toString(enrich.policyName().fold(FoldContext.small() /* TODO remove me */)),
+                BytesRefs.toString(literalValueOf(enrich.policyName())),
                 enrich.policy().getMatchField(),
                 enrich.concreteIndices(),
                 enrich.enrichFields()

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

@@ -678,3 +678,34 @@ CIDR_MATCH with non-foldable CIDRs:
             | WHERE CIDR_MATCH(ip, CONCAT(TO_STRING(ip), "/8"), "127.0.0.3/32") | LIMIT 5
   - length: { values: 3 }
 
+---
+Employee name equals:
+  - do:
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | WHERE CONCAT("Alice"," Smith") == name
+            | KEEP hire_date, salary, name
+            | LIMIT 5
+  - match: { columns.0.name: "hire_date" }
+  - match: { columns.1.name: "salary" }
+  - match: { columns.2.name: "name" }
+  - length: { values: 1 }
+  - match: { values.0.2: "Alice Smith" }
+
+---
+Employee name in list:
+  - do:
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | WHERE name in (CONCAT("Alice"," Smith"), CONCAT("fake", " name"))
+            | KEEP hire_date, salary, name
+            | LIMIT 5
+  - match: { columns.0.name: "hire_date" }
+  - match: { columns.1.name: "salary" }
+  - match: { columns.2.name: "name" }
+  - length: { values: 1 }
+  - match: { values.0.2: "Alice Smith" }