Refactored inference operator.

Author: afoucret

muted-tests.yml

@@ -315,7 +315,7 @@ tests:
 - class: org.elasticsearch.search.basic.SearchWithRandomDisconnectsIT
   method: testSearchWithRandomDisconnects
   issue: https://github.com/elastic/elasticsearch/issues/122707
-- class: org.elasticsearch.xpack.esql.inference.RerankOperatorTests
+- class: org.elasticsearch.xpack.esql.inference.rerank.RerankOperatorTests
   method: testSimpleCircuitBreaking
   issue: https://github.com/elastic/elasticsearch/issues/124337
 - class: org.elasticsearch.index.engine.ThreadPoolMergeSchedulerTests

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

@@ -8,55 +8,55 @@
 package org.elasticsearch.xpack.esql.inference;
 
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AsyncOperator;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceOperation;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceOutputBuilder;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
 
-import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
-
-abstract public class InferenceOperator<Fetched, InferenceResult extends InferenceServiceResults> extends AsyncOperator<Fetched> {
+public abstract class InferenceOperator<InferenceResult extends InferenceServiceResults> extends AsyncOperator<Page> {
 
     // Move to a setting.
     private static final int MAX_INFERENCE_WORKER = 10;
     private final InferenceRunner inferenceRunner;
     private final String inferenceId;
-    private final Class<InferenceResult> inferenceResultClass;
-
-    public InferenceOperator(
-        DriverContext driverContext,
-        ThreadContext threadContext,
-        InferenceRunner inferenceRunner,
-        String inferenceId,
-        Class<InferenceResult> inferenceResultClass
-    ) {
-        super(driverContext, threadContext, MAX_INFERENCE_WORKER);
+
+    public InferenceOperator(DriverContext driverContext, InferenceRunner inferenceRunner, String inferenceId) {
+        super(driverContext, inferenceRunner.threadContext(), MAX_INFERENCE_WORKER);
         this.inferenceRunner = inferenceRunner;
         this.inferenceId = inferenceId;
-        this.inferenceResultClass = inferenceResultClass;
+    }
 
-        assert inferenceRunner.getThreadContext() != null;
+    protected String inferenceId() {
+        return inferenceId;
     }
 
-    protected final void doInference(InferenceAction.Request inferenceRequest, ActionListener<InferenceResult> listener) {
-        inferenceRunner.doInference(inferenceRequest, listener.map(this::checkedInferenceResults));
+    @Override
+    protected void releaseFetchedOnAnyThread(Page page) {
+        releasePageOnAnyThread(page);
     }
 
-    protected String inferenceId() {
-        return inferenceId;
+    @Override
+    public Page getOutput() {
+        return fetchFromBuffer();
+    }
+
+    @Override
+    protected void performAsync(Page input, ActionListener<Page> listener) {
+        new BulkInferenceOperation<>(bulkInferenceRequestIterator(input), bulkOutputBuilder(input)).execute(inferenceRunner, listener);
     }
 
-    private InferenceResult checkedInferenceResults(InferenceAction.Response inferenceResponse) {
-        if (inferenceResultClass.isInstance(inferenceResponse.getResults())) {
-            return inferenceResultClass.cast(inferenceResponse.getResults());
-        }
-        throw new IllegalStateException(
-            format(
-                "Inference result has wrong type. Got [{}] while expecting [{}]",
-                inferenceResponse.getResults().getClass().getName(),
-                inferenceResultClass.getName()
-            )
-        );
+    protected InferenceAction.Request.Builder inferenceRequestBuilder() {
+        return InferenceAction.Request.builder(inferenceId, taskType());
     }
+
+    protected abstract TaskType taskType();
+
+    protected abstract BulkInferenceRequestIterator bulkInferenceRequestIterator(Page input);
+
+    protected abstract BulkInferenceOutputBuilder<InferenceResult, Page> bulkOutputBuilder(Page input);
 }

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

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.inference;
+
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+
+import java.util.function.Supplier;
+
+public interface InferenceRequestBuilderSupplier extends Supplier<InferenceAction.Request.Builder> {}

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

@@ -0,0 +1,13 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.inference;
+
+import org.elasticsearch.common.CheckedSupplier;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+
+public interface InferenceRequestSupplier extends CheckedSupplier<InferenceAction.Request, Exception> {}

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

@@ -12,24 +12,32 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.threadpool.ThreadPool;
 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;
 
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.ExecutorService;
 import java.util.stream.Collectors;
 
 public class InferenceRunner {
 
     private final Client client;
 
+    private final ExecutorService executorService;
+    private final ThreadContext threadContext;
+
     public InferenceRunner(Client client) {
         this.client = client;
+        // TODO: revisit the executor service instantiation and thread pool choice.
+        this.executorService = client.threadPool().executor(ThreadPool.Names.SEARCH);
+        this.threadContext = client.threadPool().getThreadContext();
     }
 
-    public ThreadContext getThreadContext() {
+    public ThreadContext threadContext() {
         return client.threadPool().getThreadContext();
     }
 
@@ -72,7 +80,20 @@ private static String planInferenceId(InferencePlan<?> plan) {
         return plan.inferenceId().fold(FoldContext.small()).toString();
     }
 
-    public void doInference(InferenceAction.Request request, ActionListener<InferenceAction.Response> listener) {
-        client.execute(InferenceAction.INSTANCE, request, listener);
+    public void doInference(InferenceRequestSupplier request, ActionListener<InferenceAction.Response> listener) {
+        try {
+            if (request == null) {
+                listener.onResponse(null);
+            }
+            executorService.submit(() -> {
+                try {
+                    client.execute(InferenceAction.INSTANCE, request.get(), listener);
+                } catch (Exception e) {
+                    listener.onFailure(e);
+                }
+            });
+        } catch (Exception e) {
+            listener.onFailure(e);
+        }
     }
 }