[8.19] Manual backport of ES|QL inference features (RERANK, COMPLETION).

docs/changelog/123074.yaml

@@ -0,0 +1,5 @@
+pr: 123074
+summary: Adding ES|QL Reranker command in snapshot builds
+area: Ranking
+type: feature
+issues: []

docs/changelog/126319.yaml

@@ -0,0 +1,5 @@
+pr: 126319
+summary: COMPLETION command grammar and logical plan
+area: ES|QL
+type: feature
+issues: []

docs/changelog/127731.yaml

@@ -0,0 +1,5 @@
+pr: 127731
+summary: ESQL - Enable telemetry for COMPLETION command
+area: Search
+type: feature
+issues: []

docs/changelog/128948.yaml

@@ -0,0 +1,6 @@
+pr: 128948
+summary: ES|QL - Add COMPLETION command as a tech preview feature
+area: ES|QL
+type: feature
+issues:
+- 124405

docs/reference/rest-api/usage.asciidoc

@@ -247,7 +247,9 @@ GET /_xpack/usage
       "lookup" : 0,
       "inlinestats" : 0,
       "lookup_join" : 0,
-      "change_point" : 0
+      "change_point" : 0,
+      "completion": 0,
+      "rerank": 0
     },
     "queries" : {
       "rest" : {

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

@@ -64,6 +64,10 @@ public static class Request extends BaseInferenceActionRequest {
         public static final ParseField TOP_N = new ParseField("top_n");
         public static final ParseField TIMEOUT = new ParseField("timeout");
 
+        public static Builder builder(String inferenceEntityId, TaskType taskType) {
+            return new Builder().setInferenceEntityId(inferenceEntityId).setTaskType(taskType);
+        }
+
         static final ObjectParser<Request.Builder, Void> PARSER = new ObjectParser<>(NAME, Request.Builder::new);
         static {
             PARSER.declareStringArray(Request.Builder::setInput, INPUT);

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

@@ -189,13 +189,11 @@ protected final List<Page> oneDriverPerPageList(Iterator<List<Page>> source, Sup
         while (source.hasNext()) {
             List<Page> in = source.next();
             try (
-                Driver d = new Driver(
-                    "test",
+                Driver d = TestDriverFactory.create(
                     driverContext(),
                     new CannedSourceOperator(in.iterator()),
                     operators.get(),
-                    new PageConsumerOperator(result::add),
-                    () -> {}
+                    new PageConsumerOperator(result::add)
                 )
             ) {
                 runDriver(d);
@@ -275,13 +273,11 @@ protected final List<Page> drive(List<Operator> operators, Iterator<Page> input,
         List<Page> results = new ArrayList<>();
         boolean success = false;
         try (
-            Driver d = new Driver(
-                "test",
+            Driver d = TestDriverFactory.create(
                 driverContext,
                 new CannedSourceOperator(input),
                 operators,
-                new TestResultPageSinkOperator(results::add),
-                () -> {}
+                new TestResultPageSinkOperator(results::add)
             )
         ) {
             runDriver(d);
@@ -303,22 +299,15 @@ public static void runDriver(List<Driver> drivers) {
         int dummyDrivers = between(0, 10);
         for (int i = 0; i < dummyDrivers; i++) {
             drivers.add(
-                new Driver(
-                    "test",
-                    "dummy-session",
-                    0,
-                    0,
+                TestDriverFactory.create(
                     new DriverContext(BigArrays.NON_RECYCLING_INSTANCE, TestBlockFactory.getNonBreakingInstance()),
-                    () -> "dummy-driver",
                     new SequenceLongBlockSourceOperator(
                         TestBlockFactory.getNonBreakingInstance(),
                         LongStream.range(0, between(1, 100)),
                         between(1, 100)
                     ),
                     List.of(),
-                    new PageConsumerOperator(page -> page.releaseBlocks()),
-                    Driver.DEFAULT_STATUS_INTERVAL,
-                    () -> {}
+                    new PageConsumerOperator(Page::releaseBlocks)
                 )
             );
         }

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

@@ -66,11 +66,13 @@
 import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
 import static org.elasticsearch.xpack.esql.CsvTestUtils.loadCsvSpecValues;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs;
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.clusterHasInferenceEndpoint;
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createInferenceEndpoint;
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoint;
+import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createInferenceEndpoints;
+import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoints;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.classpathResources;
+import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.COMPLETION;
+import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.RERANK;
+import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SEMANTIC_TEXT_FIELD_CAPS;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.cap;
 
 // This test can run very long in serverless configurations
@@ -133,8 +135,8 @@ protected EsqlSpecTestCase(
 
     @Before
     public void setup() throws IOException {
-        if (supportsInferenceTestService() && clusterHasInferenceEndpoint(client()) == false) {
-            createInferenceEndpoint(client());
+        if (supportsInferenceTestService()) {
+            createInferenceEndpoints(adminClient());
         }
 
         if (indexExists(availableDatasetsForEs(client(), supportsIndexModeLookup()).iterator().next().indexName()) == false) {
@@ -157,7 +159,8 @@ public static void wipeTestData() throws IOException {
             }
         }
 
-        deleteInferenceEndpoint(client());
+        deleteInferenceEndpoints(adminClient());
+
     }
 
     public boolean logResults() {
@@ -174,8 +177,8 @@ public final void test() throws Throwable {
     }
 
     protected void shouldSkipTest(String testName) throws IOException {
-        if (testCase.requiredCapabilities.contains("semantic_text_field_caps") || testCase.requiredCapabilities.contains("rerank")) {
-            assumeTrue("Inference test service needs to be supported for semantic_text", supportsInferenceTestService());
+        if (requiresInferenceEndpoint()) {
+            assumeTrue("Inference test service needs to be supported", supportsInferenceTestService());
         }
         checkCapabilities(adminClient(), testFeatureService, testName, testCase);
         assumeTrue("Test " + testName + " is not enabled", isEnabled(testName, instructions, Version.CURRENT));
@@ -245,6 +248,11 @@ protected boolean supportsInferenceTestService() {
         return true;
     }
 
+    protected boolean requiresInferenceEndpoint() {
+        return Stream.of(SEMANTIC_TEXT_FIELD_CAPS.capabilityName(), RERANK.capabilityName(), COMPLETION.capabilityName())
+            .anyMatch(testCase.requiredCapabilities::contains);
+    }
+
     protected boolean supportsIndexModeLookup() throws IOException {
         return true;
     }
@@ -340,6 +348,11 @@ private Object valueMapper(CsvTestUtils.Type type, Object value) {
                 return new BigDecimal(s).round(new MathContext(7, RoundingMode.DOWN)).doubleValue();
             }
         }
+        if (type == CsvTestUtils.Type.TEXT || type == CsvTestUtils.Type.KEYWORD || type == CsvTestUtils.Type.SEMANTIC_TEXT) {
+            if (value instanceof String s) {
+                value = s.replaceAll("\\\\n", "\n");
+            }
+        }
         return value.toString();
     }
 

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

@@ -780,7 +780,7 @@ public void testNamedParamsForIdentifierAndIdentifierPatterns() throws IOExcepti
             );
             error = re.getMessage();
             assertThat(error, containsString("ParsingException"));
-            assertThat(error, containsString("line 1:23: mismatched input '?cmd' expecting {'dissect', 'drop'"));
+            assertThat(error, containsString("line 1:23: mismatched input '?cmd' expecting {'completion', 'dissect'"));
         }
     }
 

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

@@ -0,0 +1,192 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.qa.rest;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.ResponseException;
+import org.elasticsearch.test.rest.ESRestTestCase;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createRerankInferenceEndpoint;
+import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteRerankInferenceEndpoint;
+import static org.hamcrest.core.StringContains.containsString;
+
+public class RestRerankTestCase extends ESRestTestCase {
+
+    @Before
+    public void skipWhenRerankDisabled() throws IOException {
+        assumeTrue(
+            "Requires RERANK capability",
+            EsqlSpecTestCase.hasCapabilities(adminClient(), List.of(EsqlCapabilities.Cap.RERANK.capabilityName()))
+        );
+    }
+
+    @Before
+    @After
+    public void assertRequestBreakerEmpty() throws Exception {
+        EsqlSpecTestCase.assertRequestBreakerEmpty();
+    }
+
+    @Before
+    public void setUpInferenceEndpoint() throws IOException {
+        createRerankInferenceEndpoint(adminClient());
+    }
+
+    @Before
+    public void setUpTestIndex() throws IOException {
+        Request request = new Request("PUT", "/rerank-test-index");
+        request.setJsonEntity("""
+            {
+              "mappings": {
+                "properties": {
+                  "title": { "type": "text" },
+                  "author": { "type": "text" }
+                }
+              }
+            }""");
+        assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
+
+        request = new Request("POST", "/rerank-test-index/_bulk");
+        request.addParameter("refresh", "true");
+        request.setJsonEntity("""
+            { "index": {"_id": 1} }
+            { "title": "The Future of Exploration", "author": "John Doe" }
+            { "index": {"_id": 2} }
+            { "title": "Deep Sea Exploration", "author": "Jane Smith" }
+            { "index": {"_id": 3} }
+            { "title": "History of Space Exploration", "author": "Alice Johnson" }
+            """);
+        assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
+    }
+
+    @After
+    public void wipeData() throws IOException {
+        try {
+            adminClient().performRequest(new Request("DELETE", "/rerank-test-index"));
+        } catch (ResponseException e) {
+            // 404 here just means we had no indexes
+            if (e.getResponse().getStatusLine().getStatusCode() != 404) {
+                throw e;
+            }
+        }
+
+        deleteRerankInferenceEndpoint(adminClient());
+    }
+
+    public void testRerankWithSingleField() throws IOException {
+        String query = """
+            FROM rerank-test-index
+            | WHERE match(title, "exploration")
+            | RERANK "exploration" ON title WITH test_reranker
+            | EVAL _score = ROUND(_score, 5)
+            """;
+
+        Map<String, Object> result = runEsqlQuery(query);
+
+        var expectedValues = List.of(
+            List.of("Jane Smith", "Deep Sea Exploration", 0.02941d),
+            List.of("John Doe", "The Future of Exploration", 0.02632d),
+            List.of("Alice Johnson", "History of Space Exploration", 0.02381d)
+        );
+
+        assertResultMap(result, defaultOutputColumns(), expectedValues);
+    }
+
+    public void testRerankWithMultipleFields() throws IOException {
+        String query = """
+            FROM rerank-test-index
+            | WHERE match(title, "exploration")
+            | RERANK "exploration" ON title, author WITH test_reranker
+            | EVAL _score = ROUND(_score, 5)
+            """;
+
+        Map<String, Object> result = runEsqlQuery(query);
+        ;
+        var expectedValues = List.of(
+            List.of("Jane Smith", "Deep Sea Exploration", 0.01818d),
+            List.of("John Doe", "The Future of Exploration", 0.01754d),
+            List.of("Alice Johnson", "History of Space Exploration", 0.01515d)
+        );
+
+        assertResultMap(result, defaultOutputColumns(), expectedValues);
+    }
+
+    public void testRerankWithPositionalParams() throws IOException {
+        String query = """
+            FROM rerank-test-index
+            | WHERE match(title, "exploration")
+            | RERANK ? ON title WITH ?
+            | EVAL _score = ROUND(_score, 5)
+            """;
+
+        Map<String, Object> result = runEsqlQuery(query, "[\"exploration\", \"test_reranker\"]");
+
+        var expectedValues = List.of(
+            List.of("Jane Smith", "Deep Sea Exploration", 0.02941d),
+            List.of("John Doe", "The Future of Exploration", 0.02632d),
+            List.of("Alice Johnson", "History of Space Exploration", 0.02381d)
+        );
+
+        assertResultMap(result, defaultOutputColumns(), expectedValues);
+    }
+
+    public void testRerankWithNamedParams() throws IOException {
+        String query = """
+            FROM rerank-test-index
+            | WHERE match(title, ?queryText)
+            | RERANK ?queryText ON title WITH ?inferenceId
+            | EVAL _score = ROUND(_score, 5)
+            """;
+
+        Map<String, Object> result = runEsqlQuery(query, "[{\"queryText\": \"exploration\"}, {\"inferenceId\": \"test_reranker\"}]");
+
+        var expectedValues = List.of(
+            List.of("Jane Smith", "Deep Sea Exploration", 0.02941d),
+            List.of("John Doe", "The Future of Exploration", 0.02632d),
+            List.of("Alice Johnson", "History of Space Exploration", 0.02381d)
+        );
+
+        assertResultMap(result, defaultOutputColumns(), expectedValues);
+    }
+
+    public void testRerankWithMissingInferenceId() {
+        String query = """
+            FROM rerank-test-index
+            | WHERE match(title, "exploration")
+            | RERANK "exploration" ON title WITH test_missing
+            | EVAL _score = ROUND(_score, 5)
+            """;
+
+        ResponseException re = expectThrows(ResponseException.class, () -> runEsqlQuery(query));
+        assertThat(re.getMessage(), containsString("Inference endpoint not found"));
+    }
+
+    private static List<Map<String, String>> defaultOutputColumns() {
+        return List.of(
+            Map.of("name", "author", "type", "text"),
+            Map.of("name", "title", "type", "text"),
+            Map.of("name", "_score", "type", "double")
+        );
+    }
+
+    private Map<String, Object> runEsqlQuery(String query) throws IOException {
+        RestEsqlTestCase.RequestObjectBuilder builder = RestEsqlTestCase.requestObjectBuilder().query(query);
+        return RestEsqlTestCase.runEsqlSync(builder);
+    }
+
+    private Map<String, Object> runEsqlQuery(String query, String params) throws IOException {
+        RestEsqlTestCase.RequestObjectBuilder builder = RestEsqlTestCase.requestObjectBuilder().query(query).params(params);
+        return RestEsqlTestCase.runEsqlSync(builder);
+    }
+}