Skip to content

Commit 2e9bf3c

Browse files
afoucretafoucret
authored
[8.19] [ES|QL] RERANK command - Updating the syntax and behavior (elastic#129488) (elastic#129718)
* [ES|QL] RERANK command - Updating the syntax and behavior (elastic#129488)
1 parent 1c30e34 commit 2e9bf3c

File tree

25 files changed

+1745
-1031
lines changed

25 files changed

+1745
-1031
lines changed

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RequestIndexFilteringTestCase.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ public void testIndicesDontExist() throws IOException {
225225
// currently we don't support remote clusters in LOOKUP JOIN
226226
// this check happens before resolving actual indices and results in a different error message
227227
RemoteClusterAware.isRemoteIndexName(pattern)
228-
? allOf(containsString("parsing_exception"), containsString("remote clusters are not supported in LOOKUP JOIN"))
228+
? allOf(containsString("parsing_exception"), containsString("remote clusters are not supported"))
229229
: allOf(containsString("verification_exception"), containsString("Unknown index [foo]"))
230230
);
231231
}

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/rerank.csv-spec‎

Lines changed: 82 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -3,21 +3,62 @@
33
// This makes the output more predictable which is helpful here.
44

55

6-
reranker using a single field
6+
reranker using a single field, overwrite existing _score column
77
required_capability: rerank
88
required_capability: match_operator_colon
99

1010
FROM books METADATA _score
1111
| WHERE title:"war and peace" AND author:"Tolstoy"
12-
| RERANK "war and peace" ON title WITH test_reranker
13-
| KEEP book_no, title, author
12+
| SORT _score DESC, book_no ASC
13+
| RERANK "war and peace" ON title WITH inferenceId=test_reranker
14+
| EVAL _score=ROUND(_score, 2)
15+
| KEEP book_no, title, author, _score
1416
;
1517

16-
book_no:keyword | title:text | author:text
17-
5327 | War and Peace | Leo Tolstoy
18-
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy]
19-
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo
20-
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy
18+
book_no:keyword | title:text | author:text | _score:double
19+
5327 | War and Peace | Leo Tolstoy | 0.08
20+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.03
21+
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo | 0.03
22+
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy | 0.02
23+
;
24+
25+
reranker using a single field, create a mew column
26+
required_capability: rerank
27+
required_capability: match_operator_colon
28+
29+
FROM books METADATA _score
30+
| WHERE title:"war and peace" AND author:"Tolstoy"
31+
| SORT _score DESC, book_no ASC
32+
| RERANK "war and peace" ON title WITH inferenceId=test_reranker, scoreColumn=rerank_score
33+
| EVAL _score=ROUND(_score, 2), rerank_score=ROUND(rerank_score, 2)
34+
| KEEP book_no, title, author, rerank_score
35+
;
36+
37+
book_no:keyword | title:text | author:text | rerank_score:double
38+
5327 | War and Peace | Leo Tolstoy | 0.08
39+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.03
40+
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo | 0.03
41+
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy | 0.02
42+
;
43+
44+
reranker using a single field, create a mew column, sort by rerank_score
45+
required_capability: rerank
46+
required_capability: match_operator_colon
47+
48+
FROM books METADATA _score
49+
| WHERE title:"war and peace" AND author:"Tolstoy"
50+
| SORT _score DESC
51+
| RERANK "war and peace" ON title WITH inferenceId=test_reranker, scoreColumn=rerank_score
52+
| EVAL _score=ROUND(_score, 2), rerank_score=ROUND(rerank_score, 2)
53+
| SORT rerank_score, _score ASC, book_no ASC
54+
| KEEP book_no, title, author, rerank_score
55+
;
56+
57+
book_no:keyword | title:text | author:text | rerank_score:double
58+
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy | 0.02
59+
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo | 0.03
60+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.03
61+
5327 | War and Peace | Leo Tolstoy | 0.08
2162
;
2263

2364

@@ -27,15 +68,17 @@ required_capability: match_operator_colon
2768

2869
FROM books METADATA _score
2970
| WHERE title:"war and peace" AND author:"Tolstoy"
30-
| RERANK "war and peace" ON title, author WITH test_reranker
31-
| KEEP book_no, title, author
71+
| RERANK "war and peace" ON title, author WITH inferenceId=test_reranker
72+
| EVAL _score=ROUND(_score, 2)
73+
| SORT _score DESC, book_no ASC
74+
| KEEP book_no, title, author, _score
3275
;
3376

34-
book_no:keyword | title:text | author:text
35-
5327 | War and Peace | Leo Tolstoy
36-
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo
37-
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy
38-
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy]
77+
book_no:keyword | title:text | author:text | _score:double
78+
5327 | War and Peace | Leo Tolstoy | 0.02
79+
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy | 0.01
80+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.01
81+
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo | 0.01
3982
;
4083

4184

@@ -45,16 +88,18 @@ required_capability: match_operator_colon
4588

4689
FROM books METADATA _score
4790
| WHERE title:"war and peace" AND author:"Tolstoy"
48-
| SORT _score DESC
91+
| SORT _score DESC, book_no ASC
4992
| LIMIT 3
50-
| RERANK "war and peace" ON title WITH test_reranker
51-
| KEEP book_no, title, author
93+
| RERANK "war and peace" ON title WITH inferenceId=test_reranker
94+
| EVAL _score=ROUND(_score, 2)
95+
| SORT _score DESC, book_no ASC
96+
| KEEP book_no, title, author, _score
5297
;
5398

54-
book_no:keyword | title:text | author:text
55-
5327 | War and Peace | Leo Tolstoy
56-
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy]
57-
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo
99+
book_no:keyword | title:text | author:text | _score:double
100+
5327 | War and Peace | Leo Tolstoy | 0.08
101+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.03
102+
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo | 0.03
58103
;
59104

60105

@@ -64,15 +109,17 @@ required_capability: match_operator_colon
64109

65110
FROM books METADATA _score
66111
| WHERE title:"war and peace" AND author:"Tolstoy"
67-
| RERANK "war and peace" ON title WITH test_reranker
68-
| KEEP book_no, title, author
112+
| RERANK "war and peace" ON title WITH inferenceId=test_reranker
113+
| EVAL _score=ROUND(_score, 2)
114+
| SORT _score DESC, book_no ASC
115+
| KEEP book_no, title, author, _score
69116
| LIMIT 3
70117
;
71118

72-
book_no:keyword | title:text | author:text
73-
5327 | War and Peace | Leo Tolstoy
74-
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy]
75-
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo
119+
book_no:keyword | title:text | author:text | _score:double
120+
5327 | War and Peace | Leo Tolstoy | 0.08
121+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.03
122+
9032 | War and Peace: A Novel (6 Volumes) | Tolstoy Leo | 0.03
76123
;
77124

78125

@@ -82,14 +129,15 @@ required_capability: match_operator_colon
82129

83130
FROM books
84131
| WHERE title:"war and peace" AND author:"Tolstoy"
85-
| RERANK "war and peace" ON title WITH test_reranker
86-
| KEEP book_no, title, author
132+
| RERANK "war and peace" ON title WITH inferenceId=test_reranker
133+
| EVAL _score=ROUND(_score, 2)
134+
| KEEP book_no, title, author, _score
87135
| SORT author, title
88136
| LIMIT 3
89137
;
90138

91-
book_no:keyword | title:text | author:text
92-
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy]
93-
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy
94-
5327 | War and Peace | Leo Tolstoy
139+
book_no:keyword | title:text | author:text | _score:double
140+
4536 | War and Peace (Signet Classics) | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.03
141+
2776 | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy | 0.02
142+
5327 | War and Peace | Leo Tolstoy | 0.08
95143
;

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -387,8 +387,21 @@ joinPredicate
387387
: valueExpression
388388
;
389389

390+
inferenceCommandOptions
391+
: inferenceCommandOption (COMMA inferenceCommandOption)*
392+
;
393+
394+
inferenceCommandOption
395+
: identifier ASSIGN inferenceCommandOptionValue
396+
;
397+
398+
inferenceCommandOptionValue
399+
: constant
400+
| identifier
401+
;
402+
390403
rerankCommand
391-
: DEV_RERANK queryText=constant ON rerankFields (WITH inferenceId=identifierOrParameter)?
404+
: DEV_RERANK queryText=constant ON rerankFields (WITH inferenceCommandOptions)?
392405
;
393406

394407
completionCommand

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -750,7 +750,11 @@ private LogicalPlan resolveRerank(Rerank rerank, List<Attribute> childrenOutput)
750750
if (rerank.scoreAttribute() instanceof UnresolvedAttribute ua) {
751751
Attribute resolved = resolveAttribute(ua, childrenOutput);
752752
if (resolved.resolved() == false || resolved.dataType() != DOUBLE) {
753-
resolved = MetadataAttribute.create(Source.EMPTY, MetadataAttribute.SCORE);
753+
if (ua.name().equals(MetadataAttribute.SCORE)) {
754+
resolved = MetadataAttribute.create(Source.EMPTY, MetadataAttribute.SCORE);
755+
} else {
756+
resolved = new ReferenceAttribute(resolved.source(), resolved.name(), DOUBLE);
757+
}
754758
}
755759
rerank = rerank.withScoreAttribute(resolved);
756760
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,9 @@
3636
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits;
3737
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineOrderBy;
3838
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineSample;
39-
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownCompletion;
4039
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEnrich;
4140
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEval;
41+
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownInferencePlan;
4242
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownRegexExtract;
4343
import org.elasticsearch.xpack.esql.optimizer.rules.logical.RemoveStatsOverride;
4444
import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateAggExpressionWithEval;
@@ -187,7 +187,7 @@ protected static Batch<LogicalPlan> operators() {
187187
new PushDownAndCombineLimits(),
188188
new PushDownAndCombineFilters(),
189189
new PushDownAndCombineSample(),
190-
new PushDownCompletion(),
190+
new PushDownInferencePlan(),
191191
new PushDownEval(),
192192
new PushDownRegexExtract(),
193193
new PushDownEnrich(),

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineFilters.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import org.elasticsearch.xpack.esql.plan.logical.Project;
2525
import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
2626
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
27-
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
27+
import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
2828
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
2929
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
3030

@@ -71,10 +71,10 @@ protected LogicalPlan rule(Filter filter) {
7171
// Push down filters that do not rely on attributes created by RegexExtract
7272
var attributes = AttributeSet.of(Expressions.asAttributes(re.extractedFields()));
7373
plan = maybePushDownPastUnary(filter, re, attributes::contains, NO_OP);
74-
} else if (child instanceof Completion completion) {
74+
} else if (child instanceof InferencePlan<?> inferencePlan) {
7575
// Push down filters that do not rely on attributes created by Cpmpletion
76-
var attributes = AttributeSet.of(completion.generatedAttributes());
77-
plan = maybePushDownPastUnary(filter, completion, attributes::contains, NO_OP);
76+
var attributes = AttributeSet.of(inferencePlan.generatedAttributes());
77+
plan = maybePushDownPastUnary(filter, inferencePlan, attributes::contains, NO_OP);
7878
} else if (child instanceof Enrich enrich) {
7979
// Push down filters that do not rely on attributes created by Enrich
8080
var attributes = AttributeSet.of(Expressions.asAttributes(enrich.enrichFields()));

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineLimits.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import org.elasticsearch.xpack.esql.plan.logical.Project;
1818
import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
1919
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
20-
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
20+
import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
2121
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
2222
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
2323

@@ -43,7 +43,7 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
4343
|| unary instanceof Project
4444
|| unary instanceof RegexExtract
4545
|| unary instanceof Enrich
46-
|| unary instanceof Completion) {
46+
|| unary instanceof InferencePlan<?>) {
4747
return unary.replaceChild(limit.replaceChild(unary.child()));
4848
} else if (unary instanceof MvExpand) {
4949
// MV_EXPAND can increase the number of rows, so we cannot just push the limit down

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownCompletion.java‎ renamed to ‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownInferencePlan.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,11 @@
88
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
99

1010
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
11-
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
11+
import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
1212

13-
public final class PushDownCompletion extends OptimizerRules.OptimizerRule<Completion> {
13+
public final class PushDownInferencePlan extends OptimizerRules.OptimizerRule<InferencePlan<?>> {
1414
@Override
15-
protected LogicalPlan rule(Completion p) {
15+
protected LogicalPlan rule(InferencePlan<?> p) {
1616
return PushDownUtils.pushGeneratingPlanPastProjectAndOrderBy(p);
1717
}
1818
}

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

Lines changed: 4 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)