Skip to content

Commit 42e14c1

Browse files
committed
Finished refactoring the inference operator implementation.
1 parent 7df8348 commit 42e14c1

18 files changed

+541
-491
lines changed

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

Lines changed: 0 additions & 179 deletions
This file was deleted.

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

Lines changed: 0 additions & 75 deletions
This file was deleted.

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,13 @@ public abstract class InferenceOperator<InferenceResult extends InferenceService
2626
private final String inferenceId;
2727
private final BlockFactory blockFactory;
2828

29-
private final BulkInferenceExecutor<InferenceResult, Page> bulkInferenceExecutor;
29+
private final BulkInferenceExecutor bulkInferenceExecutor;
3030

3131
@SuppressWarnings("this-escape")
3232
public InferenceOperator(DriverContext driverContext, InferenceRunner inferenceRunner, ThreadPool threadPool, String inferenceId) {
3333
super(driverContext, threadPool.getThreadContext(), MAX_INFERENCE_WORKER);
3434
this.blockFactory = driverContext.blockFactory();
35-
this.bulkInferenceExecutor = new BulkInferenceExecutor<>(inferenceRunner, threadPool, bulkExecutionConfig());
35+
this.bulkInferenceExecutor = new BulkInferenceExecutor(inferenceRunner, threadPool, bulkExecutionConfig());
3636
this.inferenceId = inferenceId;
3737
}
3838

@@ -75,5 +75,13 @@ protected BulkInferenceExecutionConfig bulkExecutionConfig() {
7575

7676
protected abstract BulkInferenceRequestIterator requests(Page input);
7777

78-
protected abstract BulkInferenceOutputBuilder<InferenceResult, Page> outputBuilder(Page input);
78+
protected abstract OutputBuilder<InferenceResult> outputBuilder(Page input);
79+
80+
public abstract static class OutputBuilder<InferenceResult extends InferenceServiceResults> extends BulkInferenceOutputBuilder<
81+
InferenceResult,
82+
Page> {
83+
protected void releasePageOnAnyThread(Page page) {
84+
InferenceOperator.releasePageOnAnyThread(page);
85+
}
86+
}
7987
}

0 commit comments

Comments
 (0)