Skip to content

Commit a56cb42

Browse files
afoucretafoucret
authored
[8.19] Manual backport of ES|QL inference features (RERANK, COMPLETION). (#128907)
1 parent eecc28d commit a56cb42

File tree

105 files changed

+8954
-2733
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

105 files changed

+8954
-2733
lines changed

‎docs/changelog/123074.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 123074
2+
summary: Adding ES|QL Reranker command in snapshot builds
3+
area: Ranking
4+
type: feature
5+
issues: []

‎docs/changelog/126319.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 126319
2+
summary: COMPLETION command grammar and logical plan
3+
area: ES|QL
4+
type: feature
5+
issues: []

‎docs/changelog/127731.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 127731
2+
summary: ESQL - Enable telemetry for COMPLETION command
3+
area: Search
4+
type: feature
5+
issues: []

‎docs/changelog/128948.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
pr: 128948
2+
summary: ES|QL - Add COMPLETION command as a tech preview feature
3+
area: ES|QL
4+
type: feature
5+
issues:
6+
- 124405

‎docs/reference/rest-api/usage.asciidoc‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,9 @@ GET /_xpack/usage
247247
"lookup" : 0,
248248
"inlinestats" : 0,
249249
"lookup_join" : 0,
250-
"change_point" : 0
250+
"change_point" : 0,
251+
"completion": 0,
252+
"rerank": 0
251253
},
252254
"queries" : {
253255
"rest" : {

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,10 @@ public static class Request extends BaseInferenceActionRequest {
6464
public static final ParseField TOP_N = new ParseField("top_n");
6565
public static final ParseField TIMEOUT = new ParseField("timeout");
6666

67+
public static Builder builder(String inferenceEntityId, TaskType taskType) {
68+
return new Builder().setInferenceEntityId(inferenceEntityId).setTaskType(taskType);
69+
}
70+
6771
static final ObjectParser<Request.Builder, Void> PARSER = new ObjectParser<>(NAME, Request.Builder::new);
6872
static {
6973
PARSER.declareStringArray(Request.Builder::setInput, INPUT);

‎x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/OperatorTestCase.java‎

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -189,13 +189,11 @@ protected final List<Page> oneDriverPerPageList(Iterator<List<Page>> source, Sup
189189
while (source.hasNext()) {
190190
List<Page> in = source.next();
191191
try (
192-
Driver d = new Driver(
193-
"test",
192+
Driver d = TestDriverFactory.create(
194193
driverContext(),
195194
new CannedSourceOperator(in.iterator()),
196195
operators.get(),
197-
new PageConsumerOperator(result::add),
198-
() -> {}
196+
new PageConsumerOperator(result::add)
199197
)
200198
) {
201199
runDriver(d);
@@ -275,13 +273,11 @@ protected final List<Page> drive(List<Operator> operators, Iterator<Page> input,
275273
List<Page> results = new ArrayList<>();
276274
boolean success = false;
277275
try (
278-
Driver d = new Driver(
279-
"test",
276+
Driver d = TestDriverFactory.create(
280277
driverContext,
281278
new CannedSourceOperator(input),
282279
operators,
283-
new TestResultPageSinkOperator(results::add),
284-
() -> {}
280+
new TestResultPageSinkOperator(results::add)
285281
)
286282
) {
287283
runDriver(d);
@@ -303,22 +299,15 @@ public static void runDriver(List<Driver> drivers) {
303299
int dummyDrivers = between(0, 10);
304300
for (int i = 0; i < dummyDrivers; i++) {
305301
drivers.add(
306-
new Driver(
307-
"test",
308-
"dummy-session",
309-
0,
310-
0,
302+
TestDriverFactory.create(
311303
new DriverContext(BigArrays.NON_RECYCLING_INSTANCE, TestBlockFactory.getNonBreakingInstance()),
312-
() -> "dummy-driver",
313304
new SequenceLongBlockSourceOperator(
314305
TestBlockFactory.getNonBreakingInstance(),
315306
LongStream.range(0, between(1, 100)),
316307
between(1, 100)
317308
),
318309
List.of(),
319-
new PageConsumerOperator(page -> page.releaseBlocks()),
320-
Driver.DEFAULT_STATUS_INTERVAL,
321-
() -> {}
310+
new PageConsumerOperator(Page::releaseBlocks)
322311
)
323312
);
324313
}

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

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -66,11 +66,13 @@
6666
import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
6767
import static org.elasticsearch.xpack.esql.CsvTestUtils.loadCsvSpecValues;
6868
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs;
69-
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.clusterHasInferenceEndpoint;
70-
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createInferenceEndpoint;
71-
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoint;
69+
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createInferenceEndpoints;
70+
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoints;
7271
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
7372
import static org.elasticsearch.xpack.esql.EsqlTestUtils.classpathResources;
73+
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.COMPLETION;
74+
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.RERANK;
75+
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SEMANTIC_TEXT_FIELD_CAPS;
7476
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.cap;
7577

7678
// This test can run very long in serverless configurations
@@ -133,8 +135,8 @@ protected EsqlSpecTestCase(
133135

134136
@Before
135137
public void setup() throws IOException {
136-
if (supportsInferenceTestService() && clusterHasInferenceEndpoint(client()) == false) {
137-
createInferenceEndpoint(client());
138+
if (supportsInferenceTestService()) {
139+
createInferenceEndpoints(adminClient());
138140
}
139141

140142
if (indexExists(availableDatasetsForEs(client(), supportsIndexModeLookup()).iterator().next().indexName()) == false) {
@@ -157,7 +159,8 @@ public static void wipeTestData() throws IOException {
157159
}
158160
}
159161

160-
deleteInferenceEndpoint(client());
162+
deleteInferenceEndpoints(adminClient());
163+
161164
}
162165

163166
public boolean logResults() {
@@ -174,8 +177,8 @@ public final void test() throws Throwable {
174177
}
175178

176179
protected void shouldSkipTest(String testName) throws IOException {
177-
if (testCase.requiredCapabilities.contains("semantic_text_field_caps") || testCase.requiredCapabilities.contains("rerank")) {
178-
assumeTrue("Inference test service needs to be supported for semantic_text", supportsInferenceTestService());
180+
if (requiresInferenceEndpoint()) {
181+
assumeTrue("Inference test service needs to be supported", supportsInferenceTestService());
179182
}
180183
checkCapabilities(adminClient(), testFeatureService, testName, testCase);
181184
assumeTrue("Test " + testName + " is not enabled", isEnabled(testName, instructions, Version.CURRENT));
@@ -245,6 +248,11 @@ protected boolean supportsInferenceTestService() {
245248
return true;
246249
}
247250

251+
protected boolean requiresInferenceEndpoint() {
252+
return Stream.of(SEMANTIC_TEXT_FIELD_CAPS.capabilityName(), RERANK.capabilityName(), COMPLETION.capabilityName())
253+
.anyMatch(testCase.requiredCapabilities::contains);
254+
}
255+
248256
protected boolean supportsIndexModeLookup() throws IOException {
249257
return true;
250258
}
@@ -340,6 +348,11 @@ private Object valueMapper(CsvTestUtils.Type type, Object value) {
340348
return new BigDecimal(s).round(new MathContext(7, RoundingMode.DOWN)).doubleValue();
341349
}
342350
}
351+
if (type == CsvTestUtils.Type.TEXT || type == CsvTestUtils.Type.KEYWORD || type == CsvTestUtils.Type.SEMANTIC_TEXT) {
352+
if (value instanceof String s) {
353+
value = s.replaceAll("\\\\n", "\n");
354+
}
355+
}
343356
return value.toString();
344357
}
345358

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -780,7 +780,7 @@ public void testNamedParamsForIdentifierAndIdentifierPatterns() throws IOExcepti
780780
);
781781
error = re.getMessage();
782782
assertThat(error, containsString("ParsingException"));
783-
assertThat(error, containsString("line 1:23: mismatched input '?cmd' expecting {'dissect', 'drop'"));
783+
assertThat(error, containsString("line 1:23: mismatched input '?cmd' expecting {'completion', 'dissect'"));
784784
}
785785
}
786786

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

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.rest;
9+
10+
import org.elasticsearch.client.Request;
11+
import org.elasticsearch.client.ResponseException;
12+
import org.elasticsearch.test.rest.ESRestTestCase;
13+
import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
14+
import org.junit.After;
15+
import org.junit.Before;
16+
17+
import java.io.IOException;
18+
import java.util.List;
19+
import java.util.Map;
20+
21+
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createRerankInferenceEndpoint;
22+
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteRerankInferenceEndpoint;
23+
import static org.hamcrest.core.StringContains.containsString;
24+
25+
public class RestRerankTestCase extends ESRestTestCase {
26+
27+
@Before
28+
public void skipWhenRerankDisabled() throws IOException {
29+
assumeTrue(
30+
"Requires RERANK capability",
31+
EsqlSpecTestCase.hasCapabilities(adminClient(), List.of(EsqlCapabilities.Cap.RERANK.capabilityName()))
32+
);
33+
}
34+
35+
@Before
36+
@After
37+
public void assertRequestBreakerEmpty() throws Exception {
38+
EsqlSpecTestCase.assertRequestBreakerEmpty();
39+
}
40+
41+
@Before
42+
public void setUpInferenceEndpoint() throws IOException {
43+
createRerankInferenceEndpoint(adminClient());
44+
}
45+
46+
@Before
47+
public void setUpTestIndex() throws IOException {
48+
Request request = new Request("PUT", "/rerank-test-index");
49+
request.setJsonEntity("""
50+
{
51+
"mappings": {
52+
"properties": {
53+
"title": { "type": "text" },
54+
"author": { "type": "text" }
55+
}
56+
}
57+
}""");
58+
assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
59+
60+
request = new Request("POST", "/rerank-test-index/_bulk");
61+
request.addParameter("refresh", "true");
62+
request.setJsonEntity("""
63+
{ "index": {"_id": 1} }
64+
{ "title": "The Future of Exploration", "author": "John Doe" }
65+
{ "index": {"_id": 2} }
66+
{ "title": "Deep Sea Exploration", "author": "Jane Smith" }
67+
{ "index": {"_id": 3} }
68+
{ "title": "History of Space Exploration", "author": "Alice Johnson" }
69+
""");
70+
assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
71+
}
72+
73+
@After
74+
public void wipeData() throws IOException {
75+
try {
76+
adminClient().performRequest(new Request("DELETE", "/rerank-test-index"));
77+
} catch (ResponseException e) {
78+
// 404 here just means we had no indexes
79+
if (e.getResponse().getStatusLine().getStatusCode() != 404) {
80+
throw e;
81+
}
82+
}
83+
84+
deleteRerankInferenceEndpoint(adminClient());
85+
}
86+
87+
public void testRerankWithSingleField() throws IOException {
88+
String query = """
89+
FROM rerank-test-index
90+
| WHERE match(title, "exploration")
91+
| RERANK "exploration" ON title WITH test_reranker
92+
| EVAL _score = ROUND(_score, 5)
93+
""";
94+
95+
Map<String, Object> result = runEsqlQuery(query);
96+
97+
var expectedValues = List.of(
98+
List.of("Jane Smith", "Deep Sea Exploration", 0.02941d),
99+
List.of("John Doe", "The Future of Exploration", 0.02632d),
100+
List.of("Alice Johnson", "History of Space Exploration", 0.02381d)
101+
);
102+
103+
assertResultMap(result, defaultOutputColumns(), expectedValues);
104+
}
105+
106+
public void testRerankWithMultipleFields() throws IOException {
107+
String query = """
108+
FROM rerank-test-index
109+
| WHERE match(title, "exploration")
110+
| RERANK "exploration" ON title, author WITH test_reranker
111+
| EVAL _score = ROUND(_score, 5)
112+
""";
113+
114+
Map<String, Object> result = runEsqlQuery(query);
115+
;
116+
var expectedValues = List.of(
117+
List.of("Jane Smith", "Deep Sea Exploration", 0.01818d),
118+
List.of("John Doe", "The Future of Exploration", 0.01754d),
119+
List.of("Alice Johnson", "History of Space Exploration", 0.01515d)
120+
);
121+
122+
assertResultMap(result, defaultOutputColumns(), expectedValues);
123+
}
124+
125+
public void testRerankWithPositionalParams() throws IOException {
126+
String query = """
127+
FROM rerank-test-index
128+
| WHERE match(title, "exploration")
129+
| RERANK ? ON title WITH ?
130+
| EVAL _score = ROUND(_score, 5)
131+
""";
132+
133+
Map<String, Object> result = runEsqlQuery(query, "[\"exploration\", \"test_reranker\"]");
134+
135+
var expectedValues = List.of(
136+
List.of("Jane Smith", "Deep Sea Exploration", 0.02941d),
137+
List.of("John Doe", "The Future of Exploration", 0.02632d),
138+
List.of("Alice Johnson", "History of Space Exploration", 0.02381d)
139+
);
140+
141+
assertResultMap(result, defaultOutputColumns(), expectedValues);
142+
}
143+
144+
public void testRerankWithNamedParams() throws IOException {
145+
String query = """
146+
FROM rerank-test-index
147+
| WHERE match(title, ?queryText)
148+
| RERANK ?queryText ON title WITH ?inferenceId
149+
| EVAL _score = ROUND(_score, 5)
150+
""";
151+
152+
Map<String, Object> result = runEsqlQuery(query, "[{\"queryText\": \"exploration\"}, {\"inferenceId\": \"test_reranker\"}]");
153+
154+
var expectedValues = List.of(
155+
List.of("Jane Smith", "Deep Sea Exploration", 0.02941d),
156+
List.of("John Doe", "The Future of Exploration", 0.02632d),
157+
List.of("Alice Johnson", "History of Space Exploration", 0.02381d)
158+
);
159+
160+
assertResultMap(result, defaultOutputColumns(), expectedValues);
161+
}
162+
163+
public void testRerankWithMissingInferenceId() {
164+
String query = """
165+
FROM rerank-test-index
166+
| WHERE match(title, "exploration")
167+
| RERANK "exploration" ON title WITH test_missing
168+
| EVAL _score = ROUND(_score, 5)
169+
""";
170+
171+
ResponseException re = expectThrows(ResponseException.class, () -> runEsqlQuery(query));
172+
assertThat(re.getMessage(), containsString("Inference endpoint not found"));
173+
}
174+
175+
private static List<Map<String, String>> defaultOutputColumns() {
176+
return List.of(
177+
Map.of("name", "author", "type", "text"),
178+
Map.of("name", "title", "type", "text"),
179+
Map.of("name", "_score", "type", "double")
180+
);
181+
}
182+
183+
private Map<String, Object> runEsqlQuery(String query) throws IOException {
184+
RestEsqlTestCase.RequestObjectBuilder builder = RestEsqlTestCase.requestObjectBuilder().query(query);
185+
return RestEsqlTestCase.runEsqlSync(builder);
186+
}
187+
188+
private Map<String, Object> runEsqlQuery(String query, String params) throws IOException {
189+
RestEsqlTestCase.RequestObjectBuilder builder = RestEsqlTestCase.requestObjectBuilder().query(query).params(params);
190+
return RestEsqlTestCase.runEsqlSync(builder);
191+
}
192+
}

0 commit comments

Comments
 (0)