InferenceResolution for text embedding function.

afoucret · afoucret · commit 847c99831638 · 2025-09-11T19:41:26.000+02:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java
@@ -56,6 +56,7 @@
 import org.elasticsearch.xpack.esql.expression.function.grouping.Bucket;
 import org.elasticsearch.xpack.esql.expression.function.grouping.Categorize;
 import org.elasticsearch.xpack.esql.expression.function.grouping.TBucket;
+import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Greatest;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Least;
@@ -193,7 +194,6 @@
 import org.elasticsearch.xpack.esql.expression.function.vector.L1Norm;
 import org.elasticsearch.xpack.esql.expression.function.vector.L2Norm;
 import org.elasticsearch.xpack.esql.expression.function.vector.Magnitude;
-import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 import org.elasticsearch.xpack.esql.session.Configuration;
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceResolver.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceResolver.java
@@ -15,6 +15,11 @@
 import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
+import org.elasticsearch.xpack.esql.expression.function.FunctionDefinition;
+import org.elasticsearch.xpack.esql.expression.function.UnresolvedFunction;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
 
@@ -29,14 +34,16 @@
 public class InferenceResolver {
 
     private final Client client;
+    private final EsqlFunctionRegistry functionRegistry;
 
     /**
      * Constructs a new {@code InferenceResolver}.
      *
      * @param client The Elasticsearch client for executing inference deployment lookups
      */
-    public InferenceResolver(Client client) {
+    public InferenceResolver(Client client, EsqlFunctionRegistry functionRegistry) {
         this.client = client;
+        this.functionRegistry = functionRegistry;
     }
 
     /**
@@ -71,6 +78,7 @@ public void resolveInferenceIds(LogicalPlan plan, ActionListener<InferenceResolu
      */
     void collectInferenceIds(LogicalPlan plan, Consumer<String> c) {
         collectInferenceIdsFromInferencePlans(plan, c);
+        collectInferenceIdsFromInferenceFunctions(plan, c);
     }
 
     /**
@@ -130,6 +138,28 @@ private void collectInferenceIdsFromInferencePlans(LogicalPlan plan, Consumer<St
         plan.forEachUp(InferencePlan.class, inferencePlan -> c.accept(inferenceId(inferencePlan)));
     }
 
+    /**
+     * Collects inference IDs from function expressions within the logical plan.
+     *
+     * @param plan The logical plan to scan for function expressions
+     * @param c    Consumer function to receive each discovered inference ID
+     */
+    private void collectInferenceIdsFromInferenceFunctions(LogicalPlan plan, Consumer<String> c) {
+        EsqlFunctionRegistry snapshotRegistry = functionRegistry.snapshotRegistry();
+        plan.forEachExpressionUp(UnresolvedFunction.class, f -> {
+            String functionName = snapshotRegistry.resolveAlias(f.name());
+            if (snapshotRegistry.functionExists(functionName)) {
+                FunctionDefinition def = snapshotRegistry.resolveFunction(functionName);
+                if (InferenceFunction.class.isAssignableFrom(def.clazz())) {
+                    String inferenceId = inferenceId(f, def);
+                    if (inferenceId != null) {
+                        c.accept(inferenceId);
+                    }
+                }
+            }
+        });
+    }
+
     /**
      * Extracts the inference ID from an InferencePlan object.
      *
@@ -144,6 +174,23 @@ private static String inferenceId(Expression e) {
         return BytesRefs.toString(e.fold(FoldContext.small()));
     }
 
+    public String inferenceId(UnresolvedFunction f, FunctionDefinition def) {
+        EsqlFunctionRegistry.FunctionDescription functionDescription = EsqlFunctionRegistry.description(def);
+
+        for (int i = 0; i < functionDescription.args().size(); i++) {
+            EsqlFunctionRegistry.ArgSignature arg = functionDescription.args().get(i);
+
+            if (arg.name().equals(InferenceFunction.INFERENCE_ID_PARAMETER_NAME)) {
+                Expression argValue = f.arguments().get(i);
+                if (argValue != null && argValue.foldable() && DataType.isString(argValue.dataType())) {
+                    return inferenceId(argValue);
+                }
+            }
+        }
+
+        return null;
+    }
+
     public static Factory factory(Client client) {
         return new Factory(client);
     }
@@ -155,8 +202,8 @@ private Factory(Client client) {
             this.client = client;
         }
 
-        public InferenceResolver create() {
-            return new InferenceResolver(client);
+        public InferenceResolver create(EsqlFunctionRegistry functionRegistry) {
+            return new InferenceResolver(client, functionRegistry);
         }
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceService.java
@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.inference;
 
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunnerConfig;
 
@@ -33,10 +34,12 @@ private InferenceService(InferenceResolver.Factory inferenceResolverFactory, Bul
     /**
      * Creates an inference resolver for resolving inference IDs in logical plans.
      *
+     * @param functionRegistry the function registry to resolve functions
+     *
      * @return a new inference resolver instance
      */
-    public InferenceResolver inferenceResolver() {
-        return inferenceResolverFactory.create();
+    public InferenceResolver inferenceResolver(EsqlFunctionRegistry functionRegistry) {
+        return inferenceResolverFactory.create(functionRegistry);
     }
 
     public BulkInferenceRunner bulkInferenceRunner() {
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
@@ -731,7 +731,7 @@ private void analyzeWithRetry(
     }
 
     private void resolveInferences(LogicalPlan plan, PreAnalysisResult preAnalysisResult, ActionListener<PreAnalysisResult> l) {
-        inferenceService.inferenceResolver().resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));
+        inferenceService.inferenceResolver(functionRegistry).resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));
     }
 
     private PhysicalPlan logicalPlanToPhysicalPlan(LogicalPlan optimizedPlan, EsqlQueryRequest request) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceResolverTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceResolverTests.java
@@ -22,6 +22,7 @@
 import org.elasticsearch.threadpool.FixedExecutorBuilder;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
+import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.parser.EsqlParser;
 import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 import org.junit.After;
@@ -44,6 +45,7 @@
 
 public class InferenceResolverTests extends ESTestCase {
     private TestThreadPool threadPool;
+    private EsqlFunctionRegistry functionRegistry;
 
     @Before
     public void setThreadPool() {
@@ -60,6 +62,11 @@ public void setThreadPool() {
         );
     }
 
+    @Before
+    public void setUpFunctionRegistry() {
+        functionRegistry = new EsqlFunctionRegistry();
+    }
+
     @After
     public void shutdownThreadPool() {
         terminate(threadPool);
@@ -78,6 +85,18 @@ public void testCollectInferenceIds() {
             List.of("completion-inference-id")
         );
 
+        // Test inference ID collection from an inference function
+        assertCollectInferenceIds(
+            "FROM books METADATA _score | EVAL embedding = TEXT_EMBEDDING(\"description\", \"text-embedding-inference-id\")",
+            List.of("text-embedding-inference-id")
+        );
+
+        // Test inference ID collection with nested functions
+        assertCollectInferenceIds(
+            "FROM books METADATA _score | EVAL embedding = TEXT_EMBEDDING(TEXT_EMBEDDING(\"nested\", \"nested-id\"), \"outer-id\")",
+            List.of("nested-id", "outer-id")
+        );
+
         // Multiple inference plans
         assertCollectInferenceIds("""
             FROM books METADATA _score
@@ -139,7 +158,7 @@ public void testResolveMultipleInferenceIds() throws Exception {
 
     public void testResolveMissingInferenceIds() throws Exception {
         InferenceResolver inferenceResolver = inferenceResolver();
-        List<String> inferenceIds = List.of("missing-plan");
+        List<String> inferenceIds = List.of("missing-inference-id");
 
         SetOnce<InferenceResolution> inferenceResolutionSetOnce = new SetOnce<>();
 
@@ -153,7 +172,7 @@ public void testResolveMissingInferenceIds() throws Exception {
 
             assertThat(inferenceResolution.resolvedInferences(), empty());
             assertThat(inferenceResolution.hasError(), equalTo(true));
-            assertThat(inferenceResolution.getError("missing-plan"), equalTo("inference endpoint not found"));
+            assertThat(inferenceResolution.getError("missing-inference-id"), equalTo("inference endpoint not found"));
         });
     }
 
@@ -205,7 +224,7 @@ private static ActionResponse getInferenceModelResponse(GetInferenceModelAction.
     }
 
     private InferenceResolver inferenceResolver() {
-        return new InferenceResolver(mockClient());
+        return new InferenceResolver(mockClient(), functionRegistry);
     }
 
     private static ModelConfigurations mockModelConfig(String inferenceId, TaskType taskType) {

Original file line number	Diff line number	Diff line change
`@@ -731,7 +731,7 @@ private void analyzeWithRetry(`
`731`	`731`	`}`
`732`	`732`
`733`	`733`	`private void resolveInferences(LogicalPlan plan, PreAnalysisResult preAnalysisResult, ActionListener<PreAnalysisResult> l) {`
`734`		`- inferenceService.inferenceResolver().resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));`
	`734`	`+ inferenceService.inferenceResolver(functionRegistry).resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));`
`735`	`735`	`}`
`736`	`736`
`737`	`737`	`private PhysicalPlan logicalPlanToPhysicalPlan(LogicalPlan optimizedPlan, EsqlQueryRequest request) {`