ESQL: Pass FoldContext to Translatable expressions

Author: ivancea

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/TranslationAware.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.capabilities;
 
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
@@ -27,11 +28,11 @@ public interface TranslationAware {
      * If during translation a child needs to be translated first, the handler needs to be used even if the child implements this
      * interface as well. This is to ensure that the child is wrapped in a SingleValueQuery if necessary.
      * <p>So use this:</p>
-     * <p>{@code Query childQuery = handler.asQuery(child);}</p>
+     * <p>{@code Query childQuery = handler.asQuery(child, foldContext);}</p>
      * <p>and <b>not</b> this:</p>
-     * <p>{@code Query childQuery = child.asQuery(handler);}</p>
+     * <p>{@code Query childQuery = child.asQuery(handler, foldContext);}</p>
      */
-    Query asQuery(TranslatorHandler handler);
+    Query asQuery(TranslatorHandler handler, FoldContext foldContext);
 
     /**
      * Subinterface for expressions that can only process single values (and null out on MVs).

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java

@@ -94,7 +94,7 @@ public void esql(
         });
         // Wrap it in a listener so that if we have any exceptions during execution, the listener picks it up
         // and all the metrics are properly updated
-        ActionListener.run(executeListener, l -> session.execute(request, executionInfo, planRunner, l));
+        ActionListener.run(executeListener, l -> session.execute(request, foldContext, executionInfo, planRunner, l));
     }
 
     public IndexResolver indexResolver() {

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

@@ -108,8 +108,8 @@ public Expression query() {
      *
      * @return query expression as an object
      */
-    public Object queryAsObject() {
-        Object queryAsObject = query().fold(FoldContext.small() /* TODO remove me */);
+    public Object queryAsObject(FoldContext foldContext) {
+        Object queryAsObject = query().fold(foldContext);
         return BytesRefs.toString(queryAsObject);
     }
 
@@ -148,15 +148,15 @@ public boolean translatable(LucenePushdownPredicates pushdownPredicates) {
     }
 
     @Override
-    public Query asQuery(TranslatorHandler handler) {
-        return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(handler);
+    public Query asQuery(TranslatorHandler handler, FoldContext foldContext) {
+        return queryBuilder != null ? new TranslationAwareExpressionQuery(source(), queryBuilder) : translate(handler, foldContext);
     }
 
     public QueryBuilder queryBuilder() {
         return queryBuilder;
     }
 
-    protected abstract Query translate(TranslatorHandler handler);
+    protected abstract Query translate(TranslatorHandler handler, FoldContext foldContext);
 
     public abstract Expression replaceQueryBuilder(QueryBuilder queryBuilder);
 

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -89,8 +90,8 @@ protected NodeInfo<? extends Expression> info() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
-        return new KqlQuery(source(), Objects.toString(queryAsObject()));
+    protected Query translate(TranslatorHandler handler, FoldContext foldContext) {
+        return new KqlQuery(source(), Objects.toString(queryAsObject(foldContext)));
     }
 
     @Override

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

@@ -452,8 +452,8 @@ public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
     }
 
     @Override
-    public Object queryAsObject() {
-        Object queryAsObject = query().fold(FoldContext.small() /* TODO remove me */);
+    public Object queryAsObject(FoldContext foldContext) {
+        Object queryAsObject = query().fold(foldContext);
 
         // Convert BytesRef to string for string-based values
         if (queryAsObject instanceof BytesRef bytesRef) {
@@ -478,7 +478,7 @@ public Object queryAsObject() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
+    protected Query translate(TranslatorHandler handler, FoldContext foldContext) {
         var fieldAttribute = fieldAsFieldAttribute();
         Check.notNull(fieldAttribute, "Match must have a field attribute as the first argument");
         String fieldName = fieldAttribute.name();
@@ -487,7 +487,7 @@ protected Query translate(TranslatorHandler handler) {
             fieldName = multiTypeEsField.getName();
         }
         // Make query lenient so mixed field types can be queried when a field type is incompatible with the value provided
-        return new MatchQuery(source(), fieldName, queryAsObject(), matchQueryOptions());
+        return new MatchQuery(source(), fieldName, queryAsObject(foldContext), matchQueryOptions());
     }
 
     private FieldAttribute fieldAsFieldAttribute() {

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.Rewriteable;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.util.Holder;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
@@ -26,22 +27,27 @@
 /**
  * Some {@link FullTextFunction} implementations such as {@link org.elasticsearch.xpack.esql.expression.function.fulltext.Match}
  * will be translated to a {@link QueryBuilder} that require a rewrite phase on the coordinator.
- * {@link QueryBuilderResolver#resolveQueryBuilders(LogicalPlan, TransportActionServices, ActionListener)} will rewrite the plan by
+ * {@link QueryBuilderResolver#resolveQueryBuilders(LogicalPlan, FoldContext, TransportActionServices, ActionListener)} will rewrite the plan by
  * replacing {@link FullTextFunction} expression with new ones that hold rewritten {@link QueryBuilder}s.
  */
 public final class QueryBuilderResolver {
 
     private QueryBuilderResolver() {}
 
-    public static void resolveQueryBuilders(LogicalPlan plan, TransportActionServices services, ActionListener<LogicalPlan> listener) {
+    public static void resolveQueryBuilders(
+        LogicalPlan plan,
+        FoldContext foldContext,
+        TransportActionServices services,
+        ActionListener<LogicalPlan> listener
+    ) {
         var hasFullTextFunctions = plan.anyMatch(p -> {
             Holder<Boolean> hasFullTextFunction = new Holder<>(false);
             p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
             return hasFullTextFunction.get();
         });
         if (hasFullTextFunctions) {
             Rewriteable.rewriteAndFetch(
-                new FullTextFunctionsRewritable(plan),
+                new FullTextFunctionsRewritable(plan, foldContext),
                 queryRewriteContext(services, indexNames(plan)),
                 listener.delegateFailureAndWrap((l, r) -> l.onResponse(r.plan))
             );
@@ -69,14 +75,16 @@ private static Set<String> indexNames(LogicalPlan plan) {
         return indexNames;
     }
 
-    private record FullTextFunctionsRewritable(LogicalPlan plan) implements Rewriteable<QueryBuilderResolver.FullTextFunctionsRewritable> {
+    private record FullTextFunctionsRewritable(LogicalPlan plan, FoldContext foldContext)
+        implements
+            Rewriteable<QueryBuilderResolver.FullTextFunctionsRewritable> {
         @Override
         public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException {
             Holder<IOException> exceptionHolder = new Holder<>();
             Holder<Boolean> updated = new Holder<>(false);
             LogicalPlan newPlan = plan.transformExpressionsDown(FullTextFunction.class, f -> {
                 QueryBuilder builder = f.queryBuilder(), initial = builder;
-                builder = builder == null ? f.asQuery(TranslatorHandler.TRANSLATOR_HANDLER).asBuilder() : builder;
+                builder = builder == null ? f.asQuery(TranslatorHandler.TRANSLATOR_HANDLER, foldContext).asBuilder() : builder;
                 try {
                     builder = builder.rewrite(ctx);
                 } catch (IOException e) {
@@ -89,7 +97,7 @@ public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOExc
             if (exceptionHolder.get() != null) {
                 throw exceptionHolder.get();
             }
-            return updated.get() ? new FullTextFunctionsRewritable(newPlan) : this;
+            return updated.get() ? new FullTextFunctionsRewritable(newPlan, foldContext) : this;
         }
     }
 }

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.core.querydsl.query.QueryStringQuery;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
@@ -101,8 +102,8 @@ protected NodeInfo<? extends Expression> info() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
-        return new QueryStringQuery(source(), Objects.toString(queryAsObject()), Map.of(), Map.of());
+    protected Query translate(TranslatorHandler handler, FoldContext foldContext) {
+        return new QueryStringQuery(source(), Objects.toString(queryAsObject(foldContext)), Map.of(), Map.of());
     }
 
     @Override

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.core.querydsl.query.TermQuery;
@@ -140,8 +141,8 @@ protected TypeResolutions.ParamOrdinal queryParamOrdinal() {
     }
 
     @Override
-    protected Query translate(TranslatorHandler handler) {
-        return new TermQuery(source(), ((FieldAttribute) field()).name(), queryAsObject());
+    protected Query translate(TranslatorHandler handler, FoldContext foldContext) {
+        return new TermQuery(source(), ((FieldAttribute) field()).name(), queryAsObject(foldContext));
     }
 
     @Override

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

@@ -184,12 +184,12 @@ public boolean translatable(LucenePushdownPredicates pushdownPredicates) {
     }
 
     @Override
-    public Query asQuery(TranslatorHandler handler) {
+    public Query asQuery(TranslatorHandler handler, FoldContext foldContext) {
         var fa = LucenePushdownPredicates.checkIsFieldAttribute(ipField);
         Check.isTrue(Expressions.foldable(matches), "Expected foldable matches, but got [{}]", matches);
 
         String targetFieldName = handler.nameOf(fa.exactAttribute());
-        Set<Object> set = new LinkedHashSet<>(Expressions.fold(FoldContext.small() /* TODO remove me */, matches));
+        Set<Object> set = new LinkedHashSet<>(Expressions.fold(foldContext, matches));
 
         return new TermsQuery(source(), targetFieldName, set);
     }

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

@@ -183,13 +183,13 @@ public boolean translatable(LucenePushdownPredicates pushdownPredicates) {
     }
 
     @Override
-    public Query asQuery(TranslatorHandler handler) {
+    public Query asQuery(TranslatorHandler handler, FoldContext foldContext) {
         if (left().foldable()) {
             checkSpatialRelatesFunction(left(), queryRelation());
-            return translate(handler, right(), left());
+            return translate(handler, foldContext, right(), left());
         } else {
             checkSpatialRelatesFunction(right(), queryRelation());
-            return translate(handler, left(), right());
+            return translate(handler, foldContext, left(), right());
         }
 
     }
@@ -205,12 +205,17 @@ private static void checkSpatialRelatesFunction(Expression constantExpression, S
         );
     }
 
-    private Query translate(TranslatorHandler handler, Expression spatialExpression, Expression constantExpression) {
+    private Query translate(
+        TranslatorHandler handler,
+        FoldContext foldContext,
+        Expression spatialExpression,
+        Expression constantExpression
+    ) {
         TypedAttribute attribute = LucenePushdownPredicates.checkIsPushableAttribute(spatialExpression);
         String name = handler.nameOf(attribute);
 
         try {
-            Geometry shape = SpatialRelatesUtils.makeGeometryFromLiteral(FoldContext.small() /* TODO remove me */, constantExpression);
+            Geometry shape = SpatialRelatesUtils.makeGeometryFromLiteral(foldContext, constantExpression);
             return new SpatialRelatesQuery(source(), name, queryRelation(), shape, attribute.dataType());
         } catch (IllegalArgumentException e) {
             throw new QlIllegalArgumentException(e.getMessage(), e);