Skip to content

Commit 34ccaba

Browse files
afoucretafoucret
authored
[ES|QL] RERANK command - Updating the syntax and behavior (#129488)
1 parent a0109bb commit 34ccaba

24 files changed

+1804
-1114
lines changed

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

Lines changed: 89 additions & 39 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,16 +129,17 @@ 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
;
96144

97145

@@ -105,12 +153,14 @@ FROM books METADATA _id, _index, _score
105153
| FORK ( WHERE title:"Tolkien" | SORT _score, _id DESC | LIMIT 3 )
106154
( WHERE author:"Tolkien" | SORT _score, _id DESC | LIMIT 3 )
107155
| RRF
108-
| RERANK "Tolkien" ON title WITH test_reranker
156+
| RERANK "Tolkien" ON title WITH inferenceId=test_reranker
157+
| EVAL _score=ROUND(_score, 2)
158+
| SORT _score DESC, book_no ASC
109159
| LIMIT 2
110-
| KEEP book_no, title, author
160+
| KEEP book_no, title, author, _score
111161
;
112162

113-
book_no:keyword | title:keyword | author:keyword
114-
5335 | Letters of J R R Tolkien | J.R.R. Tolkien
115-
2130 | The J. R. R. Tolkien Audio Collection | [Christopher Tolkien, John Ronald Reuel Tolkien]
163+
book_no:keyword | title:keyword | author:keyword | _score:double
164+
5335 | Letters of J R R Tolkien | J.R.R. Tolkien | 0.04
165+
2130 | The J. R. R. Tolkien Audio Collection | [Christopher Tolkien, John Ronald Reuel Tolkien] | 0.03
116166
;

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -306,8 +306,21 @@ rrfCommand
306306
: DEV_RRF
307307
;
308308

309+
inferenceCommandOptions
310+
: inferenceCommandOption (COMMA inferenceCommandOption)*
311+
;
312+
313+
inferenceCommandOption
314+
: identifier ASSIGN inferenceCommandOptionValue
315+
;
316+
317+
inferenceCommandOptionValue
318+
: constant
319+
| identifier
320+
;
321+
309322
rerankCommand
310-
: DEV_RERANK queryText=constant ON rerankFields (WITH inferenceId=identifierOrParameter)?
323+
: DEV_RERANK queryText=constant ON rerankFields (WITH inferenceCommandOptions)?
311324
;
312325

313326
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
@@ -841,7 +841,11 @@ private LogicalPlan resolveRerank(Rerank rerank, List<Attribute> childrenOutput)
841841
if (rerank.scoreAttribute() instanceof UnresolvedAttribute ua) {
842842
Attribute resolved = resolveAttribute(ua, childrenOutput);
843843
if (resolved.resolved() == false || resolved.dataType() != DOUBLE) {
844-
resolved = MetadataAttribute.create(Source.EMPTY, MetadataAttribute.SCORE);
844+
if (ua.name().equals(MetadataAttribute.SCORE)) {
845+
resolved = MetadataAttribute.create(Source.EMPTY, MetadataAttribute.SCORE);
846+
} else {
847+
resolved = new ReferenceAttribute(resolved.source(), resolved.name(), DOUBLE);
848+
}
845849
}
846850
rerank = rerank.withScoreAttribute(resolved);
847851
}

‎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
@@ -38,9 +38,9 @@
3838
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits;
3939
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineOrderBy;
4040
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineSample;
41-
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownCompletion;
4241
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEnrich;
4342
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEval;
43+
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownInferencePlan;
4444
import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownRegexExtract;
4545
import org.elasticsearch.xpack.esql.optimizer.rules.logical.RemoveStatsOverride;
4646
import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateAggExpressionWithEval;
@@ -192,7 +192,7 @@ protected static Batch<LogicalPlan> operators() {
192192
new PushDownAndCombineLimits(),
193193
new PushDownAndCombineFilters(),
194194
new PushDownAndCombineSample(),
195-
new PushDownCompletion(),
195+
new PushDownInferencePlan(),
196196
new PushDownEval(),
197197
new PushDownRegexExtract(),
198198
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.InlineJoin;
2929
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
3030
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
@@ -72,10 +72,10 @@ protected LogicalPlan rule(Filter filter) {
7272
// Push down filters that do not rely on attributes created by RegexExtract
7373
var attributes = AttributeSet.of(Expressions.asAttributes(re.extractedFields()));
7474
plan = maybePushDownPastUnary(filter, re, attributes::contains, NO_OP);
75-
} else if (child instanceof Completion completion) {
75+
} else if (child instanceof InferencePlan<?> inferencePlan) {
7676
// Push down filters that do not rely on attributes created by Cpmpletion
77-
var attributes = AttributeSet.of(completion.generatedAttributes());
78-
plan = maybePushDownPastUnary(filter, completion, attributes::contains, NO_OP);
77+
var attributes = AttributeSet.of(inferencePlan.generatedAttributes());
78+
plan = maybePushDownPastUnary(filter, inferencePlan, attributes::contains, NO_OP);
7979
} else if (child instanceof Enrich enrich) {
8080
// Push down filters that do not rely on attributes created by Enrich
8181
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)