ES|QL Reranker command

docs/changelog/123074.yaml

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

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

@@ -62,6 +62,10 @@ public static class Request extends BaseInferenceActionRequest {
         public static final ParseField QUERY = new ParseField("query");
         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/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -52,6 +52,7 @@
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.JOIN_LOOKUP_V12;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.JOIN_PLANNING_V1;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.METADATA_FIELDS_REMOTE_TEST;
+import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.RERANK;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.UNMAPPED_FIELDS;
 import static org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase.Mode.SYNC;
 import static org.mockito.ArgumentMatchers.any;
@@ -130,6 +131,8 @@ protected void shouldSkipTest(String testName) throws IOException {
         assumeFalse("LOOKUP JOIN not yet supported in CCS", testCase.requiredCapabilities.contains(JOIN_LOOKUP_V12.capabilityName()));
         // Unmapped fields require a coorect capability response from every cluster, which isn't currently implemented.
         assumeFalse("UNMAPPED FIELDS not yet supported in CCS", testCase.requiredCapabilities.contains(UNMAPPED_FIELDS.capabilityName()));
+        // Need to do additional developmnet to get CSS support for the rerank coammnd
+        assumeFalse("RERANK not yet supported in CCS", testCase.requiredCapabilities.contains(RERANK.capabilityName()));
     }
 
     private TestFeatureService remoteFeaturesService() throws IOException {

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

@@ -66,7 +66,7 @@
 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.createInferenceEndpoints;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoint;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.classpathResources;
@@ -130,7 +130,7 @@ protected EsqlSpecTestCase(
     @Before
     public void setup() throws IOException {
         if (supportsInferenceTestService() && clusterHasInferenceEndpoint(client()) == false) {
-            createInferenceEndpoint(client());
+            createInferenceEndpoints(client());
         }
 
         boolean supportsLookup = supportsIndexModeLookup();

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java

@@ -357,7 +357,7 @@ private static void loadDataSetIntoEs(
     }
 
     /** The semantic_text mapping type require an inference endpoint that needs to be setup before creating the index. */
-    public static void createInferenceEndpoint(RestClient client) throws IOException {
+    public static void createInferenceEndpoints(RestClient client) throws IOException {
         Request request = new Request("PUT", "_inference/sparse_embedding/test_sparse_inference");
         request.setJsonEntity("""
                   {
@@ -371,6 +371,21 @@ public static void createInferenceEndpoint(RestClient client) throws IOException
                  }
             """);
         client.performRequest(request);
+
+        request = new Request("PUT", "_inference/rerank/test_reranker");
+        request.setJsonEntity("""
+            {
+                "service": "test_reranking_service",
+                "service_settings": {
+                    "model_id": "my_model",
+                    "api_key": "abc64"
+                },
+                "task_settings": {
+                    "use_text_length": true
+                }
+            }
+            """);
+        client.performRequest(request);
     }
 
     public static void deleteInferenceEndpoint(RestClient client) throws IOException {

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -11,6 +11,7 @@
 import org.apache.lucene.sandbox.document.HalfFloatPoint;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
@@ -67,6 +68,8 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
 import org.elasticsearch.xpack.esql.index.EsIndex;
+import org.elasticsearch.xpack.esql.inference.InferenceResolution;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.parser.QueryParam;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -148,6 +151,8 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 
 public final class EsqlTestUtils {
@@ -375,9 +380,21 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         null,
         mock(ClusterService.class),
         mock(IndexNameExpressionResolver.class),
-        null
+        null,
+        mockInferenceService()
     );
 
+    @SuppressWarnings("unchecked")
+    private static InferenceService mockInferenceService() {
+        InferenceService inferenceService = mock(InferenceService.class);
+        doAnswer(i -> {
+            i.getArgument(1, ActionListener.class).onResponse(emptyInferenceResolution());
+            return null;
+        }).when(inferenceService).resolveInferences(any(), any());
+
+        return inferenceService;
+    }
+
     private EsqlTestUtils() {}
 
     public static Configuration configuration(QueryPragmas pragmas, String query) {
@@ -453,6 +470,10 @@ public static EnrichResolution emptyPolicyResolution() {
         return new EnrichResolution();
     }
 
+    public static InferenceResolution emptyInferenceResolution() {
+        return InferenceResolution.EMPTY;
+    }
+
     public static SearchStats statsForExistingField(String... names) {
         return fieldMatchingExistOrMissing(true, names);
     }

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

@@ -0,0 +1,95 @@
+// Note:
+// The "test_reranker" service scores the row from the inputText length and does not really score by relevance.
+// This makes the output more predictable which is helpful here.
+
+
+reranker using a single field
+required_capability: rerank
+required_capability: match_function
+
+FROM books METADATA _score
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| RERANK "war and peace" ON title WITH "test_reranker"
+| KEEP book_no, title, author, _score
+;
+
+book_no:keyword | title:text                                            | author:text                                 | _score:double
+5327            | War and Peace                                         | Leo Tolstoy                                 | 0.03846153989434242
+4536            | War and Peace (Signet Classics)                       | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.02222222276031971
+9032            | War and Peace: A Novel (6 Volumes)                    | Tolstoy Leo                                 | 0.02083333395421505
+2776            | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy                                 | 0.01515151560306549
+;
+
+
+reranker using multiple fields
+required_capability: rerank
+required_capability: match_function
+
+FROM books METADATA _score
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| RERANK "war and peace" ON title, author WITH "test_reranker"
+| KEEP book_no, title, author, _score
+;
+
+book_no:keyword | title:text                                            | author:text                                 | _score:double
+5327            | War and Peace                                         | Leo Tolstoy                                 | 0.02083333395421505
+9032            | War and Peace: A Novel (6 Volumes)                    | Tolstoy Leo                                 | 0.014285714365541935
+2776            | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy                                 | 0.011363636702299118
+4536            | War and Peace (Signet Classics)                       | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.009523809887468815
+;
+
+
+reranker after a limit
+required_capability: rerank
+required_capability: match_function
+
+FROM books METADATA _score
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| SORT _score DESC
+| LIMIT 3
+| RERANK "war and peace" ON title WITH "test_reranker"
+| KEEP book_no, title, author, _score
+;
+
+book_no:keyword | title:text                         | author:text                                 | _score:double
+5327            | War and Peace                      | Leo Tolstoy                                 | 0.03846153989434242
+4536            | War and Peace (Signet Classics)    | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.02222222276031971
+9032            | War and Peace: A Novel (6 Volumes) | Tolstoy Leo                                 | 0.02083333395421505
+;
+
+
+reranker before a limit
+required_capability: rerank
+required_capability: match_function
+
+FROM books METADATA _score
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| RERANK "war and peace" ON title WITH "test_reranker"
+| KEEP book_no, title, author, _score
+| LIMIT 3
+;
+
+book_no:keyword | title:text                         | author:text                                 | _score:double
+5327            | War and Peace                      | Leo Tolstoy                                 | 0.03846153989434242
+4536            | War and Peace (Signet Classics)    | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.02222222276031971
+9032            | War and Peace: A Novel (6 Volumes) | Tolstoy Leo                                 | 0.02083333395421505
+;
+
+
+reranker add the _score column when missing
+required_capability: rerank
+required_capability: match_function
+
+FROM books
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| RERANK "war and peace" ON title WITH "test_reranker"
+| KEEP book_no, title, author, _score
+;
+
+
+book_no:keyword | title:text                                            | author:text                                 | _score:double
+5327            | War and Peace                                         | Leo Tolstoy                                 | 0.03846153989434242
+4536            | War and Peace (Signet Classics)                       | [John Hockenberry, Leo Tolstoy, Pat Conroy] | 0.02222222276031971
+9032            | War and Peace: A Novel (6 Volumes)                    | Tolstoy Leo                                 | 0.02083333395421505
+2776            | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy                                 | 0.01515151560306549
+;