ESQL: Track memory in evaluators (#133392) (#133587)

If you write very very large ESQL queries you can spend a lot of memory
on the expression evaluators themselves. You can certainly do it in real
life, but our tests do something like:
```
FROM foo
| EVAL a0001 = n + 1
| EVAL a0002 = a0001 + 1
| EVAL a0003 = a0002 + 1
...
| EVAL a5000 = a4999 + 1
| STATS MAX(a5000)
```

Each evaluator costs like 200 bytes a pop. For thousands of evaluators
this adds up. So! We have to track it.

Nhat had suggested charging a flat 200 bytes a pop. I thought about it
and decided that it'd be pretty easy to get the actual size. Most of the
evaluators are generated and it's a fairly small generated change to
pick that up. So I did.

We *do* build the evaluators before we cost them, but that's fine
because they are very very small. So long as we account for them, I
think it's safe.

Author: nik9000

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

@@ -666,6 +666,11 @@ public Block eval(Page page) {
                 return mask;
             }
 
+            @Override
+            public long baseRamBytesUsed() {
+                return 0;
+            }
+
             @Override
             public void close() {
                 mask.close();

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

@@ -144,7 +144,7 @@ static void selfTest() {
     public String operation;
 
     private static Operator operator(String operation) {
-        return new EvalOperator(driverContext.blockFactory(), evaluator(operation));
+        return new EvalOperator(driverContext, evaluator(operation));
     }
 
     private static EvalOperator.ExpressionEvaluator evaluator(String operation) {

docs/changelog/133392.yaml

@@ -0,0 +1,5 @@
+pr: 133392
+summary: Track memory in evaluators
+area: ES|QL
+type: bug
+issues: []

docs/reference/query-languages/esql/_snippets/functions/description/md5.md

@@ -6,7 +6,9 @@
  */
 package org.elasticsearch.xpack.esql.core.expression;
 
+import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -35,7 +37,9 @@
 /**
  * Literal or constant.
  */
-public class Literal extends LeafExpression {
+public class Literal extends LeafExpression implements Accountable {
+    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(Literal.class);
+
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Expression.class,
         "Literal",
@@ -169,6 +173,17 @@ public String nodeString() {
         return toString() + "[" + dataType + "]";
     }
 
+    @Override
+    public long ramBytesUsed() {
+        long ramBytesUsed = BASE_RAM_BYTES_USED;
+        if (value instanceof BytesRef b) {
+            ramBytesUsed += b.length;
+        } else {
+            ramBytesUsed += RamUsageEstimator.sizeOfObject(value);
+        }
+        return ramBytesUsed;
+    }
+
     /**
      * Utility method for creating a literal out of a foldable expression.
      * Throws an exception if the expression is not foldable.

docs/reference/query-languages/esql/_snippets/functions/parameters/categorize.md

@@ -23,6 +23,7 @@
 import javax.lang.model.type.TypeMirror;
 import javax.lang.model.util.Elements;
 
+import static org.elasticsearch.compute.gen.EvaluatorImplementer.baseRamBytesUsed;
 import static org.elasticsearch.compute.gen.Methods.buildFromFactory;
 import static org.elasticsearch.compute.gen.Methods.getMethod;
 import static org.elasticsearch.compute.gen.Types.ABSTRACT_CONVERT_FUNCTION_EVALUATOR;
@@ -98,6 +99,7 @@ private TypeSpec type() {
         builder.addJavadoc("This class is generated. Edit {@code " + getClass().getSimpleName() + "} instead.");
         builder.addModifiers(Modifier.PUBLIC, Modifier.FINAL);
         builder.superclass(ABSTRACT_CONVERT_FUNCTION_EVALUATOR);
+        builder.addField(baseRamBytesUsed(implementation));
 
         for (EvaluatorImplementer.ProcessFunctionArg a : processFunction.args) {
             a.declareField(builder);
@@ -113,6 +115,7 @@ private TypeSpec type() {
         }
         builder.addMethod(processFunction.toStringMethod(implementation));
         builder.addMethod(processFunction.close());
+        builder.addMethod(processFunction.baseRamBytesUsed());
         builder.addType(factory());
         return builder.build();
     }

docs/reference/query-languages/esql/kibana/definition/functions/md5.json

@@ -9,6 +9,7 @@
 
 import com.squareup.javapoet.ArrayTypeName;
 import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
@@ -47,6 +48,7 @@
 import static org.elasticsearch.compute.gen.Types.INT_BLOCK;
 import static org.elasticsearch.compute.gen.Types.LONG_BLOCK;
 import static org.elasticsearch.compute.gen.Types.PAGE;
+import static org.elasticsearch.compute.gen.Types.RAM_USAGE_ESIMATOR;
 import static org.elasticsearch.compute.gen.Types.RELEASABLE;
 import static org.elasticsearch.compute.gen.Types.RELEASABLES;
 import static org.elasticsearch.compute.gen.Types.SOURCE;
@@ -96,6 +98,7 @@ private TypeSpec type() {
         builder.addJavadoc("This class is generated. Edit {@code " + getClass().getSimpleName() + "} instead.");
         builder.addModifiers(Modifier.PUBLIC, Modifier.FINAL);
         builder.addSuperinterface(EXPRESSION_EVALUATOR);
+        builder.addField(baseRamBytesUsed(implementation));
         builder.addType(factory());
 
         builder.addField(SOURCE, "source", Modifier.PRIVATE, Modifier.FINAL);
@@ -106,6 +109,7 @@ private TypeSpec type() {
 
         builder.addMethod(ctor());
         builder.addMethod(eval());
+        builder.addMethod(processFunction.baseRamBytesUsed());
 
         if (processOutputsMultivalued) {
             if (processFunction.args.stream().anyMatch(x -> x instanceof FixedProcessFunctionArg == false)) {
@@ -123,6 +127,19 @@ private TypeSpec type() {
         return builder.build();
     }
 
+    static FieldSpec baseRamBytesUsed(ClassName implementation) {
+        FieldSpec.Builder builder = FieldSpec.builder(
+            TypeName.LONG,
+            "BASE_RAM_BYTES_USED",
+            Modifier.PRIVATE,
+            Modifier.STATIC,
+            Modifier.FINAL
+        );
+        builder.initializer("$T.shallowSizeOfInstance($T.class)", RAM_USAGE_ESIMATOR, implementation);
+
+        return builder.build();
+    }
+
     private MethodSpec ctor() {
         MethodSpec.Builder builder = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
         builder.addParameter(SOURCE, "source");
@@ -411,6 +428,11 @@ interface ProcessFunctionArg {
          * The string to close this argument or {@code null}.
          */
         String closeInvocation();
+
+        /**
+         * Invokes {@code baseRamBytesUsed} on sub-expressions an
+         */
+        void sumBaseRamBytesUsed(MethodSpec.Builder builder);
     }
 
     record StandardProcessFunctionArg(TypeName type, String name) implements ProcessFunctionArg {
@@ -535,6 +557,11 @@ public void buildToStringInvocation(StringBuilder pattern, List<Object> args, St
         public String closeInvocation() {
             return name;
         }
+
+        @Override
+        public void sumBaseRamBytesUsed(MethodSpec.Builder builder) {
+            builder.addStatement("baseRamBytesUsed += $L.baseRamBytesUsed()", name);
+        }
     }
 
     private record ArrayProcessFunctionArg(TypeName componentType, String name) implements ProcessFunctionArg {
@@ -667,6 +694,13 @@ public void buildToStringInvocation(StringBuilder pattern, List<Object> args, St
         public String closeInvocation() {
             return "() -> Releasables.close(" + name + ")";
         }
+
+        @Override
+        public void sumBaseRamBytesUsed(MethodSpec.Builder builder) {
+            builder.beginControlFlow("for ($T e : $L)", EXPRESSION_EVALUATOR, name);
+            builder.addStatement("baseRamBytesUsed += e.baseRamBytesUsed()");
+            builder.endControlFlow();
+        }
     }
 
     record FixedProcessFunctionArg(TypeName type, String name, boolean includeInToString, Scope scope, boolean releasable)
@@ -769,6 +803,9 @@ public void buildToStringInvocation(StringBuilder pattern, List<Object> args, St
         public String closeInvocation() {
             return releasable ? name : null;
         }
+
+        @Override
+        public void sumBaseRamBytesUsed(MethodSpec.Builder builder) {}
     }
 
     private record BuilderProcessFunctionArg(ClassName type, String name) implements ProcessFunctionArg {
@@ -853,6 +890,9 @@ public void buildToStringInvocation(StringBuilder pattern, List<Object> args, St
         public String closeInvocation() {
             return null;
         }
+
+        @Override
+        public void sumBaseRamBytesUsed(MethodSpec.Builder builder) {}
     }
 
     private record BlockProcessFunctionArg(TypeName type, String name) implements ProcessFunctionArg {
@@ -940,6 +980,11 @@ public void buildToStringInvocation(StringBuilder pattern, List<Object> args, St
         public String closeInvocation() {
             return name;
         }
+
+        @Override
+        public void sumBaseRamBytesUsed(MethodSpec.Builder builder) {
+            builder.addStatement("baseRamBytesUsed += $L.baseRamBytesUsed()", name);
+        }
     }
 
     static class ProcessFunction {
@@ -1085,6 +1130,18 @@ MethodSpec close() {
             }
             return builder.build();
         }
+
+        MethodSpec baseRamBytesUsed() {
+            MethodSpec.Builder builder = MethodSpec.methodBuilder("baseRamBytesUsed").addAnnotation(Override.class);
+            builder.addModifiers(Modifier.PUBLIC).returns(TypeName.LONG);
+
+            builder.addStatement("long baseRamBytesUsed = BASE_RAM_BYTES_USED");
+            for (ProcessFunctionArg arg : args) {
+                arg.sumBaseRamBytesUsed(builder);
+            }
+            builder.addStatement("return baseRamBytesUsed");
+            return builder.build();
+        }
     }
 
     static boolean isBlockType(TypeName type) {