Skip to content

Commit e1b241e

Browse files
committed
[ES|QL] RERANK command default inferenceId (#128685)
* [ES|QL] Make inferenceId optional for the `RERANK` command (default to `.rerank-v1-elasticsearch`) * lint * Remove old test. * Fix failing test
1 parent e01703d commit e1b241e

File tree

7 files changed

+22
-7354
lines changed

7 files changed

+22
-7354
lines changed

x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -378,7 +378,7 @@ joinPredicate
378378
;
379379

380380
rerankCommand
381-
: DEV_RERANK queryText=constant ON rerankFields WITH inferenceId=identifierOrParameter
381+
: DEV_RERANK queryText=constant ON rerankFields (WITH inferenceId=identifierOrParameter)?
382382
;
383383

384384
completionCommand

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseParser.interp

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

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

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

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -662,7 +662,11 @@ public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
662662
);
663663
}
664664

665-
return p -> new Rerank(source, p, inferenceId(ctx.inferenceId), queryText, visitRerankFields(ctx.rerankFields()));
665+
Literal inferenceId = ctx.inferenceId != null
666+
? inferenceId(ctx.inferenceId)
667+
: new Literal(source, Rerank.DEFAULT_INFERENCE_ID, KEYWORD);
668+
669+
return p -> new Rerank(source, p, inferenceId, queryText, visitRerankFields(ctx.rerankFields()));
666670
}
667671

668672
@Override

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
public class Rerank extends InferencePlan<Rerank> implements SortAgnostic, SurrogateLogicalPlan, TelemetryAware {
4242

4343
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Rerank", Rerank::new);
44+
public static final Object DEFAULT_INFERENCE_ID = ".rerank-v1-elasticsearch";
4445
private final Attribute scoreAttribute;
4546
private final Expression queryText;
4647
private final List<Alias> rerankFields;

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3179,7 +3179,10 @@ public void testResolveRerankFields() {
31793179
| RERANK "food" ON title, SUBSTRING(description, 0, 100), yearRenamed=year WITH `reranking-inference-id`
31803180
""", "mapping-books.json");
31813181
} catch (ParsingException ex) {
3182-
assertThat(ex.getMessage(), containsString("line 3:36: mismatched input '(' expecting {'=', ',', '.', 'with'}"));
3182+
assertThat(
3183+
ex.getMessage(),
3184+
containsString("line 3:36: mismatched input '(' expecting {<EOF>, '|', '=', ',', '.', 'with'}")
3185+
);
31833186
}
31843187
}
31853188

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3220,6 +3220,17 @@ public void testInvalidJoinPatterns() {
32203220
}
32213221
}
32223222

3223+
public void testRerankDefaultInferenceId() {
3224+
assumeTrue("RERANK requires corresponding capability", EsqlCapabilities.Cap.RERANK.isEnabled());
3225+
3226+
var plan = processingCommand("RERANK \"query text\" ON title");
3227+
var rerank = as(plan, Rerank.class);
3228+
3229+
assertThat(rerank.inferenceId(), equalTo(literalString(".rerank-v1-elasticsearch")));
3230+
assertThat(rerank.queryText(), equalTo(literalString("query text")));
3231+
assertThat(rerank.rerankFields(), equalTo(List.of(alias("title", attribute("title")))));
3232+
}
3233+
32233234
public void testRerankSingleField() {
32243235
assumeTrue("RERANK requires corresponding capability", EsqlCapabilities.Cap.RERANK.isEnabled());
32253236

@@ -3295,10 +3306,7 @@ public void testRerankWithNamedParameters() {
32953306
public void testInvalidRerank() {
32963307
assumeTrue("RERANK requires corresponding capability", EsqlCapabilities.Cap.RERANK.isEnabled());
32973308
expectError("FROM foo* | RERANK ON title WITH inferenceId", "line 1:20: mismatched input 'ON' expecting {QUOTED_STRING");
3298-
32993309
expectError("FROM foo* | RERANK \"query text\" WITH inferenceId", "line 1:33: mismatched input 'WITH' expecting 'on'");
3300-
3301-
expectError("FROM foo* | RERANK \"query text\" ON title", "line 1:41: mismatched input '<EOF>' expecting {'=', ',', '.', 'with'}");
33023310
}
33033311

33043312
public void testCompletionUsingFieldAsPrompt() {

0 commit comments

Comments
 (0)