Fix an error with InferenceFunctionEvaluator when outputing a unidime… (#136210)

Author: afoucret

muted-tests.yml

@@ -615,9 +615,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
   method: test {csv-spec:semantic_text.Repeat}
   issue: https://github.com/elastic/elasticsearch/issues/136150
-- class: org.elasticsearch.xpack.esql.inference.InferenceFunctionEvaluatorTests
-  method: testFoldTextEmbeddingFunction
-  issue: https://github.com/elastic/elasticsearch/issues/136154
 - class: org.elasticsearch.xpack.remotecluster.CrossClusterEsqlRCS1UnavailableRemotesIT
   method: testEsqlRcs1UnavailableRemoteScenarios
   issue: https://github.com/elastic/elasticsearch/issues/136157

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbedding.java

@@ -109,7 +109,7 @@ public boolean foldable() {
 
     @Override
     public DataType dataType() {
-        return DataType.DENSE_VECTOR;
+        return inputText.dataType() == DataType.NULL ? DataType.NULL : DataType.DENSE_VECTOR;
     }
 
     @Override

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

@@ -24,11 +24,14 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.evaluator.EvalMapper;
 import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
 import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
 import org.elasticsearch.xpack.esql.inference.textembedding.TextEmbeddingOperator;
 
+import java.util.List;
+
 /**
  * Evaluator for inference functions that performs constant folding by executing inference operations
  * at optimization time and replacing them with their computed results.
@@ -76,6 +79,11 @@ public void fold(InferenceFunction<?> f, ActionListener<Expression> listener) {
             listener.onFailure(new IllegalArgumentException("Inference function must be foldable"));
             return;
         }
+        if (f.dataType() == DataType.NULL) {
+            // If the function's return type is NULL, we can directly return a NULL literal without executing anything.
+            listener.onResponse(Literal.of(f, null));
+            return;
+        }
 
         // Set up a DriverContext for executing the inference operator.
         // This follows the same pattern as EvaluatorMapper but in a simplified context
@@ -129,7 +137,7 @@ public CircuitBreakerStats stats(String name) {
                         }
 
                         // Convert the operator result back to an ESQL expression (Literal)
-                        l.onResponse(Literal.of(f, BlockUtils.toJavaObject(output.getBlock(0), 0)));
+                        l.onResponse(Literal.of(f, processValue(f.dataType(), BlockUtils.toJavaObject(output.getBlock(0), 0))));
                     } finally {
                         Releasables.close(inferenceOperator);
                         if (output != null) {
@@ -148,6 +156,14 @@ public CircuitBreakerStats stats(String name) {
         }
     }
 
+    private Object processValue(DataType dataType, Object value) {
+        if (dataType == DataType.DENSE_VECTOR && value instanceof List == false) {
+            value = List.of(value);
+        }
+
+        return value;
+    }
+
     /**
      * Functional interface for providing inference operators.
      * <p>

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

@@ -32,6 +32,7 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.nullValue;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -54,8 +55,8 @@ public void testFoldTextEmbeddingFunction() throws Exception {
         // Create a mock TextEmbedding function
         TextEmbedding textEmbeddingFunction = new TextEmbedding(
             Source.EMPTY,
-            Literal.keyword(Source.EMPTY, "test-model"),
-            Literal.keyword(Source.EMPTY, "test input")
+            Literal.keyword(Source.EMPTY, "test input"),
+            Literal.keyword(Source.EMPTY, "test-model")
         );
 
         // Create a mock operator that returns a result
@@ -94,12 +95,46 @@ public void testFoldTextEmbeddingFunction() throws Exception {
         allBreakersEmpty();
     }
 
+    public void testFoldTextEmbeddingFunctionWithNullInput() throws Exception {
+        // Create a mock TextEmbedding function
+        TextEmbedding textEmbeddingFunction = new TextEmbedding(Source.EMPTY, Literal.NULL, Literal.keyword(Source.EMPTY, "test-model"));
+
+        // Create a mock operator that returns a result
+        Operator operator = mock(Operator.class);
+
+        Float[] embedding = randomArray(1, 100, Float[]::new, ESTestCase::randomFloat);
+
+        when(operator.getOutput()).thenAnswer(i -> {
+            FloatBlock.Builder outputBlockBuilder = blockFactory().newFloatBlockBuilder(1);
+            outputBlockBuilder.appendNull();
+            return new Page(outputBlockBuilder.build());
+        });
+
+        InferenceFunctionEvaluator.InferenceOperatorProvider inferenceOperatorProvider = (f, driverContext) -> operator;
+
+        // Execute the fold operation
+        InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(FoldContext.small(), inferenceOperatorProvider);
+
+        AtomicReference<Expression> resultExpression = new AtomicReference<>();
+        evaluator.fold(textEmbeddingFunction, ActionListener.wrap(resultExpression::set, ESTestCase::fail));
+
+        assertBusy(() -> {
+            assertNotNull(resultExpression.get());
+            Literal result = as(resultExpression.get(), Literal.class);
+            assertThat(result.dataType(), equalTo(DataType.NULL));
+            assertThat(result.value(), nullValue());
+        });
+
+        // Check all breakers are empty after the operation is executed
+        allBreakersEmpty();
+    }
+
     public void testFoldWithNonFoldableFunction() {
         // A function with a non-literal argument is not foldable.
         TextEmbedding textEmbeddingFunction = new TextEmbedding(
             Source.EMPTY,
             mock(Attribute.class),
-            Literal.keyword(Source.EMPTY, "test input")
+            Literal.keyword(Source.EMPTY, "test model")
         );
 
         InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(
@@ -118,8 +153,8 @@ public void testFoldWithNonFoldableFunction() {
     public void testFoldWithAsyncFailure() throws Exception {
         TextEmbedding textEmbeddingFunction = new TextEmbedding(
             Source.EMPTY,
-            Literal.keyword(Source.EMPTY, "test-model"),
-            Literal.keyword(Source.EMPTY, "test input")
+            Literal.keyword(Source.EMPTY, "test input"),
+            Literal.keyword(Source.EMPTY, "test-model")
         );
 
         // Mock an operator that will trigger an async failure
@@ -146,8 +181,8 @@ public void testFoldWithAsyncFailure() throws Exception {
     public void testFoldWithNullOutputPage() throws Exception {
         TextEmbedding textEmbeddingFunction = new TextEmbedding(
             Source.EMPTY,
-            Literal.keyword(Source.EMPTY, "test-model"),
-            Literal.keyword(Source.EMPTY, "test input")
+            Literal.keyword(Source.EMPTY, "test input"),
+            Literal.keyword(Source.EMPTY, "test-model")
         );
 
         Operator operator = mock(Operator.class);