[ES|QL] Implement new syntax for RERANK and COMPLETION (elastic#129716)

Author: afoucret

docs/reference/query-languages/esql/_snippets/commands/layout/completion.md

@@ -9,10 +9,26 @@ The `COMPLETION` command allows you to send prompts and context to a Large Langu
 
 **Syntax**
 
+::::{tab-set}
+
+:::{tab-item} 9.2.0+
+
 ```esql
-COMPLETION [column =] prompt WITH inference_id
+COMPLETION [column =] prompt WITH { "inference_id" : "my_inference_endpoint" }
 ```
 
+:::
+
+:::{tab-item} 9.1.x only
+
+```esql
+COMPLETION [column =] prompt WITH my_inference_endpoint
+```
+
+:::
+
+::::
+
 **Parameters**
 
 `column`
@@ -24,7 +40,7 @@ COMPLETION [column =] prompt WITH inference_id
 :   The input text or expression used to prompt the LLM.
     This can be a string literal or a reference to a column containing text.
 
-`inference_id`
+`my_inference_endpoint`
 :   The ID of the [inference endpoint](docs-content://explore-analyze/elastic-inference/inference-api.md) to use for the task.
     The inference endpoint must be configured with the `completion` task type.
 
@@ -75,7 +91,7 @@ How you increase the timeout depends on your deployment type:
 If you don't want to increase the timeout limit, try the following:
 
 * Reduce data volume with `LIMIT` or more selective filters before the `COMPLETION` command
-* Split complex operations into multiple simpler queries 
+* Split complex operations into multiple simpler queries
 * Configure your HTTP client's response timeout (Refer to [HTTP client configuration](/reference/elasticsearch/configuration-reference/networking-settings.md#_http_client_configuration))
 
 
@@ -85,7 +101,7 @@ Use the default column name (results stored in `completion` column):
 
 ```esql
 ROW question = "What is Elasticsearch?"
-| COMPLETION question WITH test_completion_model
+| COMPLETION question WITH { "inference_id" : "my_inference_endpoint" }
 | KEEP question, completion
 ```
 
@@ -97,7 +113,7 @@ Specify the output column (results stored in `answer` column):
 
 ```esql
 ROW question = "What is Elasticsearch?"
-| COMPLETION answer = question WITH test_completion_model
+| COMPLETION answer = question WITH { "inference_id" : "my_inference_endpoint" }
 | KEEP question, answer
 ```
 
@@ -117,7 +133,7 @@ FROM movies
    "Synopsis: ", synopsis, "\n",
    "Actors: ", MV_CONCAT(actors, ", "), "\n",
   )
-| COMPLETION summary = prompt WITH test_completion_model
+| COMPLETION summary = prompt WITH { "inference_id" : "my_inference_endpoint" }
 | KEEP title, summary, rating
 ```
 

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/MapExpression.java

@@ -120,10 +120,18 @@ public Expression get(Object key) {
             return map.get(key);
         } else {
             // the key(literal) could be converted to BytesRef by ConvertStringToByteRef
-            return keyFoldedMap.containsKey(key) ? keyFoldedMap.get(key) : keyFoldedMap.get(new BytesRef(key.toString()));
+            return keyFoldedMap.containsKey(key) ? keyFoldedMap.get(key) : keyFoldedMap.get(foldKey(key));
         }
     }
 
+    public boolean containsKey(Object key) {
+        return keyFoldedMap.containsKey(key) || keyFoldedMap.containsKey(foldKey(key));
+    }
+
+    private BytesRef foldKey(Object key) {
+        return new BytesRef(key.toString());
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj) {

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

@@ -93,7 +93,7 @@ public void testRerankWithSingleField() throws IOException {
         String query = """
             FROM rerank-test-index
             | WHERE match(title, "exploration")
-            | RERANK "exploration" ON title WITH test_reranker
+            | RERANK "exploration" ON title WITH { "inference_id" : "test_reranker" }
             | EVAL _score = ROUND(_score, 5)
             """;
 
@@ -112,7 +112,7 @@ public void testRerankWithMultipleFields() throws IOException {
         String query = """
             FROM rerank-test-index
             | WHERE match(title, "exploration")
-            | RERANK "exploration" ON title, author WITH test_reranker
+            | RERANK "exploration" ON title, author WITH { "inference_id" : "test_reranker" }
             | EVAL _score = ROUND(_score, 5)
             """;
 
@@ -131,7 +131,7 @@ public void testRerankWithPositionalParams() throws IOException {
         String query = """
             FROM rerank-test-index
             | WHERE match(title, "exploration")
-            | RERANK ? ON title WITH ?
+            | RERANK ? ON title WITH { "inference_id" : ? }
             | EVAL _score = ROUND(_score, 5)
             """;
 
@@ -150,7 +150,7 @@ public void testRerankWithNamedParams() throws IOException {
         String query = """
             FROM rerank-test-index
             | WHERE match(title, ?queryText)
-            | RERANK ?queryText ON title WITH ?inferenceId
+            | RERANK ?queryText ON title WITH { "inference_id" : ?inferenceId }
             | EVAL _score = ROUND(_score, 5)
             """;
 
@@ -169,7 +169,7 @@ public void testRerankWithMissingInferenceId() {
         String query = """
             FROM rerank-test-index
             | WHERE match(title, "exploration")
-            | RERANK "exploration" ON title WITH test_missing
+            | RERANK "exploration" ON title WITH { "inference_id" : "test_missing" }
             | EVAL _score = ROUND(_score, 5)
             """;
 

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

@@ -6,7 +6,7 @@ completion using a ROW source operator
 required_capability: completion
 
 ROW prompt="Who is Victor Hugo?"
-| COMPLETION completion_output = prompt WITH test_completion
+| COMPLETION completion_output = prompt WITH { "inference_id" : "test_completion" }
 ;
 
 prompt:keyword      | completion_output:keyword
@@ -18,7 +18,7 @@ completion using a ROW source operator and prompt is a multi-valued field
 required_capability: completion
 
 ROW prompt=["Answer the following question:", "Who is Victor Hugo?"]
-| COMPLETION completion_output = prompt WITH test_completion
+| COMPLETION completion_output = prompt WITH { "inference_id" : "test_completion" }
 ;
 
 prompt:keyword                                        | completion_output:keyword
@@ -34,7 +34,7 @@ FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
 | SORT _score DESC
 | LIMIT 2
-| COMPLETION title WITH test_completion
+| COMPLETION title WITH { "inference_id" : "test_completion" }
 | KEEP title, completion
 ;
 
@@ -51,7 +51,7 @@ FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
 | SORT _score DESC
 | LIMIT 2
-| COMPLETION CONCAT("This is a prompt: ", title) WITH test_completion
+| COMPLETION CONCAT("This is a prompt: ", title) WITH { "inference_id" : "test_completion" }
 | KEEP title, completion
 ;
 

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

@@ -809,7 +809,7 @@ FROM employees
 | KEEP emp_no, first_name, last_name
 | FORK (WHERE emp_no == 10048 OR emp_no == 10081)
        (WHERE emp_no == 10081 OR emp_no == 10087)
-| COMPLETION x = CONCAT(first_name, " ", last_name) WITH test_completion
+| COMPLETION x=CONCAT(first_name, " ", last_name) WITH { "inference_id" : "test_completion" }
 | SORT _fork, emp_no
 ;
 
@@ -827,7 +827,7 @@ required_capability: completion
 FROM employees
 | KEEP emp_no, first_name, last_name
 | FORK (WHERE emp_no == 10048 OR emp_no == 10081
-       | COMPLETION x = CONCAT(first_name, " ", last_name) WITH test_completion)
+       | COMPLETION x=CONCAT(first_name, " ", last_name) WITH { "inference_id" : "test_completion" })
        (WHERE emp_no == 10081 OR emp_no == 10087)
 | SORT _fork, emp_no
 ;
@@ -845,7 +845,7 @@ required_capability: completion
 
 FROM employees
 | KEEP emp_no, first_name, last_name
-| COMPLETION x = CONCAT(first_name, " ", last_name) WITH test_completion
+| COMPLETION x=CONCAT(first_name, " ", last_name) WITH { "inference_id" : "test_completion" }
 | FORK (WHERE emp_no == 10048 OR emp_no == 10081)
        (WHERE emp_no == 10081 OR emp_no == 10087)
 | SORT _fork, emp_no

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

@@ -10,7 +10,7 @@ required_capability: match_operator_colon
 FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
 | SORT _score DESC, book_no ASC
-| RERANK "war and peace" ON title WITH inferenceId=test_reranker
+| RERANK "war and peace" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2)
 | KEEP book_no, title, author, _score
 ;
@@ -29,7 +29,7 @@ required_capability: match_operator_colon
 FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
 | SORT _score DESC, book_no ASC
-| RERANK "war and peace" ON title WITH inferenceId=test_reranker, scoreColumn=rerank_score
+| RERANK rerank_score="war and peace" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2), rerank_score=ROUND(rerank_score, 2)
 | KEEP book_no, title, author, rerank_score
 ;
@@ -48,7 +48,7 @@ required_capability: match_operator_colon
 FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
 | SORT _score DESC
-| RERANK "war and peace" ON title WITH inferenceId=test_reranker, scoreColumn=rerank_score
+| RERANK rerank_score="war and peace" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2), rerank_score=ROUND(rerank_score, 2)
 | SORT rerank_score, _score ASC, book_no ASC
 | KEEP book_no, title, author, rerank_score
@@ -68,7 +68,7 @@ required_capability: match_operator_colon
 
 FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
-| RERANK "war and peace" ON title, author WITH inferenceId=test_reranker
+| RERANK "war and peace" ON title, author WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2)
 | SORT _score DESC, book_no ASC
 | KEEP book_no, title, author, _score
@@ -90,7 +90,7 @@ FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
 | SORT _score DESC, book_no ASC
 | LIMIT 3
-| RERANK "war and peace" ON title WITH inferenceId=test_reranker
+| RERANK "war and peace" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2)
 | SORT _score DESC, book_no ASC
 | KEEP book_no, title, author, _score
@@ -109,7 +109,7 @@ required_capability: match_operator_colon
 
 FROM books METADATA _score
 | WHERE title:"war and peace" AND author:"Tolstoy"
-| RERANK "war and peace" ON title WITH inferenceId=test_reranker
+| RERANK "war and peace" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2)
 | SORT _score DESC, book_no ASC
 | KEEP book_no, title, author, _score
@@ -129,7 +129,7 @@ required_capability: match_operator_colon
 
 FROM books
 | WHERE title:"war and peace" AND author:"Tolstoy"
-| RERANK "war and peace" ON title WITH inferenceId=test_reranker
+| RERANK "war and peace" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2)
 | KEEP book_no, title, author, _score
 | SORT author, title 
@@ -153,7 +153,7 @@ FROM books METADATA _id, _index, _score
 | FORK ( WHERE title:"Tolkien" | SORT _score, _id DESC | LIMIT 3 )
        ( WHERE author:"Tolkien" | SORT _score, _id DESC | LIMIT 3 )
 | FUSE
-| RERANK "Tolkien" ON title WITH inferenceId=test_reranker
+| RERANK "Tolkien" ON title WITH { "inference_id" : "test_reranker" }
 | EVAL _score=ROUND(_score, 2)
 | SORT _score DESC, book_no ASC
 | LIMIT 2

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

@@ -219,23 +219,28 @@ renameClause:
     ;
 
 dissectCommand
-    : DISSECT primaryExpression string commandOptions?
+    : DISSECT primaryExpression string dissectCommandOptions?
     ;
 
-grokCommand
-    : GROK primaryExpression string
+dissectCommandOptions
+    : dissectCommandOption (COMMA dissectCommandOption)*
     ;
 
-mvExpandCommand
-    : MV_EXPAND qualifiedName
+dissectCommandOption
+    : identifier ASSIGN constant
     ;
 
-commandOptions
-    : commandOption (COMMA commandOption)*
+
+commandNamedParameters
+    : (WITH mapExpression)?
     ;
 
-commandOption
-    : identifier ASSIGN constant
+grokCommand
+    : GROK primaryExpression string
+    ;
+
+mvExpandCommand
+    : MV_EXPAND qualifiedName
     ;
 
 explainCommand
@@ -293,7 +298,7 @@ forkSubQueryProcessingCommand
     ;
 
 completionCommand
-    : COMPLETION (targetField=qualifiedName ASSIGN)? prompt=primaryExpression WITH inferenceId=identifierOrParameter
+    : COMPLETION (targetField=qualifiedName ASSIGN)? prompt=primaryExpression commandNamedParameters
     ;
 
 //
@@ -315,19 +320,6 @@ fuseCommand
     : DEV_FUSE
     ;
 
-inferenceCommandOptions
-    : inferenceCommandOption (COMMA inferenceCommandOption)*
-    ;
-
-inferenceCommandOption
-    : identifier ASSIGN inferenceCommandOptionValue
-    ;
-
-inferenceCommandOptionValue
-    : constant
-    | identifier
-    ;
-
 rerankCommand
-    : DEV_RERANK queryText=constant ON rerankFields (WITH inferenceCommandOptions)?
+    : DEV_RERANK (targetField=qualifiedName ASSIGN)? queryText=constant ON rerankFields commandNamedParameters
     ;

x-pack/plugin/esql/src/main/antlr/parser/Expression.g4

@@ -57,7 +57,7 @@ functionName
     ;
 
 mapExpression
-    : LEFT_BRACES entryExpression (COMMA entryExpression)* RIGHT_BRACES
+    : LEFT_BRACES (entryExpression (COMMA entryExpression)*)? RIGHT_BRACES
     ;
 
 entryExpression

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -1405,7 +1405,7 @@ private BitSet gatherPreAnalysisMetrics(LogicalPlan plan, BitSet b) {
     private static class ImplicitCasting extends ParameterizedRule<LogicalPlan, LogicalPlan, AnalyzerContext> {
         @Override
         public LogicalPlan apply(LogicalPlan plan, AnalyzerContext context) {
-            // do implicit casting for function arguments
+            // do implicit casting for named parameters
             return plan.transformExpressionsUp(
                 org.elasticsearch.xpack.esql.core.expression.function.Function.class,
                 e -> ImplicitCasting.cast(e, context.functionRegistry().snapshotRegistry())