ESQL: Limit memory usage of fold (#118602)

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.core.TimeValue;
 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.Literal;
 import org.elasticsearch.xpack.esql.core.expression.predicate.regex.RLikePattern;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -71,12 +72,11 @@ public class EvalBenchmark {
         BigArrays.NON_RECYCLING_INSTANCE
     );
 
+    private static final FoldContext FOLD_CONTEXT = FoldContext.small();
+
     private static final int BLOCK_LENGTH = 8 * 1024;
 
-    static final DriverContext driverContext = new DriverContext(
-        BigArrays.NON_RECYCLING_INSTANCE,
-        BlockFactory.getInstance(new NoopCircuitBreaker("noop"), BigArrays.NON_RECYCLING_INSTANCE)
-    );
+    static final DriverContext driverContext = new DriverContext(BigArrays.NON_RECYCLING_INSTANCE, blockFactory);
 
     static {
         // Smoke test all the expected values and force loading subclasses more like prod
@@ -114,18 +114,20 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
         return switch (operation) {
             case "abs" -> {
                 FieldAttribute longField = longField();
-                yield EvalMapper.toEvaluator(new Abs(Source.EMPTY, longField), layout(longField)).get(driverContext);
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, new Abs(Source.EMPTY, longField), layout(longField)).get(driverContext);
             }
             case "add" -> {
                 FieldAttribute longField = longField();
                 yield EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
                     new Add(Source.EMPTY, longField, new Literal(Source.EMPTY, 1L, DataType.LONG)),
                     layout(longField)
                 ).get(driverContext);
             }
             case "add_double" -> {
                 FieldAttribute doubleField = doubleField();
                 yield EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
                     new Add(Source.EMPTY, doubleField, new Literal(Source.EMPTY, 1D, DataType.DOUBLE)),
                     layout(doubleField)
                 ).get(driverContext);
@@ -140,7 +142,8 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
                     lhs = new Add(Source.EMPTY, lhs, new Literal(Source.EMPTY, 1L, DataType.LONG));
                     rhs = new Add(Source.EMPTY, rhs, new Literal(Source.EMPTY, 1L, DataType.LONG));
                 }
-                yield EvalMapper.toEvaluator(new Case(Source.EMPTY, condition, List.of(lhs, rhs)), layout(f1, f2)).get(driverContext);
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, new Case(Source.EMPTY, condition, List.of(lhs, rhs)), layout(f1, f2))
+                    .get(driverContext);
             }
             case "date_trunc" -> {
                 FieldAttribute timestamp = new FieldAttribute(
@@ -149,35 +152,37 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
                     new EsField("timestamp", DataType.DATETIME, Map.of(), true)
                 );
                 yield EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
                     new DateTrunc(Source.EMPTY, new Literal(Source.EMPTY, Duration.ofHours(24), DataType.TIME_DURATION), timestamp),
                     layout(timestamp)
                 ).get(driverContext);
             }
             case "equal_to_const" -> {
                 FieldAttribute longField = longField();
                 yield EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
                     new Equals(Source.EMPTY, longField, new Literal(Source.EMPTY, 100_000L, DataType.LONG)),
                     layout(longField)
                 ).get(driverContext);
             }
             case "long_equal_to_long" -> {
                 FieldAttribute lhs = longField();
                 FieldAttribute rhs = longField();
-                yield EvalMapper.toEvaluator(new Equals(Source.EMPTY, lhs, rhs), layout(lhs, rhs)).get(driverContext);
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, new Equals(Source.EMPTY, lhs, rhs), layout(lhs, rhs)).get(driverContext);
             }
             case "long_equal_to_int" -> {
                 FieldAttribute lhs = longField();
                 FieldAttribute rhs = intField();
-                yield EvalMapper.toEvaluator(new Equals(Source.EMPTY, lhs, rhs), layout(lhs, rhs)).get(driverContext);
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, new Equals(Source.EMPTY, lhs, rhs), layout(lhs, rhs)).get(driverContext);
             }
             case "mv_min", "mv_min_ascending" -> {
                 FieldAttribute longField = longField();
-                yield EvalMapper.toEvaluator(new MvMin(Source.EMPTY, longField), layout(longField)).get(driverContext);
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, new MvMin(Source.EMPTY, longField), layout(longField)).get(driverContext);
             }
             case "rlike" -> {
                 FieldAttribute keywordField = keywordField();
                 RLike rlike = new RLike(Source.EMPTY, keywordField, new RLikePattern(".ar"));
-                yield EvalMapper.toEvaluator(rlike, layout(keywordField)).get(driverContext);
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, rlike, layout(keywordField)).get(driverContext);
             }
             default -> throw new UnsupportedOperationException();
         };

docs/changelog/118602.yaml

@@ -0,0 +1,5 @@
+pr: 118602
+summary: Limit memory usage of `fold`
+area: ES|QL
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/common/settings/Setting.java

@@ -1740,7 +1740,7 @@ public static <T extends Enum<T>> Setting<T> enumSetting(
      *
      * @param key the key for the setting
      * @param defaultValue the default value for this setting
-     * @param properties properties properties for this setting like scope, filtering...
+     * @param properties properties for this setting like scope, filtering...
      * @return the setting object
      */
     public static Setting<ByteSizeValue> memorySizeSetting(String key, ByteSizeValue defaultValue, Property... properties) {

test/external-modules/esql-heap-attack/src/javaRestTest/java/org/elasticsearch/xpack/esql/heap_attack/HeapAttackIT.java

@@ -194,6 +194,16 @@ private void assertCircuitBreaks(ThrowingRunnable r) throws IOException {
         );
     }
 
+    private void assertFoldCircuitBreaks(ThrowingRunnable r) throws IOException {
+        ResponseException e = expectThrows(ResponseException.class, r);
+        Map<?, ?> map = responseAsMap(e.getResponse());
+        logger.info("expected fold circuit breaking {}", map);
+        assertMap(
+            map,
+            matchesMap().entry("status", 400).entry("error", matchesMap().extraOk().entry("type", "fold_too_much_memory_exception"))
+        );
+    }
+
     private void assertParseFailure(ThrowingRunnable r) throws IOException {
         ResponseException e = expectThrows(ResponseException.class, r);
         Map<?, ?> map = responseAsMap(e.getResponse());
@@ -325,11 +335,23 @@ public void testManyConcatFromRow() throws IOException {
         assertManyStrings(resp, strings);
     }
 
+    /**
+     * Hits a circuit breaker by building many moderately long strings.
+     */
+    public void testHugeManyConcatFromRow() throws IOException {
+        assertFoldCircuitBreaks(
+            () -> manyConcat(
+                "ROW a=9999999999999, b=99999999999999999, c=99999999999999999, d=99999999999999999, e=99999999999999999",
+                5000
+            )
+        );
+    }
+
     /**
      * Fails to parse a huge huge query.
      */
     public void testHugeHugeManyConcatFromRow() throws IOException {
-        assertParseFailure(() -> manyConcat("ROW a=9999, b=9999, c=9999, d=9999, e=9999", 50000));
+        assertParseFailure(() -> manyConcat("ROW a=9999, b=9999, c=9999, d=9999, e=9999", 6000));
     }
 
     /**
@@ -387,13 +409,20 @@ public void testHugeManyRepeat() throws IOException {
      * Returns many moderately long strings.
      */
     public void testManyRepeatFromRow() throws IOException {
-        int strings = 10000;
+        int strings = 300;
         Response resp = manyRepeat("ROW a = 99", strings);
         assertManyStrings(resp, strings);
     }
 
     /**
-     * Fails to parse a huge huge query.
+     * Hits a circuit breaker by building many moderately long strings.
+     */
+    public void testHugeManyRepeatFromRow() throws IOException {
+        assertFoldCircuitBreaks(() -> manyRepeat("ROW a = 99", 400));
+    }
+
+    /**
+     * Fails to parse a huge, huge query.
      */
     public void testHugeHugeManyRepeatFromRow() throws IOException {
         assertParseFailure(() -> manyRepeat("ROW a = 99", 100000));

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

@@ -78,12 +78,20 @@ public Expression(Source source, List<Expression> children) {
         super(source, children);
     }
 
-    // whether the expression can be evaluated statically (folded) or not
+    /**
+     * Whether the expression can be evaluated statically, aka "folded", or not.
+     */
     public boolean foldable() {
         return false;
     }
 
-    public Object fold() {
+    /**
+     * Evaluate this expression statically to a constant. It is an error to call
+     * this if {@link #foldable} returns false.
+     */
+    public Object fold(FoldContext ctx) {
+        // TODO After removing FoldContext.unbounded from non-test code examine all calls
+        // for places we should use instanceof Literal instead
         throw new QlIllegalArgumentException("Should not fold expression");
     }
 

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

@@ -107,10 +107,10 @@ public static boolean foldable(List<? extends Expression> exps) {
         return true;
     }
 
-    public static List<Object> fold(List<? extends Expression> exps) {
+    public static List<Object> fold(FoldContext ctx, List<? extends Expression> exps) {
         List<Object> folded = new ArrayList<>(exps.size());
         for (Expression exp : exps) {
-            folded.add(exp.fold());
+            folded.add(exp.fold(ctx));
         }
 
         return folded;
@@ -135,7 +135,7 @@ public static String name(Expression e) {
     /**
      * Is this {@linkplain Expression} <strong>guaranteed</strong> to have
      * only the {@code null} value. {@linkplain Expression}s that
-     * {@link Expression#fold()} to {@code null} <strong>may</strong>
+     * {@link Expression#fold} to {@code null} <strong>may</strong>
      * return {@code false} here, but should <strong>eventually</strong> be folded
      * into a {@link Literal} containing {@code null} which will return
      * {@code true} from here.