Inference execution.

Author: afoucret

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java

@@ -60,6 +60,10 @@ public static class Request extends BaseInferenceActionRequest {
         public static final ParseField QUERY = new ParseField("query");
         public static final ParseField TIMEOUT = new ParseField("timeout");
 
+        public static Builder builder(String inferenceEntityId, TaskType taskType) {
+            return new Builder().setInferenceEntityId(inferenceEntityId).setTaskType(taskType);
+        }
+
         static final ObjectParser<Request.Builder, Void> PARSER = new ObjectParser<>(NAME, Request.Builder::new);
         static {
             PARSER.declareStringArray(Request.Builder::setInput, INPUT);

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
 
@@ -72,4 +73,8 @@ public void resolveInferences(List<InferencePlan> plans, ActionListener<Inferenc
 
         listener.onResponse(inferenceResolution);
     }
+
+    public void doInference(InferenceAction.Request request, ActionListener<InferenceAction.Response> listener) {
+        client.execute(InferenceAction.INSTANCE, request, listener);
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/RerankOperator.java

@@ -7,18 +7,31 @@
 
 package org.elasticsearch.xpack.esql.plan.physical.inference;
 
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AsyncOperator;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
 import org.elasticsearch.xpack.esql.inference.InferenceService;
 
+import java.io.IOException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 public class RerankOperator extends AsyncOperator<Page> {
@@ -101,7 +114,17 @@ public RerankOperator(
 
     @Override
     protected void performAsync(Page inputPage, ActionListener<Page> listener) {
-        listener.onResponse(inputPage);
+        try {
+            inferenceService.doInference(
+                buildInferenceRequest(inputPage),
+                ActionListener.wrap(
+                    (inferenceResponse) -> listener.onResponse(buildOutput(inputPage, inferenceResponse)),
+                    listener::onFailure
+                )
+            );
+        } catch (IOException e) {
+            listener.onFailure(e);
+        }
     }
 
     @Override
@@ -118,4 +141,78 @@ protected void releaseFetchedOnAnyThread(Page page) {
     public Page getOutput() {
         return fetchFromBuffer();
     }
+
+    private Page buildOutput(Page inputPage, InferenceAction.Response inferenceResponse) {
+        int blockCount = inputPage.getBlockCount();
+        Block.Builder[] blocksBuilders = new Block.Builder[blockCount];
+
+        for (int b = 0; b < blockCount; b++) {
+            if (b == scoreChannel) {
+                blocksBuilders[b] = ElementType.DOUBLE.newBlockBuilder(inputPage.getPositionCount(), blockFactory);
+            } else {
+                blocksBuilders[b] = inputPage.getBlock(b).elementType().newBlockBuilder(inputPage.getPositionCount(), blockFactory);
+            }
+        }
+
+        if (inferenceResponse.getResults() instanceof RankedDocsResults rankedDocsResults) {
+            for (var rankedDoc : rankedDocsResults.getRankedDocs()) {
+                for (int b = 0; b < blockCount; b++) {
+                    if (b == scoreChannel) {
+                        if (blocksBuilders[b] instanceof DoubleBlock.Builder scoreBlockBuilder) {
+                            scoreBlockBuilder.beginPositionEntry().appendDouble(rankedDoc.relevanceScore()).endPositionEntry();
+                        }
+                    } else {
+                        blocksBuilders[b].copyFrom(inputPage.getBlock(b), rankedDoc.index(), rankedDoc.index() + 1);
+                    }
+                }
+            }
+
+            return new Page(Block.Builder.buildAll(blocksBuilders));
+        }
+
+        throw new IllegalStateException(
+            "Inference result has wrong type. Got ["
+                + inferenceResponse.getResults().getClass()
+                + "] while expecting ["
+                + RankedDocsResults.class
+                + "]"
+        );
+    }
+
+    private InferenceAction.Request buildInferenceRequest(Page inputPage) throws IOException {
+        String[] inputs = new String[inputPage.getPositionCount()];
+        Map<String, Block> inputBlocks = new HashMap<>();
+
+        for (var entry : rerankFieldsEvaluator.entrySet()) {
+            inputBlocks.put(entry.getKey(), entry.getValue().eval(inputPage));
+        }
+
+        for (int pos = 0; pos < inputPage.getPositionCount(); pos++) {
+            inputs[pos] = toYaml(inputBlocks, pos);
+        }
+
+        return InferenceAction.Request.builder(inferenceId, TaskType.RERANK).setInput(List.of(inputs)).setQuery(queryText).build();
+    }
+
+    private String toYaml(Map<String, Block> inputBlocks, int position) throws IOException {
+        try (XContentBuilder yamlBuilder = XContentFactory.yamlBuilder().startObject()) {
+            for (var blockEntry : inputBlocks.entrySet()) {
+                String fieldName = blockEntry.getKey();
+                Block currentBlock = blockEntry.getValue();
+                if (currentBlock.isNull(position)) {
+                    continue;
+                }
+                yamlBuilder.field(fieldName, toYaml(BlockUtils.toJavaObject(currentBlock, position)));
+            }
+            return Strings.toString(yamlBuilder.endObject());
+        }
+    }
+
+    private Object toYaml(Object value) {
+        return switch (value) {
+            case BytesRef b -> b.utf8ToString();
+            case List<?> l -> l.stream().map(this::toYaml).toList();
+            default -> value;
+        };
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -537,8 +537,6 @@ private PhysicalOperation planRerank(RerankExec rerank, LocalExecutionPlannerCon
             }
         }
 
-        logger.warn("layout {}", source.layout);
-
         return source.with(
             new RerankOperator.Factory(inferenceService, inferenceId, queryText, rerankFieldsEvaluatorSuppliers, scoreChannel),
             source.layout