[ES|QL] COMPLETION command grammar and logical plan (#126319)

afoucret · afoucret · commit ac63795e929f · 2025-06-04T16:19:41.000+02:00
diff --git a/docs/changelog/126319.yaml b/docs/changelog/126319.yaml
@@ -0,0 +1,5 @@
+pr: 126319
+summary: COMPLETION command grammar and logical plan
+area: ES|QL
+type: feature
+issues: []
diff --git a/x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4 b/x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4
@@ -54,6 +54,7 @@ processingCommand
     | joinCommand
     | changePointCommand
     // in development
+    | {this.isDevVersion()}? completionCommand
     | {this.isDevVersion()}? inlinestatsCommand
     | {this.isDevVersion()}? lookupCommand
     | {this.isDevVersion()}? rerankCommand
@@ -371,3 +372,7 @@ joinPredicate
 rerankCommand
     : DEV_RERANK queryText=constant ON fields WITH inferenceId=identifierOrParameter
     ;
+
+completionCommand
+    : DEV_COMPLETION prompt=primaryExpression WITH inferenceId=identifierOrParameter (AS targetField=qualifiedName)?
+    ;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
@@ -771,6 +771,11 @@ public enum Cap {
          */
         RERANK(Build.current().isSnapshot()),
 
+        /**
+         * Support for COMPLETION command
+         */
+        COMPLETION(Build.current().isSnapshot()),
+
         /**
          * Allow mixed numeric types in conditional functions - case, greatest and least
          */
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java
@@ -385,9 +385,9 @@ private static NamedExpression createEnrichFieldExpression(
         }
     }
 
-    private static class ResolveInference extends ParameterizedAnalyzerRule<InferencePlan, AnalyzerContext> {
+    private static class ResolveInference extends ParameterizedAnalyzerRule<InferencePlan<?>, AnalyzerContext> {
         @Override
-        protected LogicalPlan rule(InferencePlan plan, AnalyzerContext context) {
+        protected LogicalPlan rule(InferencePlan<?> plan, AnalyzerContext context) {
             assert plan.inferenceId().resolved() && plan.inferenceId().foldable();
 
             String inferenceId = plan.inferenceId().fold(FoldContext.small()).toString();
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/PreAnalyzer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/PreAnalyzer.java
@@ -28,10 +28,10 @@ public static class PreAnalysis {
 
         public final List<TableInfo> indices;
         public final List<Enrich> enriches;
-        public final List<InferencePlan> inferencePlans;
+        public final List<InferencePlan<?>> inferencePlans;
         public final List<TableInfo> lookupIndices;
 
-        public PreAnalysis(List<TableInfo> indices, List<Enrich> enriches, List<InferencePlan> inferencePlans, List<TableInfo> lookupIndices) {
+        public PreAnalysis(List<TableInfo> indices, List<Enrich> enriches, List<InferencePlan<?>> inferencePlans, List<TableInfo> lookupIndices) {
             this.indices = indices;
             this.enriches = enriches;
             this.inferencePlans = inferencePlans;
@@ -51,7 +51,7 @@ protected PreAnalysis doPreAnalyze(LogicalPlan plan) {
         List<TableInfo> indices = new ArrayList<>();
         List<Enrich> unresolvedEnriches = new ArrayList<>();
         List<TableInfo> lookupIndices = new ArrayList<>();
-        List<InferencePlan> unresolvedInferencePlans = new ArrayList<>();
+        List<InferencePlan<?>> unresolvedInferencePlans = new ArrayList<>();
 
         plan.forEachUp(UnresolvedRelation.class, p -> {
             List<TableInfo> list = p.indexMode() == IndexMode.LOOKUP ? lookupIndices : indices;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceRunner.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceRunner.java
@@ -33,7 +33,7 @@ public ThreadContext getThreadContext() {
         return client.threadPool().getThreadContext();
     }
 
-    public void resolveInferenceIds(List<InferencePlan> plans, ActionListener<InferenceResolution> listener) {
+    public void resolveInferenceIds(List<InferencePlan<?>> plans, ActionListener<InferenceResolution> listener) {
         resolveInferenceIds(plans.stream().map(InferenceRunner::planInferenceId).collect(Collectors.toSet()), listener);
 
     }
@@ -68,7 +68,7 @@ private void resolveInferenceIds(Set<String> inferenceIds, ActionListener<Infere
         }
     }
 
-    private static String planInferenceId(InferencePlan plan) {
+    private static String planInferenceId(InferencePlan<?> plan) {
         return plan.inferenceId().fold(FoldContext.small()).toString();
     }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java
@@ -61,6 +61,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Rename;
 import org.elasticsearch.xpack.esql.plan.logical.Row;
 import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
 import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
 import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
 import org.elasticsearch.xpack.esql.plan.logical.show.ShowInfo;
@@ -643,12 +644,7 @@ public PlanFactory visitJoinCommand(EsqlBaseParser.JoinCommandContext ctx) {
 
     @Override
     public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
-        var source = source(ctx);
-
-        if (false == EsqlCapabilities.Cap.RERANK.isEnabled()) {
-            throw new ParsingException(source, "RERANK is in preview and only available in SNAPSHOT build");
-        }
-
+        Source source = source(ctx);
         Expression queryText = expression(ctx.queryText);
         if (queryText instanceof Literal queryTextLiteral && DataType.isString(queryText.dataType())) {
             if (queryTextLiteral.value() == null) {
@@ -669,6 +665,18 @@ public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
         return p -> new Rerank(source, p, inferenceId(ctx.inferenceId), queryText, visitFields(ctx.fields()));
     }
 
+    @Override
+    public PlanFactory visitCompletionCommand(EsqlBaseParser.CompletionCommandContext ctx) {
+        Source source = source(ctx);
+        Expression prompt = expression(ctx.prompt);
+        Literal inferenceId = inferenceId(ctx.inferenceId);
+        Attribute targetField = ctx.targetField == null
+            ? new UnresolvedAttribute(source, Completion.DEFAULT_OUTPUT_FIELD_NAME)
+            : visitQualifiedName(ctx.targetField);
+
+        return p -> new Completion(source, p, inferenceId, prompt, targetField);
+    }
+
     public Literal inferenceId(EsqlBaseParser.IdentifierOrParameterContext ctx) {
         if (ctx.identifier() != null) {
             return new Literal(source(ctx), visitIdentifier(ctx.identifier()), KEYWORD);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java
@@ -22,6 +22,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
 import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
 import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
@@ -65,6 +66,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
     public static List<NamedWriteableRegistry.Entry> logical() {
         return List.of(
             Aggregate.ENTRY,
+            Completion.ENTRY,
             Dissect.ENTRY,
             Enrich.ENTRY,
             EsRelation.ENTRY,
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/Completion.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/Completion.java
@@ -0,0 +1,152 @@
+/*
+ * 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.plan.logical.inference;
+
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.NameId;
+import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.SortAgnostic;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
+
+public class Completion extends InferencePlan<Completion> implements GeneratingPlan<Completion>, SortAgnostic {
+
+    public static final String DEFAULT_OUTPUT_FIELD_NAME = "completion";
+
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
+        LogicalPlan.class,
+        "Completion",
+        Completion::new
+    );
+    private final Expression prompt;
+    private final Attribute targetField;
+    private List<Attribute> lazyOutput;
+
+    public Completion(Source source, LogicalPlan child, Expression inferenceId, Expression prompt, Attribute targetField) {
+        super(source, child, inferenceId);
+        this.prompt = prompt;
+        this.targetField = targetField;
+    }
+
+    public Completion(StreamInput in) throws IOException {
+        this(
+            Source.readFrom((PlanStreamInput) in),
+            in.readNamedWriteable(LogicalPlan.class),
+            in.readNamedWriteable(Expression.class),
+            in.readNamedWriteable(Expression.class),
+            in.readNamedWriteable(Attribute.class)
+        );
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeNamedWriteable(prompt);
+        out.writeNamedWriteable(targetField);
+    }
+
+    public Expression prompt() {
+        return prompt;
+    }
+
+    public Attribute targetField() {
+        return targetField;
+    }
+
+    @Override
+    public Completion withInferenceId(Expression newInferenceId) {
+        return new Completion(source(), child(), newInferenceId, prompt, targetField);
+    }
+
+    @Override
+    public Completion replaceChild(LogicalPlan newChild) {
+        return new Completion(source(), newChild, inferenceId(), prompt, targetField);
+    }
+
+    @Override
+    public TaskType taskType() {
+        return TaskType.COMPLETION;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
+    @Override
+    public List<Attribute> output() {
+        if (lazyOutput == null) {
+            lazyOutput = mergeOutputAttributes(List.of(targetField), child().output());
+        }
+
+        return lazyOutput;
+    }
+
+    @Override
+    public List<Attribute> generatedAttributes() {
+        return List.of(targetField);
+    }
+
+    @Override
+    public Completion withGeneratedNames(List<String> newNames) {
+        checkNumberOfNewNames(newNames);
+        return new Completion(source(), child(), inferenceId(), prompt, this.renameTargetField(newNames.get(0)));
+    }
+
+    private Attribute renameTargetField(String newName) {
+        if (newName.equals(targetField.name())) {
+            return targetField;
+        }
+
+        return targetField.withName(newName).withId(new NameId());
+    }
+
+    @Override
+    protected AttributeSet computeReferences() {
+        return prompt.references();
+    }
+
+    @Override
+    public boolean expressionsResolved() {
+        return super.expressionsResolved() && prompt.resolved();
+    }
+
+    @Override
+    protected NodeInfo<? extends LogicalPlan> info() {
+        return NodeInfo.create(this, Completion::new, child(), inferenceId(), prompt, targetField);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        if (super.equals(o) == false) return false;
+        Completion completion = (Completion) o;
+
+        return Objects.equals(prompt, completion.prompt) && Objects.equals(targetField, completion.targetField);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), prompt, targetField);
+    }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/InferencePlan.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/InferencePlan.java
@@ -18,7 +18,7 @@
 import java.io.IOException;
 import java.util.Objects;
 
-public abstract class InferencePlan extends UnaryPlan {
+public abstract class InferencePlan<PlanType extends InferencePlan<PlanType>> extends UnaryPlan {
 
     private final Expression inferenceId;
 
@@ -48,7 +48,7 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         if (super.equals(o) == false) return false;
-        InferencePlan other = (InferencePlan) o;
+        InferencePlan<?> other = (InferencePlan<?>) o;
         return Objects.equals(inferenceId(), other.inferenceId());
     }
 
@@ -59,9 +59,9 @@ public int hashCode() {
 
     public abstract TaskType taskType();
 
-    public abstract InferencePlan withInferenceId(Expression newInferenceId);
+    public abstract PlanType withInferenceId(Expression newInferenceId);
 
-    public InferencePlan withInferenceResolutionError(String inferenceId, String error) {
+    public PlanType withInferenceResolutionError(String inferenceId, String error) {
         return withInferenceId(new UnresolvedAttribute(inferenceId().source(), inferenceId, error));
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/Rerank.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/Rerank.java
@@ -37,7 +37,7 @@
 import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public class Rerank extends InferencePlan implements SortAgnostic, SurrogateLogicalPlan {
+public class Rerank extends InferencePlan<Rerank> implements SortAgnostic, SurrogateLogicalPlan {
 
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Rerank", Rerank::new);
     private final Attribute scoreAttribute;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java
@@ -546,7 +546,7 @@ private static void resolveFieldNames(LogicalPlan parsed, EnrichResolution enric
     }
 
     private void resolveInferences(
-        List<InferencePlan> inferencePlans,
+        List<InferencePlan<?>> inferencePlans,
         PreAnalysisResult preAnalysisResult,
         ActionListener<PreAnalysisResult> l
     ) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceRunnerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceRunnerTests.java
@@ -36,7 +36,7 @@
 public class InferenceRunnerTests extends ESTestCase {
     public void testResolveInferenceIds() throws Exception {
         InferenceRunner inferenceRunner = new InferenceRunner(mockClient());
-        List<InferencePlan> inferencePlans = List.of(mockInferencePlan("rerank-plan"));
+        List<InferencePlan<?>> inferencePlans = List.of(mockInferencePlan("rerank-plan"));
         SetOnce<InferenceResolution> inferenceResolutionSetOnce = new SetOnce<>();
 
         inferenceRunner.resolveInferenceIds(inferencePlans, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
@@ -53,7 +53,7 @@ public void testResolveInferenceIds() throws Exception {
 
     public void testResolveMultipleInferenceIds() throws Exception {
         InferenceRunner inferenceRunner = new InferenceRunner(mockClient());
-        List<InferencePlan> inferencePlans = List.of(
+        List<InferencePlan<?>> inferencePlans = List.of(
             mockInferencePlan("rerank-plan"),
             mockInferencePlan("rerank-plan"),
             mockInferencePlan("completion-plan")
@@ -81,7 +81,7 @@ public void testResolveMultipleInferenceIds() throws Exception {
 
     public void testResolveMissingInferenceIds() throws Exception {
         InferenceRunner inferenceRunner = new InferenceRunner(mockClient());
-        List<InferencePlan> inferencePlans = List.of(mockInferencePlan("missing-plan"));
+        List<InferencePlan<?>> inferencePlans = List.of(mockInferencePlan("missing-plan"));
 
         SetOnce<InferenceResolution> inferenceResolutionSetOnce = new SetOnce<>();
 
@@ -138,8 +138,8 @@ private static ModelConfigurations mockModelConfig(String inferenceId, TaskType
         return new ModelConfigurations(inferenceId, taskType, randomIdentifier(), mock(ServiceSettings.class));
     }
 
-    private static InferencePlan mockInferencePlan(String inferenceId) {
-        InferencePlan plan = mock(InferencePlan.class);
+    private static InferencePlan<?> mockInferencePlan(String inferenceId) {
+        InferencePlan<?> plan = mock(InferencePlan.class);
         when(plan.inferenceId()).thenReturn(new Literal(Source.EMPTY, inferenceId, DataType.KEYWORD));
         return plan;
     }
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/GrammarInDevelopmentParsingTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/GrammarInDevelopmentParsingTests.java
@@ -31,7 +31,11 @@ public void testDevelopmentMatch() throws Exception {
     }
 
     public void testDevelopmentRerank() {
-        parse("row a = 1 | rerank \"foo\" ON title WITH reranker", "rerank");
+        parse("row a = 1 | rerank \"foo\" on title with reranker", "rerank");
+    }
+
+    public void testDevelopmentCompletion() {
+        parse("row a = 1 | completion concat(\"test\", \"a\") with inferenceId as fieldName", "completion");
     }
 
     void parse(String query, String errorMessage) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/StatementParserTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/StatementParserTests.java
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/inference/CompletionSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/inference/CompletionSerializationTests.java

Original file line number	Diff line number	Diff line change
`@@ -385,9 +385,9 @@ private static NamedExpression createEnrichFieldExpression(`
`385`	`385`	`}`
`386`	`386`	`}`
`387`	`387`
`388`		`- private static class ResolveInference extends ParameterizedAnalyzerRule<InferencePlan, AnalyzerContext> {`
	`388`	`+ private static class ResolveInference extends ParameterizedAnalyzerRule<InferencePlan<?>, AnalyzerContext> {`
`389`	`389`	`@Override`
`390`		`- protected LogicalPlan rule(InferencePlan plan, AnalyzerContext context) {`
	`390`	`+ protected LogicalPlan rule(InferencePlan<?> plan, AnalyzerContext context) {`
`391`	`391`	`assert plan.inferenceId().resolved() && plan.inferenceId().foldable();`
`392`	`392`
`393`	`393`	`String inferenceId = plan.inferenceId().fold(FoldContext.small()).toString();`
Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ public ThreadContext getThreadContext() {`
`33`	`33`	`return client.threadPool().getThreadContext();`
`34`	`34`	`}`
`35`	`35`
`36`		`- public void resolveInferenceIds(List<InferencePlan> plans, ActionListener<InferenceResolution> listener) {`
	`36`	`+ public void resolveInferenceIds(List<InferencePlan<?>> plans, ActionListener<InferenceResolution> listener) {`
`37`	`37`	`resolveInferenceIds(plans.stream().map(InferenceRunner::planInferenceId).collect(Collectors.toSet()), listener);`
`38`	`38`
`39`	`39`	`}`
`@@ -68,7 +68,7 @@ private void resolveInferenceIds(Set<String> inferenceIds, ActionListener<Infere`
`68`	`68`	`}`
`69`	`69`	`}`
`70`	`70`
`71`		`- private static String planInferenceId(InferencePlan plan) {`
	`71`	`+ private static String planInferenceId(InferencePlan<?> plan) {`
`72`	`72`	`return plan.inferenceId().fold(FoldContext.small()).toString();`
`73`	`73`	`}`
`74`	`74`