[ES|QL] Fix broken CSV tests for text embedding (#136300)

* Fix bug in CSV tests when several function embedding are present in the same plan.

* Fix GenerativeForkIT text-embedding tests.

* lint

Author: afoucret

muted-tests.yml

@@ -591,9 +591,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
   method: test {csv-spec:math.PowIntInt}
   issue: https://github.com/elastic/elasticsearch/issues/136106
-- class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
-  method: test {csv-spec:text-embedding.Text_embedding with knn on semantic_text_dense_field}
-  issue: https://github.com/elastic/elasticsearch/issues/136108
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
   method: test {csv-spec:string.ContainsFail}
   issue: https://github.com/elastic/elasticsearch/issues/136112
@@ -606,9 +603,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
   method: test {csv-spec:bucket.BucketByWeekInString}
   issue: https://github.com/elastic/elasticsearch/issues/136136
-- class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
-  method: test {csv-spec:text-embedding.Text_embedding with knn (inline) on semantic_text_dense_field}
-  issue: https://github.com/elastic/elasticsearch/issues/136142
 - class: org.elasticsearch.versioning.ConcurrentSeqNoVersioningIT
   method: testSeqNoCASLinearizability
   issue: https://github.com/elastic/elasticsearch/issues/117249
@@ -624,15 +618,9 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
   method: test {csv-spec:math.LeastGreatestMany}
   issue: https://github.com/elastic/elasticsearch/issues/136161
-- class: org.elasticsearch.xpack.esql.qa.*.EsqlSpecIT
-  method: test {csv-spec:text-embedding.*}
-  issue: https://github.com/elastic/elasticsearch/issues/136090
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
   method: test {csv-spec:categorize.Values aggs}
   issue: https://github.com/elastic/elasticsearch/issues/136230
-- class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
-  method: test {csv-spec:text-embedding.Text_embedding using a row source operator}
-  issue: https://github.com/elastic/elasticsearch/issues/136234
 - class: org.elasticsearch.test.rest.yaml.CcsCommonYamlTestSuiteIT
   method: test {p0=field_caps/10_basic/Field caps for number field with only doc values}
   issue: https://github.com/elastic/elasticsearch/issues/136244

x-pack/plugin/esql/qa/testFixtures/src/main/resources/text-embedding.csv-spec

@@ -69,6 +69,25 @@ live long and prosper
 all we have to decide is what to do with the time that is given to us
 ;
 
+text_embedding with multiple knn queries
+required_capability: text_embedding_function
+required_capability: dense_vector_field_type_released
+required_capability: knn_function_v5
+required_capability: fork_v9
+required_capability: semantic_text_field_caps
+
+FROM semantic_text METADATA _score, _id
+| WHERE KNN(semantic_text_dense_field, TEXT_EMBEDDING("be excellent to each other", "test_dense_inference")) OR KNN(semantic_text_dense_field, TEXT_EMBEDDING("live long and prosper", "test_dense_inference")) 
+| SORT _score DESC, _id
+| LIMIT 10
+| KEEP semantic_text_field
+;
+
+semantic_text_field:text                                             
+live long and prosper
+be excellent to each other
+all we have to decide is what to do with the time that is given to us                  
+;
 
 text_embedding with multiple knn queries in fork
 required_capability: text_embedding_function
@@ -91,5 +110,5 @@ live long and prosper                                                 | [50.0, 5
 live long and prosper                                                 | [45.0, 55.0, 54.0]           | fork1        
 be excellent to each other                                            | [50.0, 57.0, 56.0]           | fork2        
 all we have to decide is what to do with the time that is given to us | [45.0, 55.0, 54.0]           | fork1        
-all we have to decide is what to do with the time that is given to us | [50.0, 57.0, 56.0]           | fork2        
+all we have to decide is what to do with the time that is given to us | [50.0, 57.0, 56.0]           | fork2
 ;

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.BlockFactory;
@@ -88,7 +89,7 @@ public void fold(InferenceFunction<?> f, ActionListener<Expression> listener) {
         // Set up a DriverContext for executing the inference operator.
         // This follows the same pattern as EvaluatorMapper but in a simplified context
         // suitable for constant folding during optimization.
-        CircuitBreaker breaker = foldContext.circuitBreakerView(f.source());
+        CircuitBreaker breaker = new NoopCircuitBreaker(CircuitBreaker.REQUEST);
         BigArrays bigArrays = new BigArrays(null, new CircuitBreakerService() {
             @Override
             public CircuitBreaker getBreaker(String name) {
@@ -123,26 +124,28 @@ public CircuitBreakerStats stats(String name) {
                 // Execute the inference operation asynchronously and handle the result
                 // The operator will perform the actual inference call and return a page with the result
                 driverContext.waitForAsyncActions(listener.delegateFailureIgnoreResponseAndWrap(l -> {
-                    Page output = inferenceOperator.getOutput();
-
                     try {
+                        Page output = inferenceOperator.getOutput();
+
                         if (output == null) {
                             l.onFailure(new IllegalStateException("Expected output page from inference operator"));
                             return;
                         }
 
+                        output.allowPassingToDifferentDriver();
+                        l = ActionListener.releaseBefore(output, l);
+
                         if (output.getPositionCount() != 1 || output.getBlockCount() != 1) {
                             l.onFailure(new IllegalStateException("Expected a single block with a single value from inference operator"));
                             return;
                         }
 
                         // Convert the operator result back to an ESQL expression (Literal)
                         l.onResponse(Literal.of(f, processValue(f.dataType(), BlockUtils.toJavaObject(output.getBlock(0), 0))));
+                    } catch (Exception e) {
+                        l.onFailure(e);
                     } finally {
                         Releasables.close(inferenceOperator);
-                        if (output != null) {
-                            output.releaseBlocks();
-                        }
                     }
                 }));
             } catch (Exception e) {

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

@@ -188,6 +188,10 @@ private static String inferenceId(UnresolvedFunction f, FunctionDefinition def)
 
         for (int i = 0; i < functionDescription.args().size(); i++) {
             EsqlFunctionRegistry.ArgSignature arg = functionDescription.args().get(i);
+            if (i >= f.arguments().size()) {
+                // Argument is missing. We will fail later during verifier, so just return null here.
+                return null;
+            }
 
             if (arg.name().equals(InferenceFunction.INFERENCE_ID_PARAMETER_NAME)) {
                 Expression inferenceId = f.arguments().get(i);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceFunctionEvaluatorTests.java

@@ -12,8 +12,11 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
@@ -22,6 +25,7 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
 import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
 import org.junit.After;
 import org.junit.Before;
 
@@ -33,6 +37,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.nullValue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -51,6 +56,7 @@ public void tearDownThreadPool() {
         terminate(threadPool);
     }
 
+    @SuppressWarnings("unchecked")
     public void testFoldTextEmbeddingFunction() throws Exception {
         // Create a mock TextEmbedding function
         TextEmbedding textEmbeddingFunction = new TextEmbedding(
@@ -62,7 +68,23 @@ public void testFoldTextEmbeddingFunction() throws Exception {
         // Create a mock operator that returns a result
         Operator operator = mock(Operator.class);
 
-        Float[] embedding = randomArray(1, 100, Float[]::new, ESTestCase::randomFloat);
+        float[] embedding = randomEmbedding(between(1, 100));
+
+        InferenceService inferenceService = mock(InferenceService.class);
+        BulkInferenceRunner bulkInferenceRunner = mock(BulkInferenceRunner.class);
+
+        doAnswer(i -> {
+            threadPool.schedule(
+                () -> i.getArgument(1, ActionListener.class).onResponse(List.of(inferenceResponse(embedding))),
+                TimeValue.timeValueMillis(between(1, 10)),
+                threadPool.generic()
+            );
+
+            return null;
+        }).when(bulkInferenceRunner).executeBulk(any(), any());
+        when(bulkInferenceRunner.threadPool()).thenReturn(threadPool);
+
+        when(inferenceService.bulkInferenceRunner()).thenReturn(bulkInferenceRunner);
 
         when(operator.getOutput()).thenAnswer(i -> {
             FloatBlock.Builder outputBlockBuilder = blockFactory().newFloatBlockBuilder(1).beginPositionEntry();
@@ -79,7 +101,7 @@ public void testFoldTextEmbeddingFunction() throws Exception {
         InferenceFunctionEvaluator.InferenceOperatorProvider inferenceOperatorProvider = (f, driverContext) -> operator;
 
         // Execute the fold operation
-        InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(FoldContext.small(), inferenceOperatorProvider);
+        InferenceFunctionEvaluator evaluator = InferenceFunctionEvaluator.factory().create(FoldContext.small(), inferenceService);
 
         AtomicReference<Expression> resultExpression = new AtomicReference<>();
         evaluator.fold(textEmbeddingFunction, ActionListener.wrap(resultExpression::set, ESTestCase::fail));
@@ -218,4 +240,17 @@ public void testFoldWithUnsupportedFunction() throws Exception {
 
         allBreakersEmpty();
     }
+
+    private float[] randomEmbedding(int length) {
+        float[] embedding = new float[length];
+        for (int i = 0; i < length; i++) {
+            embedding[i] = randomFloat();
+        }
+        return embedding;
+    }
+
+    private InferenceAction.Response inferenceResponse(float[] embedding) {
+        TextEmbeddingFloatResults.Embedding embeddingResult = new TextEmbeddingFloatResults.Embedding(embedding);
+        return new InferenceAction.Response(new TextEmbeddingFloatResults(List.of(embeddingResult)));
+    }
 }