Add row limit to Completion and Rerank plan.

Author: afoucret

server/src/main/resources/transport/definitions/referable/esql_completion_usage_limit.csv

@@ -0,0 +1 @@
+9223000

server/src/main/resources/transport/definitions/referable/esql_rerank_usage_limit.csv

@@ -0,0 +1 @@
+9222000

server/src/main/resources/transport/upper_bounds/9.3.csv

@@ -1 +1 @@
-index_created_transport_version,9221000
+esql_completion_usage_limit,9223000

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -600,7 +600,7 @@ private LogicalPlan resolveCompletion(Completion p, List<Attribute> childrenOutp
                 prompt = prompt.transformUp(UnresolvedAttribute.class, ua -> maybeResolveAttribute(ua, childrenOutput));
             }
 
-            return new Completion(p.source(), p.child(), p.inferenceId(), prompt, targetField);
+            return new Completion(p.source(), p.child(), p.inferenceId(), prompt, targetField, p.rowLimit());
         }
 
         private LogicalPlan resolveMvExpand(MvExpand p, List<Attribute> childrenOutput) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceCommandConfig.java

@@ -9,15 +9,17 @@
 
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
 
 import java.util.Map;
 
 public record InferenceCommandConfig(boolean enabled, int rowLimit) {
 
     public static final Setting<Boolean> COMPLETION_ENABLED_SETTING = commandEnabledSetting("completion");
-    public static final Setting<Integer> COMPLETION_ROW_LIMIT_SETTING = rowLimitSetting("completion", 100);
+    public static final Setting<Integer> COMPLETION_ROW_LIMIT_SETTING = rowLimitSetting("completion", Completion.DEFAULT_MAX_ROW_LIMIT);
     public static final Setting<Boolean> RERANK_ENABLED_SETTING = commandEnabledSetting("rerank");
-    public static final Setting<Integer> RERANK_ROW_LIMIT_SETTING = rowLimitSetting("rerank", 1000);
+    public static final Setting<Integer> RERANK_ROW_LIMIT_SETTING = rowLimitSetting("rerank", Rerank.DEFAULT_MAX_ROW_LIMIT);
 
     public static InferenceCommandConfig completionCommandConfig(Settings settings) {
         return new InferenceCommandConfig(COMPLETION_ENABLED_SETTING.get(settings), COMPLETION_ROW_LIMIT_SETTING.get(settings));
@@ -52,4 +54,5 @@ private static Setting<Integer> rowLimitSetting(String commandName, int defaultV
             Setting.Property.Dynamic
         );
     }
+
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceCommandConfigProvider.java

@@ -0,0 +1,13 @@
+/*
+ * 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.inference;
+
+@FunctionalInterface
+public interface InferenceCommandConfigProvider {
+    InferenceCommandConfig get(String commandName);
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlParser.java

@@ -175,7 +175,10 @@ private <T> T invokeParser(
                 log.trace("Parse tree: {}", tree.toStringTree());
             }
 
-            return result.apply(new AstBuilder(new ExpressionBuilder.ParsingContext(params, metrics)), tree);
+            return result.apply(
+                new AstBuilder(new ExpressionBuilder.ParsingContext(params, metrics, config::inferenceCommandConfig)),
+                tree
+            );
         } catch (StackOverflowError e) {
             throw new ParsingException("ESQL statement is too large, causing stack overflow when generating the parsing tree: [{}]", query);
             // likely thrown by an invalid popMode (such as extra closing parenthesis)

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/ExpressionBuilder.java

@@ -67,6 +67,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.InsensitiveEquals;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
+import org.elasticsearch.xpack.esql.inference.InferenceCommandConfigProvider;
 import org.elasticsearch.xpack.esql.telemetry.PlanTelemetry;
 import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
 
@@ -124,7 +125,11 @@ public abstract class ExpressionBuilder extends IdentifierBuilder {
 
     protected final ParsingContext context;
 
-    public record ParsingContext(QueryParams params, PlanTelemetry telemetry) {}
+    public record ParsingContext(
+        QueryParams params,
+        PlanTelemetry telemetry,
+        InferenceCommandConfigProvider inferenceCommandConfigProvider
+    ) {}
 
     ExpressionBuilder(ParsingContext context) {
         this.context = context;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -1109,9 +1109,14 @@ public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
             throw qualifiersUnsupportedInFieldDefinitions(scoreAttribute.source(), ctx.targetField.getText());
         }
 
+        Literal rowLimit = Literal.integer(Source.EMPTY, context.inferenceCommandConfigProvider().get("rerank").rowLimit());
+
         return p -> {
             checkForRemoteClusters(p, source, "RERANK");
-            return applyRerankOptions(new Rerank(source, p, queryText, rerankFields, scoreAttribute), ctx.commandNamedParameters());
+            return applyRerankOptions(
+                new Rerank(source, p, queryText, rerankFields, scoreAttribute, rowLimit),
+                ctx.commandNamedParameters()
+            );
         };
     }
 
@@ -1150,9 +1155,11 @@ public PlanFactory visitCompletionCommand(EsqlBaseParser.CompletionCommandContex
             throw qualifiersUnsupportedInFieldDefinitions(targetField.source(), ctx.targetField.getText());
         }
 
+        Literal rowLimit = Literal.integer(Source.EMPTY, context.inferenceCommandConfigProvider().get("completion").rowLimit());
+
         return p -> {
             checkForRemoteClusters(p, source, "COMPLETION");
-            return applyCompletionOptions(new Completion(source, p, prompt, targetField), ctx.commandNamedParameters());
+            return applyCompletionOptions(new Completion(source, p, prompt, targetField, rowLimit), ctx.commandNamedParameters());
         };
     }
 

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.plan.logical.inference;
 
+import org.elasticsearch.TransportVersion;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -42,18 +43,32 @@ public class Completion extends InferencePlan<Completion> implements TelemetryAw
         "Completion",
         Completion::new
     );
+    public static final int DEFAULT_MAX_ROW_LIMIT = 100;
+
+    private static final TransportVersion ESQL_COMPLETION_USAGE_LIMIT = TransportVersion.fromName("esql_completion_usage_limit");
+
     private final Expression prompt;
     private final Attribute targetField;
+    private final Expression rowLimit;
+
     private List<Attribute> lazyOutput;
 
-    public Completion(Source source, LogicalPlan p, Expression prompt, Attribute targetField) {
-        this(source, p, Literal.keyword(Source.EMPTY, DEFAULT_OUTPUT_FIELD_NAME), prompt, targetField);
+    public Completion(Source source, LogicalPlan p, Expression prompt, Attribute targetField, Expression rowLimit) {
+        this(source, p, Literal.keyword(Source.EMPTY, DEFAULT_OUTPUT_FIELD_NAME), prompt, targetField, rowLimit);
     }
 
-    public Completion(Source source, LogicalPlan child, Expression inferenceId, Expression prompt, Attribute targetField) {
+    public Completion(
+        Source source,
+        LogicalPlan child,
+        Expression inferenceId,
+        Expression prompt,
+        Attribute targetField,
+        Expression rowLimit
+    ) {
         super(source, child, inferenceId);
         this.prompt = prompt;
         this.targetField = targetField;
+        this.rowLimit = rowLimit;
     }
 
     public Completion(StreamInput in) throws IOException {
@@ -62,7 +77,10 @@ public Completion(StreamInput in) throws IOException {
             in.readNamedWriteable(LogicalPlan.class),
             in.readNamedWriteable(Expression.class),
             in.readNamedWriteable(Expression.class),
-            in.readNamedWriteable(Attribute.class)
+            in.readNamedWriteable(Attribute.class),
+            in.getTransportVersion().supports(ESQL_COMPLETION_USAGE_LIMIT)
+                ? in.readNamedWriteable(Expression.class)
+                : Literal.integer(Source.EMPTY, DEFAULT_MAX_ROW_LIMIT)
         );
     }
 
@@ -71,6 +89,9 @@ public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeNamedWriteable(prompt);
         out.writeNamedWriteable(targetField);
+        if (out.getTransportVersion().supports(ESQL_COMPLETION_USAGE_LIMIT)) {
+            out.writeNamedWriteable(rowLimit);
+        }
     }
 
     public Expression prompt() {
@@ -81,18 +102,22 @@ public Attribute targetField() {
         return targetField;
     }
 
+    public Expression rowLimit() {
+        return rowLimit;
+    }
+
     @Override
     public Completion withInferenceId(Expression newInferenceId) {
         if (inferenceId().equals(newInferenceId)) {
             return this;
         }
 
-        return new Completion(source(), child(), newInferenceId, prompt, targetField);
+        return new Completion(source(), child(), newInferenceId, prompt, targetField, rowLimit);
     }
 
     @Override
     public Completion replaceChild(LogicalPlan newChild) {
-        return new Completion(source(), newChild, inferenceId(), prompt, targetField);
+        return new Completion(source(), newChild, inferenceId(), prompt, targetField, rowLimit);
     }
 
     @Override
@@ -122,7 +147,7 @@ public List<Attribute> generatedAttributes() {
     @Override
     public Completion withGeneratedNames(List<String> newNames) {
         checkNumberOfNewNames(newNames);
-        return new Completion(source(), child(), inferenceId(), prompt, this.renameTargetField(newNames.get(0)));
+        return new Completion(source(), child(), inferenceId(), prompt, this.renameTargetField(newNames.get(0)), rowLimit);
     }
 
     private Attribute renameTargetField(String newName) {
@@ -157,7 +182,7 @@ public void postAnalysisVerification(Failures failures) {
 
     @Override
     protected NodeInfo<? extends LogicalPlan> info() {
-        return NodeInfo.create(this, Completion::new, child(), inferenceId(), prompt, targetField);
+        return NodeInfo.create(this, Completion::new, child(), inferenceId(), prompt, targetField, rowLimit);
     }
 
     @Override
@@ -167,11 +192,13 @@ public boolean equals(Object o) {
         if (super.equals(o) == false) return false;
         Completion completion = (Completion) o;
 
-        return Objects.equals(prompt, completion.prompt) && Objects.equals(targetField, completion.targetField);
+        return Objects.equals(prompt, completion.prompt)
+            && Objects.equals(targetField, completion.targetField)
+            && Objects.equals(rowLimit, completion.rowLimit);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), prompt, targetField);
+        return Objects.hash(super.hashCode(), prompt, targetField, rowLimit);
     }
 }