[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/compute/test/src/main/java/org/elasticsearch/compute/test/TestDriverFactory.java

@@ -0,0 +1,75 @@
+/*
+ * 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.compute.test;
+
+import org.elasticsearch.common.Randomness;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.LocalCircuitBreaker;
+import org.elasticsearch.compute.operator.Driver;
+import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.compute.operator.SinkOperator;
+import org.elasticsearch.compute.operator.SourceOperator;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.indices.CrankyCircuitBreakerService;
+
+import java.util.List;
+
+public class TestDriverFactory {
+
+    /**
+     * This is a convenience factory to create a test driver instance
+     * that accepts fewer parameters
+     */
+    public static Driver create(
+        DriverContext driverContext,
+        SourceOperator source,
+        List<Operator> intermediateOperators,
+        SinkOperator sink
+    ) {
+        return create(driverContext, source, intermediateOperators, sink, () -> {});
+    }
+
+    public static Driver create(
+        DriverContext driverContext,
+        SourceOperator source,
+        List<Operator> intermediateOperators,
+        SinkOperator sink,
+        Releasable releasable
+    ) {
+        // Do not wrap the local breaker for small local breakers, as the output mights not match expectations.
+        if (driverContext.breaker() instanceof CrankyCircuitBreakerService.CrankyCircuitBreaker == false
+            && driverContext.breaker() instanceof LocalCircuitBreaker == false
+            && driverContext.breaker().getLimit() >= ByteSizeValue.ofMb(100).getBytes()
+            && Randomness.get().nextBoolean()) {
+            final int overReservedBytes = Randomness.get().nextInt(1024 * 1024);
+            final int maxOverReservedBytes = overReservedBytes + Randomness.get().nextInt(1024 * 1024);
+            var localBreaker = new LocalCircuitBreaker(driverContext.breaker(), overReservedBytes, maxOverReservedBytes);
+            BlockFactory localBlockFactory = driverContext.blockFactory().newChildFactory(localBreaker);
+            driverContext = new DriverContext(localBlockFactory.bigArrays(), localBlockFactory);
+        }
+        if (driverContext.breaker() instanceof LocalCircuitBreaker localBreaker) {
+            releasable = Releasables.wrap(releasable, localBreaker);
+        }
+        return new Driver(
+            "unset",
+            "test-task",
+            System.currentTimeMillis(),
+            System.nanoTime(),
+            driverContext,
+            () -> null,
+            source,
+            intermediateOperators,
+            sink,
+            Driver.DEFAULT_STATUS_INTERVAL,
+            releasable
+        );
+    }
+}