Skip to content

Commit 2490d7e

Browse files
afoucretafoucret
committed
Addressing PR feedbacks
1 parent eb84e24 commit 2490d7e

File tree

5 files changed

+56
-44
lines changed

5 files changed

+56
-44
lines changed

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/MapExpression.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,6 @@ public class MapExpression extends Expression {
3636
MapExpression::readFrom
3737
);
3838

39-
public static final MapExpression EMPTY = new MapExpression(Source.EMPTY, List.of());
40-
4139
private final List<EntryExpression> entryExpressions;
4240

4341
private final Map<Expression, Expression> map;

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

Lines changed: 30 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@
7272
import org.elasticsearch.xpack.esql.plan.logical.TimeSeriesAggregate;
7373
import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
7474
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
75+
import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
7576
import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
7677
import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
7778
import org.elasticsearch.xpack.esql.plan.logical.show.ShowInfo;
@@ -759,16 +760,12 @@ public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
759760

760761
if (queryText instanceof Literal queryTextLiteral && DataType.isString(queryText.dataType())) {
761762
if (queryTextLiteral.value() == null) {
762-
throw new ParsingException(
763-
source(ctx.queryText),
764-
"Query text cannot be null or undefined in RERANK",
765-
ctx.queryText.getText()
766-
);
763+
throw new ParsingException(source(ctx.queryText), "Query cannot be null or undefined in RERANK", ctx.queryText.getText());
767764
}
768765
} else {
769766
throw new ParsingException(
770767
source(ctx.queryText),
771-
"RERANK only support string as query text but [{}] cannot be used as string",
768+
"Query must be a valid string in RERANK, found [{}]",
772769
ctx.queryText.getText()
773770
);
774771
}
@@ -790,26 +787,16 @@ private Rerank applyRerankOptions(Rerank rerank, EsqlBaseParser.CommandNamedPara
790787
Expression inferenceId = optionsMap.remove(Rerank.INFERENCE_ID_OPTION_NAME);
791788

792789
if (inferenceId != null) {
793-
if (inferenceId instanceof Literal inferenceIdLiteral && DataType.isString(inferenceId.dataType())) {
794-
if (inferenceIdLiteral.value() == null) {
795-
throw new ParsingException(
796-
inferenceId.source(),
797-
"[{}] option cannot be null or undefined in RERANK",
798-
Rerank.INFERENCE_ID_OPTION_NAME
799-
);
800-
}
801-
rerank = rerank.withInferenceId(inferenceId);
802-
} else {
803-
throw new ParsingException(
804-
inferenceId.source(),
805-
"Option [{}] only support string in RERANK but [{}] cannot be used as string",
806-
Rerank.INFERENCE_ID_OPTION_NAME
807-
);
808-
}
790+
rerank = applyInferenceId(rerank, inferenceId);
809791
}
810792

811793
if (optionsMap.isEmpty() == false) {
812-
throw new ParsingException(source(ctx), "Unknown option [{}] in RERANK command", optionsMap.keySet().stream().findAny().get());
794+
throw new ParsingException(
795+
source(ctx),
796+
"Inavalid option [{}] in RERANK, expected one of [{}]",
797+
optionsMap.keySet().stream().findAny().get(),
798+
rerank.validOptionNames()
799+
);
813800
}
814801

815802
return rerank;
@@ -840,35 +827,37 @@ private Completion applyCompletionOptions(Completion completion, EsqlBaseParser.
840827

841828
Expression inferenceId = optionsMap.remove(Completion.INFERENCE_ID_OPTION_NAME);
842829
if (inferenceId != null) {
843-
if (inferenceId instanceof Literal inferenceIdLiteral && DataType.isString(inferenceId.dataType())) {
844-
if (inferenceIdLiteral.value() == null) {
845-
throw new ParsingException(
846-
inferenceId.source(),
847-
"[{}] option cannot be null or undefined in COMPLETION",
848-
Completion.INFERENCE_ID_OPTION_NAME
849-
);
850-
}
851-
completion = completion.withInferenceId(inferenceId);
852-
} else {
853-
throw new ParsingException(
854-
inferenceId.source(),
855-
"Option [{}] only support string in COMPLETION but [{}] cannot be used as string",
856-
Completion.INFERENCE_ID_OPTION_NAME
857-
);
858-
}
830+
completion = applyInferenceId(completion, inferenceId);
859831
}
860832

861833
if (optionsMap.isEmpty() == false) {
862834
throw new ParsingException(
863835
source(ctx),
864-
"Unknown option [{}] in Completion command",
865-
optionsMap.keySet().stream().findAny().get()
836+
"Inavalid option [{}] in COMPLETION, expected one of [{}]",
837+
optionsMap.keySet().stream().findAny().get(),
838+
completion.validOptionNames()
866839
);
867840
}
868841

869842
return completion;
870843
}
871844

845+
private <InferencePlanType extends InferencePlan<InferencePlanType>> InferencePlanType applyInferenceId(
846+
InferencePlanType inferencePlan,
847+
Expression inferenceId
848+
) {
849+
if ((inferenceId instanceof Literal && DataType.isString(inferenceId.dataType())) == false) {
850+
throw new ParsingException(
851+
inferenceId.source(),
852+
"Option [{}] must be a valid string, found [{}]",
853+
Completion.INFERENCE_ID_OPTION_NAME,
854+
inferenceId.source().text()
855+
);
856+
}
857+
858+
return inferencePlan.withInferenceId(inferenceId);
859+
}
860+
872861
public PlanFactory visitSampleCommand(EsqlBaseParser.SampleCommandContext ctx) {
873862
Source source = source(ctx);
874863
Object val = expression(ctx.probability).fold(FoldContext.small() /* TODO remove me */);

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
1919

2020
import java.io.IOException;
21+
import java.util.List;
2122
import java.util.Objects;
2223

2324
public abstract class InferencePlan<PlanType extends InferencePlan<PlanType>> extends UnaryPlan
@@ -26,6 +27,8 @@ public abstract class InferencePlan<PlanType extends InferencePlan<PlanType>> ex
2627
GeneratingPlan<InferencePlan<PlanType>> {
2728

2829
public static final String INFERENCE_ID_OPTION_NAME = "inference_id";
30+
public static final List<String> VALID_INFERENCE_OPTION_NAMES = List.of(INFERENCE_ID_OPTION_NAME);
31+
2932
private final Expression inferenceId;
3033

3134
protected InferencePlan(Source source, LogicalPlan child, Expression inferenceId) {
@@ -70,4 +73,8 @@ public int hashCode() {
7073
public PlanType withInferenceResolutionError(String inferenceId, String error) {
7174
return withInferenceId(new UnresolvedAttribute(inferenceId().source(), inferenceId, error));
7275
}
76+
77+
public List<String> validOptionNames() {
78+
return VALID_INFERENCE_OPTION_NAMES;
79+
}
7380
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5515,7 +5515,7 @@ record PushdownShadowingGeneratingPlanTestCase(
55155515
),
55165516
new PushDownEnrich()
55175517
),
5518-
// | COMPLETION y =CONCAT(some text, x) INTO WITH { "inference_id" : "inferenceID" }
5518+
// | COMPLETION y =CONCAT(some text, x) WITH { "inference_id" : "inferenceID" }
55195519
new PushdownShadowingGeneratingPlanTestCase(
55205520
(plan, attr) -> new Completion(
55215521
EMPTY,

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3868,6 +3868,15 @@ public void testRerankWithNamedParameters() {
38683868

38693869
public void testInvalidRerank() {
38703870
assumeTrue("RERANK requires corresponding capability", EsqlCapabilities.Cap.RERANK.isEnabled());
3871+
expectError(
3872+
"FROM foo* | RERANK \"query text\" ON title WITH { \"inference_id\": 3 }",
3873+
"line 1:65: Option [inference_id] must be a valid string, found [3]"
3874+
);
3875+
expectError(
3876+
"FROM foo* | RERANK \"query text\" ON title WITH { \"inference_id\": \"inferenceId\", \"unknown_option\": 3 }",
3877+
"line 1:42: Inavalid option [unknown_option] in RERANK, expected one of [[inference_id]]"
3878+
);
3879+
expectError("FROM foo* | RERANK 45 ON title", "Query must be a valid string in RERANK, found [45]");
38713880
expectError("FROM foo* | RERANK ON title WITH inferenceId", "line 1:20: extraneous input 'ON' expecting {QUOTED_STRING");
38723881
expectError("FROM foo* | RERANK \"query text\" WITH inferenceId", "line 1:33: mismatched input 'WITH' expecting 'on'");
38733882

@@ -3952,6 +3961,15 @@ public void testCompletionWithNamedParameters() {
39523961
}
39533962

39543963
public void testInvalidCompletion() {
3964+
expectError(
3965+
"FROM foo* | COMPLETION prompt WITH { \"inference_id\": 3 }",
3966+
"line 1:54: Option [inference_id] must be a valid string, found [3]"
3967+
);
3968+
expectError(
3969+
"FROM foo* | COMPLETION prompt WITH { \"inference_id\": \"inferenceId\", \"unknown_option\": 3 }",
3970+
"line 1:42: Inavalid option [unknown_option] in COMPLETION, expected one of [[inference_id]]"
3971+
);
3972+
39553973
expectError("FROM foo* | COMPLETION WITH inferenceId", "line 1:24: extraneous input 'WITH' expecting {");
39563974

39573975
expectError("FROM foo* | COMPLETION completion=prompt WITH", "ine 1:46: mismatched input '<EOF>' expecting '{'");

0 commit comments

Comments
 (0)