Draft CompletionOperator.

Author: afoucret

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

@@ -7,11 +7,22 @@
 
 package org.elasticsearch.xpack.esql.inference;
 
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.CountDownActionListener;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.ChatCompletionResults;
+
+import java.util.List;
 
 public class CompletionOperator extends InferenceOperator<Page> {
 
@@ -30,6 +41,7 @@ public Operator get(DriverContext driverContext) {
     }
 
     private final ExpressionEvaluator promptEvaluator;
+    private final BlockFactory blockFactory;
 
     public CompletionOperator(
         DriverContext driverContext,
@@ -39,12 +51,66 @@ public CompletionOperator(
     ) {
         super(driverContext, inferenceRunner.getThreadContext(), inferenceRunner, inferenceId);
         this.promptEvaluator = promptEvaluator;
+        this.blockFactory = driverContext.blockFactory();
     }
 
     @Override
     protected void performAsync(Page inputPage, ActionListener<Page> listener) {
-        Page outputPage = inputPage.appendBlock(promptEvaluator.eval(inputPage));
-        listener.onResponse(outputPage);
+        int pageSize = inputPage.getPositionCount();
+        String[] responses = new String[pageSize];
+
+        CountDownActionListener countDownListener = new CountDownActionListener(
+            inputPage.getPositionCount(),
+            listener.delegateFailureIgnoreResponseAndWrap(l -> {
+                try(BytesRefBlock.Builder outputBlockBuilder = blockFactory.newBytesRefBlockBuilder(pageSize)) {
+                    BytesRefBuilder bytesRefBuilder = new BytesRefBuilder();
+                    for (int pos = 0; pos < pageSize; pos++) {
+                        if (responses[pos] == null) {
+                            outputBlockBuilder.appendNull();
+                        } else {
+                            bytesRefBuilder.copyChars(responses[pos]);
+                            outputBlockBuilder.appendBytesRef(bytesRefBuilder.get());
+                        }
+                    }
+
+                    l.onResponse(inputPage.appendBlock(outputBlockBuilder.build()));
+                }
+            })
+        );
+
+        try (BytesRefBlock promptBlock = (BytesRefBlock) promptEvaluator.eval(inputPage)) {
+            BytesRef readBuffer = new BytesRef();
+            for (int pos = 0; pos < pageSize; pos++) {
+                final int currentPos = pos;
+                if (promptBlock.isNull(pos)) {
+                    countDownListener.onResponse(null);
+                } else {
+                    StringBuilder promptBuilder = new StringBuilder();
+                    for (int valueIndex = 0; valueIndex < promptBlock.getValueCount(pos); valueIndex++) {
+                        readBuffer = promptBlock.getBytesRef(promptBlock.getFirstValueIndex(pos) + valueIndex, readBuffer);
+                        promptBuilder.append(readBuffer.utf8ToString()).append("\n");
+
+
+                    InferenceAction.Request request = InferenceAction.Request.builder(inferenceId(), TaskType.COMPLETION)
+                        .setInput(List.of(promptBuilder.toString())).build();
+
+                    doInference(request, countDownListener.delegateFailureAndWrap((l, r) -> {
+                        if (r.getResults() instanceof ChatCompletionResults completionResults) {
+                            responses[currentPos] = completionResults.results().getFirst().content();
+                            l.onResponse(null);
+                        } else {
+                            l.onFailure(new IllegalStateException(
+                                "Inference result has wrong type. Got ["
+                                    + r.getResults().getClass()
+                                    + "] while expecting ["
+                                    + ChatCompletionResults.class
+                                    + "]"
+                            ));
+                        }
+                    }));
+                }
+            }
+        }
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/Completion.java

@@ -136,7 +136,7 @@ protected AttributeSet computeReferences() {
 
     @Override
     public boolean expressionsResolved() {
-        return super.expressionsResolved() && prompt.resolved();
+        return super.expressionsResolved() && prompt.resolved() && targetField.resolved();
     }
 
     @Override