Skip to content

Commit dbaa437

Browse files
afoucretafoucret
committed
Refactoring.
1 parent 8188b15 commit dbaa437

File tree

14 files changed

+70
-76
lines changed

14 files changed

+70
-76
lines changed

‎x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
6262
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
6363
import org.elasticsearch.xpack.esql.index.EsIndex;
64+
import org.elasticsearch.xpack.esql.inference.InferenceService;
6465
import org.elasticsearch.xpack.esql.parser.QueryParam;
6566
import org.elasticsearch.xpack.esql.plan.logical.Enrich;
6667
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
@@ -136,6 +137,7 @@
136137
import static org.hamcrest.Matchers.instanceOf;
137138
import static org.junit.Assert.assertNotNull;
138139
import static org.junit.Assert.assertNull;
140+
import static org.mockito.Mockito.mock;
139141

140142
public final class EsqlTestUtils {
141143

@@ -353,6 +355,7 @@ public String toString() {
353355
public static final Verifier TEST_VERIFIER = new Verifier(new Metrics(new EsqlFunctionRegistry()), new XPackLicenseState(() -> 0L));
354356

355357
public static final QueryBuilderResolver MOCK_QUERY_BUILDER_RESOLVER = new MockQueryBuilderResolver();
358+
public static final InferenceService MOCK_INFERENCE_SERVICE = mock(InferenceService.class);
356359

357360
private EsqlTestUtils() {}
358361

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ public InferenceService(Client client) {
3030
}
3131

3232
public void resolveInference(InferencePlan inferencePlan, ActionListener<InferenceResolution> listener) {
33-
String inferenceId = inferencePlan.inferenceId().fold().toString();
33+
String inferenceId = inferencePlan.inferenceId();
3434
TaskType taskType = inferencePlan.taskType();
3535

3636
client.execute(

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -581,10 +581,6 @@ public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
581581
throw new ParsingException(source, "RERANK is in preview and only available in SNAPSHOT build");
582582
}
583583

584-
Literal queryText = visitStringOrParameter(ctx.queryText);
585-
Expression input = expression(ctx.input);
586-
Literal inferenceId = new Literal(source(ctx.inferenceId), visitStringOrParameter(ctx.inferenceId), DataType.KEYWORD);
587-
588-
return p -> new Rerank(source, p, queryText, input, inferenceId);
584+
return p -> new Rerank(source, p, visitStringOrParameter(ctx.inferenceId).fold().toString(), visitStringOrParameter(ctx.queryText).fold().toString(), expression(ctx.input));
589585
}
590586
}

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

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,32 @@
77

88
package org.elasticsearch.xpack.esql.plan.logical.inference;
99

10+
import org.elasticsearch.common.io.stream.StreamOutput;
1011
import org.elasticsearch.inference.TaskType;
11-
import org.elasticsearch.xpack.esql.core.expression.Expression;
1212
import org.elasticsearch.xpack.esql.core.tree.Source;
1313
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
1414
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
1515

16+
import java.io.IOException;
1617
import java.util.Objects;
1718

1819
public abstract class InferencePlan extends UnaryPlan {
1920

20-
private final Expression inferenceId;
21+
private final String inferenceId;
2122

22-
protected InferencePlan(Source source, LogicalPlan child, Expression inferenceId) {
23+
protected InferencePlan(Source source, LogicalPlan child, String inferenceId) {
2324
super(source, child);
2425
this.inferenceId = inferenceId;
2526
}
2627

27-
public Expression inferenceId() {
28+
@Override
29+
public void writeTo(StreamOutput out) throws IOException {
30+
Source.EMPTY.writeTo(out);
31+
out.writeNamedWriteable(child());
32+
out.writeString(inferenceId());
33+
}
34+
35+
public String inferenceId() {
2836
return inferenceId;
2937
}
3038

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

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,10 @@
2424
public class Rerank extends InferencePlan {
2525

2626
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Rerank", Rerank::new);
27-
private final Expression queryText;
27+
private final String queryText;
2828
private final Expression input;
2929

30-
public Rerank(Source source, LogicalPlan child, Expression queryText, Expression input, Expression inferenceId) {
30+
public Rerank(Source source, LogicalPlan child, String inferenceId, String queryText, Expression input) {
3131
super(source, child, inferenceId);
3232
this.queryText = queryText;
3333
this.input = input;
@@ -37,22 +37,20 @@ public Rerank(StreamInput in) throws IOException {
3737
this(
3838
Source.readFrom((PlanStreamInput) in),
3939
in.readNamedWriteable(LogicalPlan.class),
40-
in.readNamedWriteable(Expression.class),
41-
in.readNamedWriteable(Expression.class),
40+
in.readString(),
41+
in.readString(),
4242
in.readNamedWriteable(Expression.class)
4343
);
4444
}
4545

4646
@Override
4747
public void writeTo(StreamOutput out) throws IOException {
48-
Source.EMPTY.writeTo(out);
49-
out.writeNamedWriteable(child());
50-
out.writeNamedWriteable(queryText);
48+
super.writeTo(out);
49+
out.writeString(queryText);
5150
out.writeNamedWriteable(input);
52-
out.writeNamedWriteable(inferenceId());
5351
}
5452

55-
public Expression queryText() {
53+
public String queryText() {
5654
return queryText;
5755
}
5856

@@ -77,17 +75,17 @@ public String commandName() {
7775

7876
@Override
7977
public UnaryPlan replaceChild(LogicalPlan newChild) {
80-
return new Rerank(source(), newChild, queryText, input, inferenceId());
78+
return new Rerank(source(), newChild, inferenceId(), queryText, input);
8179
}
8280

8381
@Override
8482
public boolean expressionsResolved() {
85-
return queryText.resolved() && input.resolved() && inferenceId().resolved();
83+
return input.resolved();
8684
}
8785

8886
@Override
8987
protected NodeInfo<? extends LogicalPlan> info() {
90-
return NodeInfo.create(this, Rerank::new, child(), queryText, input, inferenceId());
88+
return NodeInfo.create(this, Rerank::new, child(), inferenceId(), queryText, input);
9189
}
9290

9391
@Override

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

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,17 @@ public class RerankExec extends UnaryExec {
2828
RerankExec::new
2929
);
3030

31-
private final Expression queryText;
32-
31+
private final String inferenceId;
32+
private final String queryText;
3333
private final Expression input;
34-
private final Expression inferenceId;
3534

36-
public RerankExec(Source source, PhysicalPlan child, Expression queryText, Expression input, Expression inferenceId) {
35+
public RerankExec(
36+
Source source,
37+
PhysicalPlan child,
38+
String inferenceId,
39+
String queryText,
40+
Expression input
41+
) {
3742
super(source, child);
3843
this.queryText = queryText;
3944
this.input = input;
@@ -44,24 +49,24 @@ public RerankExec(StreamInput in) throws IOException {
4449
this(
4550
Source.readFrom((PlanStreamInput) in),
4651
in.readNamedWriteable(PhysicalPlan.class),
47-
in.readNamedWriteable(Expression.class),
48-
in.readNamedWriteable(Expression.class),
52+
in.readString(),
53+
in.readString(),
4954
in.readNamedWriteable(Expression.class)
5055
);
5156
}
5257

53-
public Expression queryText() {
58+
public String inferenceId() {
59+
return inferenceId;
60+
}
61+
62+
public String queryText() {
5463
return queryText;
5564
}
5665

5766
public Expression input() {
5867
return input;
5968
}
6069

61-
public Expression inferenceId() {
62-
return inferenceId;
63-
}
64-
6570
@Override
6671
public String getWriteableName() {
6772
return ENTRY.name;
@@ -71,19 +76,19 @@ public String getWriteableName() {
7176
public void writeTo(StreamOutput out) throws IOException {
7277
Source.EMPTY.writeTo(out);
7378
out.writeNamedWriteable(child());
74-
out.writeNamedWriteable(queryText());
79+
out.writeString(inferenceId());
80+
out.writeString(queryText());
7581
out.writeNamedWriteable(input());
76-
out.writeNamedWriteable(inferenceId());
7782
}
7883

7984
@Override
8085
protected NodeInfo<? extends PhysicalPlan> info() {
81-
return NodeInfo.create(this, RerankExec::new, child(), queryText, input, inferenceId);
86+
return NodeInfo.create(this, RerankExec::new, child(), inferenceId, queryText, input);
8287
}
8388

8489
@Override
8590
public UnaryExec replaceChild(PhysicalPlan newChild) {
86-
return new RerankExec(source(), newChild, queryText, input, inferenceId);
91+
return new RerankExec(source(), newChild, inferenceId, queryText, input);
8792
}
8893

8994
@Override

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
package org.elasticsearch.xpack.esql.planner;
99

10-
import org.apache.lucene.util.BytesRef;
1110
import org.elasticsearch.common.settings.Settings;
1211
import org.elasticsearch.common.util.BigArrays;
1312
import org.elasticsearch.compute.Describable;
@@ -680,8 +679,8 @@ private PhysicalOperation planRerank(RerankExec rerank, LocalExecutionPlannerCon
680679

681680
ExpressionEvaluator.Factory inputEvaluatorSupplier = EvalMapper.toEvaluator(rerank.input(), source.layout);
682681

683-
String inferenceId = rerank.inferenceId().fold().toString();
684-
String queryText = ((BytesRef) rerank.queryText().fold()).utf8ToString();
682+
String inferenceId = rerank.inferenceId();
683+
String queryText = rerank.queryText();
685684

686685
int scoreChannel = -1;
687686

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ static PhysicalPlan mapUnary(UnaryPlan p, PhysicalPlan child) {
8686
}
8787

8888
if (p instanceof Rerank rerank) {
89-
return new RerankExec(rerank.source(), child, rerank.queryText(), rerank.input(), rerank.inferenceId());
89+
return new RerankExec(rerank.source(), child, rerank.inferenceId(), rerank.queryText(), rerank.input());
9090
}
9191

9292
if (p instanceof Enrich enrich) {

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,8 @@ private ActualResults executePlan(BigArrays bigArrays) throws Exception {
506506
TEST_VERIFIER,
507507
new PlanningMetrics(),
508508
null,
509-
EsqlTestUtils.MOCK_QUERY_BUILDER_RESOLVER
509+
EsqlTestUtils.MOCK_QUERY_BUILDER_RESOLVER,
510+
EsqlTestUtils.MOCK_INFERENCE_SERVICE
510511
);
511512
TestPhysicalOperationProviders physicalOperationProviders = testOperationProviders(testDatasets);
512513

@@ -619,6 +620,7 @@ void executeSubPlan(
619620
exchangeSink,
620621
Mockito.mock(EnrichLookupService.class),
621622
Mockito.mock(LookupFromIndexService.class),
623+
EsqlTestUtils.MOCK_INFERENCE_SERVICE,
622624
physicalOperationProviders
623625
);
624626

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7562,6 +7562,7 @@ private LocalExecutionPlanner.LocalExecutionPlan physicalOperationsFromPhysicalP
75627562
new ExchangeSinkHandler(null, 10, () -> 10),
75637563
null,
75647564
null,
7565+
EsqlTestUtils.MOCK_INFERENCE_SERVICE,
75657566
new EsPhysicalOperationProviders(List.of(), null)
75667567
);
75677568

0 commit comments

Comments
 (0)