Rerank command analysis and verification.

Author: afoucret

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -66,6 +66,8 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
 import org.elasticsearch.xpack.esql.index.EsIndex;
+import org.elasticsearch.xpack.esql.inference.InferenceResolution;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.parser.QueryParam;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -374,7 +376,8 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         null,
         mock(ClusterService.class),
         mock(IndexNameExpressionResolver.class),
-        null
+        null,
+        mock(InferenceService.class)
     );
 
     private EsqlTestUtils() {}
@@ -452,6 +455,10 @@ public static EnrichResolution emptyPolicyResolution() {
         return new EnrichResolution();
     }
 
+    public static InferenceResolution emptyInferenceResolution() {
+        return InferenceResolution.EMPTY;
+    }
+    
     public static SearchStats statsForExistingField(String... names) {
         return fieldMatchingExistOrMissing(true, names);
     }

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

@@ -64,6 +64,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.In;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.index.IndexResolution;
+import org.elasticsearch.xpack.esql.inference.ResolvedInference;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 import org.elasticsearch.xpack.esql.plan.IndexPattern;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
@@ -81,6 +82,8 @@
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.Rename;
 import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
+import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
 import org.elasticsearch.xpack.esql.plan.logical.join.JoinConfig;
 import org.elasticsearch.xpack.esql.plan.logical.join.JoinType;
@@ -156,6 +159,7 @@ public class Analyzer extends ParameterizedRuleExecutor<LogicalPlan, AnalyzerCon
             Limiter.ONCE,
             new ResolveTable(),
             new ResolveEnrich(),
+            new ResolveInference(),
             new ResolveLookupTables(),
             new ResolveFunctions(),
             new ResolveForkFunctions()
@@ -393,6 +397,34 @@ private static NamedExpression createEnrichFieldExpression(
         }
     }
 
+    private static class ResolveInference extends ParameterizedAnalyzerRule<InferencePlan, AnalyzerContext> {
+        @Override
+        protected LogicalPlan rule(InferencePlan plan, AnalyzerContext context) {
+            assert plan.inferenceId().resolved() && plan.inferenceId().foldable();
+
+            String inferenceId = plan.inferenceId().fold(FoldContext.small()).toString();
+            ResolvedInference resolvedInference = context.inferenceResolution().getResolvedInference(inferenceId);
+
+            if (resolvedInference != null && resolvedInference.taskType() == plan.taskType()) {
+                return plan;
+            } else if (resolvedInference != null) {
+                String error = "cannot use inference endpoint ["
+                    + inferenceId
+                    + "] with task type ["
+                    + resolvedInference.taskType()
+                    + "] within a "
+                    + plan.nodeName()
+                    + " command. Only inference endpoints with the task type ["
+                    + plan.taskType()
+                    + "] are supported.";
+                return plan.withInferenceResolutionError(inferenceId, error);
+            } else {
+                String error = context.inferenceResolution().getError(inferenceId);
+                return plan.withInferenceResolutionError(inferenceId, error);
+            }
+        }
+    }
+
     private static class ResolveLookupTables extends ParameterizedAnalyzerRule<Lookup, AnalyzerContext> {
 
         @Override
@@ -498,6 +530,10 @@ protected LogicalPlan rule(LogicalPlan plan, AnalyzerContext context) {
                 return resolveFork(f, context);
             }
 
+            if (plan instanceof Rerank r) {
+                return resolveRerank(r, childrenOutput);
+            }
+
             return plan.transformExpressionsOnly(UnresolvedAttribute.class, ua -> maybeResolveAttribute(ua, childrenOutput));
         }
 
@@ -687,6 +723,17 @@ private LogicalPlan resolveFork(Fork fork, AnalyzerContext context) {
             return new Fork(fork.source(), fork.child(), newSubPlans);
         }
 
+        private LogicalPlan resolveRerank(Rerank rerank, List<Attribute> childOutput) {
+            List<Alias> newFields = new ArrayList<>();
+            boolean changed = false;
+            for (Alias field : rerank.rerankFields()) {
+                Alias result = (Alias) field.transformUp(UnresolvedAttribute.class, ua -> resolveAttribute(ua, childOutput));
+                changed |= result != field;
+            }
+
+            return changed ? new Rerank(rerank.source(), rerank.child(), rerank.inferenceId(), rerank.queryText(), newFields) : rerank;
+        }
+
         private List<Attribute> resolveUsingColumns(List<Attribute> cols, List<Attribute> output, String side) {
             List<Attribute> resolved = new ArrayList<>(cols.size());
             for (Attribute col : cols) {

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.index.IndexResolution;
+import org.elasticsearch.xpack.esql.inference.InferenceResolution;
 import org.elasticsearch.xpack.esql.session.Configuration;
 
 import java.util.Map;
@@ -18,16 +19,18 @@ public record AnalyzerContext(
     EsqlFunctionRegistry functionRegistry,
     IndexResolution indexResolution,
     Map<String, IndexResolution> lookupResolution,
-    EnrichResolution enrichResolution
+    EnrichResolution enrichResolution,
+    InferenceResolution inferenceResolution
 ) {
     // Currently for tests only, since most do not test lookups
     // TODO: make this even simpler, remove the enrichResolution for tests that do not require it (most tests)
     public AnalyzerContext(
         Configuration configuration,
         EsqlFunctionRegistry functionRegistry,
         IndexResolution indexResolution,
-        EnrichResolution enrichResolution
+        EnrichResolution enrichResolution,
+        InferenceResolution inferenceResolution
     ) {
-        this(configuration, functionRegistry, indexResolution, Map.of(), enrichResolution);
+        this(configuration, functionRegistry, indexResolution, Map.of(), enrichResolution, inferenceResolution);
     }
 }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
+import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -23,15 +24,22 @@
 public class PreAnalyzer {
 
     public static class PreAnalysis {
-        public static final PreAnalysis EMPTY = new PreAnalysis(emptyList(), emptyList(), emptyList());
+        public static final PreAnalysis EMPTY = new PreAnalysis(emptyList(), emptyList(), emptyList(), emptyList());
 
         public final List<TableInfo> indices;
         public final List<Enrich> enriches;
+        public final List<InferencePlan> inferencePlans;
         public final List<TableInfo> lookupIndices;
 
-        public PreAnalysis(List<TableInfo> indices, List<Enrich> enriches, 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;
             this.lookupIndices = lookupIndices;
         }
     }
@@ -47,17 +55,19 @@ public PreAnalysis preAnalyze(LogicalPlan plan) {
     protected PreAnalysis doPreAnalyze(LogicalPlan plan) {
         List<TableInfo> indices = new ArrayList<>();
         List<Enrich> unresolvedEnriches = new ArrayList<>();
+        List<InferencePlan> unresolvedInferencePlans = new ArrayList<>();
         List<TableInfo> lookupIndices = new ArrayList<>();
 
         plan.forEachUp(UnresolvedRelation.class, p -> {
             List<TableInfo> list = p.indexMode() == IndexMode.LOOKUP ? lookupIndices : indices;
             list.add(new TableInfo(p.indexPattern()));
         });
         plan.forEachUp(Enrich.class, unresolvedEnriches::add);
+        plan.forEachUp(InferencePlan.class, unresolvedInferencePlans::add);
 
         // mark plan as preAnalyzed (if it were marked, there would be no analysis)
         plan.forEachUp(LogicalPlan::setPreAnalyzed);
 
-        return new PreAnalysis(indices, unresolvedEnriches, lookupIndices);
+        return new PreAnalysis(indices, unresolvedEnriches, unresolvedInferencePlans, lookupIndices);
     }
 }

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

@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
+
+import java.util.Collection;
+import java.util.Map;
+
+public class InferenceResolution {
+
+    public static final InferenceResolution EMPTY = new InferenceResolution();
+
+    private final Map<String, ResolvedInference> resolvedInferences = ConcurrentCollections.newConcurrentMap();
+
+    private final Map<String, String> errors = ConcurrentCollections.newConcurrentMap();
+
+    public ResolvedInference getResolvedInference(String inferenceId) {
+        return resolvedInferences.get(inferenceId);
+    }
+
+    public Collection<ResolvedInference> resolvedInferences() {
+        return resolvedInferences.values();
+    }
+
+    public String getError(String inferenceId) {
+        final String error = errors.get(inferenceId);
+        if (error != null) {
+            return error;
+        } else {
+            assert false : "unresolved inference [" + inferenceId + "]";
+            return "unresolved inference [" + inferenceId + "]";
+        }
+    }
+
+    public void addResolvedInference(ResolvedInference resolvedInference) {
+        resolvedInferences.putIfAbsent(resolvedInference.inferenceId(), resolvedInference);
+    }
+
+    public void addError(String inferenceId, String reason) {
+        errors.putIfAbsent(inferenceId, reason);
+    }
+}

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

@@ -0,0 +1,70 @@
+/*
+ * 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;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.OriginSettingClient;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
+
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.stream.Collectors;
+
+import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
+
+public class InferenceService {
+
+    private final Client client;
+
+    public InferenceService(Client client) {
+        this.client = new OriginSettingClient(client, ML_ORIGIN);
+    }
+
+    public void resolveInferences(List<InferencePlan> plans, ActionListener<InferenceResolution> listener) {
+
+        if (plans.isEmpty()) {
+            listener.onResponse(InferenceResolution.EMPTY);
+            return;
+        }
+
+        Set<String> inferenceIds = plans.stream()
+            .map(p -> p.inferenceId().fold(FoldContext.small()).toString())
+            .collect(Collectors.toSet());
+
+        CountDownLatch countDownLatch = new CountDownLatch(inferenceIds.size());
+        InferenceResolution inferenceResolution = new InferenceResolution();
+
+        for (var inferenceId : inferenceIds) {
+            client.execute(
+                GetInferenceModelAction.INSTANCE,
+                new GetInferenceModelAction.Request(inferenceId, TaskType.ANY),
+                ActionListener.wrap(r -> {
+                    ResolvedInference resolvedInference = new ResolvedInference(inferenceId, r.getEndpoints().getFirst().getTaskType());
+                    inferenceResolution.addResolvedInference(resolvedInference);
+                    countDownLatch.countDown();
+                }, e -> {
+                    inferenceResolution.addError(inferenceId, e.getMessage());
+                    countDownLatch.countDown();
+                })
+            );
+        }
+
+        try {
+            countDownLatch.await();
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+
+        listener.onResponse(inferenceResolution);
+    }
+}

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

@@ -0,0 +1,28 @@
+/*
+ * 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;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.inference.TaskType;
+
+import java.io.IOException;
+
+public record ResolvedInference(String inferenceId, TaskType taskType) implements Writeable {
+
+    public ResolvedInference(StreamInput in) throws IOException {
+        this(in.readString(), in.readEnum(TaskType.class));
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(inferenceId);
+        out.writeEnum(taskType);
+    }
+}

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

@@ -195,7 +195,7 @@ public Literal visitIntegerValue(EsqlBaseParser.IntegerValueContext ctx) {
     }
 
     @Override
-    public String visitStringOrParameter(EsqlBaseParser.StringOrParameterContext ctx) {
+    public Literal visitStringOrParameter(EsqlBaseParser.StringOrParameterContext ctx) {
         if (ctx.parameter() != null) {
             if (expression(ctx.parameter()) instanceof Literal lit) {
                 if (lit.value() == null) {
@@ -205,7 +205,7 @@ public String visitStringOrParameter(EsqlBaseParser.StringOrParameterContext ctx
                         ctx.parameter().getText()
                     );
                 }
-                return lit.value().toString();
+                return lit;
             }
 
             throw new ParsingException(
@@ -215,7 +215,7 @@ public String visitStringOrParameter(EsqlBaseParser.StringOrParameterContext ctx
             );
         }
 
-        return unquote(source(ctx.string()));
+        return visitString(ctx.string());
     }
 
     @Override