[ES|QL] Rerank operator improvements (#132318)

Author: afoucret

x-pack/plugin/esql/qa/testFixtures/src/main/resources/data/books.csv

@@ -25,6 +25,9 @@
     },
     "year": {
       "type": "integer"
+    },
+    "collection": {
+      "type": "text"
     }
   }
 }

x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-books.json

@@ -61,6 +61,45 @@ book_no:keyword | title:text                                            | author
 5327            | War and Peace                                         | Leo Tolstoy                                 | 0.08         
 ;
 
+reranker using a non text fields
+required_capability: rerank
+required_capability: match_operator_colon
+
+FROM books METADATA _score
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| RERANK "war and peace" ON ratings WITH { "inference_id" : "test_reranker" }
+| EVAL _score=ROUND(_score, 2), ratings = ROUND(ratings, 2)
+| SORT _score DESC, book_no ASC
+| KEEP book_no, title, ratings, _score
+;
+
+book_no:keyword | title:text                                            | ratings:double | _score:double
+2776            | The Devil and Other Stories (Oxford World's Classics) | 5.0            | 0.33
+4536            | War and Peace (Signet Classics)                       | 4.75           | 0.25
+5327            | War and Peace                                         | 3.84           | 0.06
+9032            | War and Peace: A Novel (6 Volumes)                    | 3.81           | 0.06
+;
+
+
+reranker using a sparse input field
+required_capability: rerank
+required_capability: match_operator_colon
+
+FROM books METADATA _score
+| WHERE MATCH(title, "lord of the rings", {"minimum_should_match": "100%"}) AND author:"tolkien"
+| RERANK rerank_score="war and peace" ON collection WITH { "inference_id" : "test_reranker" }
+| EVAL _score=ROUND(_score, 2), rerank_score=ROUND(rerank_score, 2)
+| SORT rerank_score DESC NULLS LAST, _score  DESC
+| LIMIT 3
+| KEEP book_no, title, author, collection, rerank_score, _score
+;
+
+book_no:keyword | title:text                                                                 | author:text                  | collection:text       | rerank_score:double | _score:double
+2714            | Return of the King Being the Third Part of The Lord of the Rings           | J. R. R. Tolkien             | The Lord of the Rings | 0.05                | 8.56
+2675            | The Lord of the Rings - Boxed Set                                          | J.R.R. Tolkien               | The Lord of the Rings | 0.05                | 8.35
+7140            | The Lord of the Rings Poster Collection: Six Paintings by Alan Lee (No. 1) | [Alan Lee, J. R. R. Tolkien] | null                  | null                | 5.5
+;
+
 
 reranker using multiple fields
 required_capability: rerank
@@ -82,6 +121,26 @@ book_no:keyword | title:text                                            | author
 ;
 
 
+reranker using multiple fields with some non text fields
+required_capability: rerank
+required_capability: match_operator_colon
+
+FROM books METADATA _score
+| WHERE title:"war and peace" AND author:"Tolstoy"
+| RERANK "war and peace" ON title, ratings WITH { "inference_id" : "test_reranker" }
+| EVAL _score=ROUND(_score, 2), ratings = ROUND(ratings, 2)
+| SORT _score DESC, book_no ASC
+| KEEP book_no, title, ratings, _score
+;
+
+book_no:keyword | title:text                                            | ratings:double | _score:double
+4536            | War and Peace (Signet Classics)                       | 4.75           | 0.02
+5327            | War and Peace                                         | 3.84           | 0.02
+2776            | The Devil and Other Stories (Oxford World's Classics) | 5.0            | 0.01
+9032            | War and Peace: A Novel (6 Volumes)                    | 3.81           | 0.01     
+;
+
+
 reranker after a limit
 required_capability: rerank
 required_capability: match_operator_colon

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

@@ -351,8 +351,8 @@ from books metadata _score
 
 book_no:keyword | c_score:double
 8875            | 3.0
+7480            | 1.0
 7350            | 1.0
-7140            | 3.0
 ;
 
 QstrScoreManipulation

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

@@ -79,6 +79,7 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDouble;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToInteger;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToString;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToUnsignedLong;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.expression.function.vector.VectorFunction;
@@ -813,11 +814,23 @@ private LogicalPlan resolveRerank(Rerank rerank, List<Attribute> childrenOutput)
             List<Alias> newFields = new ArrayList<>();
             boolean changed = false;
 
+            // Do not need to cast as string if there are multiple rerank fields since it will be converted to YAML.
+            boolean castRerankFieldsAsString = rerank.rerankFields().size() < 2;
+
             // First resolving fields used in expression
             for (Alias field : rerank.rerankFields()) {
-                Alias result = (Alias) field.transformUp(UnresolvedAttribute.class, ua -> resolveAttribute(ua, childrenOutput));
-                newFields.add(result);
-                changed |= result != field;
+                Alias resolved = (Alias) field.transformUp(UnresolvedAttribute.class, ua -> resolveAttribute(ua, childrenOutput));
+
+                if (resolved.resolved()) {
+                    if (castRerankFieldsAsString
+                        && rerank.isValidRerankField(resolved)
+                        && DataType.isString(resolved.dataType()) == false) {
+                        resolved = resolved.replaceChild(new ToString(resolved.child().source(), resolved.child()));
+                    }
+                }
+
+                newFields.add(resolved);
+                changed |= resolved != field;
             }
 
             if (changed) {

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

@@ -77,10 +77,10 @@ public void close() {
     @Override
     public BytesRefBlock eval(Page page) {
         Block[] fieldValueBlocks = new Block[fieldsValueEvaluators.length];
+
         try (
             BytesRefStreamOutput outputStream = new BytesRefStreamOutput();
-            XContentBuilder xContentBuilder = XContentFactory.contentBuilder(xContentType, outputStream);
-            BytesRefBlock.Builder outputBlockBuilder = blockFactory.newBytesRefBlockBuilder(page.getPositionCount());
+            BytesRefBlock.Builder outputBlockBuilder = blockFactory.newBytesRefBlockBuilder(page.getPositionCount())
         ) {
 
             PositionToXContent[] toXContents = new PositionToXContent[fieldsValueEvaluators.length];
@@ -90,18 +90,28 @@ public BytesRefBlock eval(Page page) {
             }
 
             for (int pos = 0; pos < page.getPositionCount(); pos++) {
-                xContentBuilder.startObject();
-                for (int i = 0; i < fieldValueBlocks.length; i++) {
-                    String fieldName = columnsInfo[i].name();
-                    Block currentBlock = fieldValueBlocks[i];
-                    if (currentBlock.isNull(pos) || currentBlock.getValueCount(pos) < 1) {
-                        continue;
+                try (XContentBuilder xContentBuilder = XContentFactory.contentBuilder(xContentType, outputStream)) {
+
+                    xContentBuilder.startObject();
+                    boolean hasNullsOnly = true;
+                    for (int i = 0; i < fieldValueBlocks.length; i++) {
+                        String fieldName = columnsInfo[i].name();
+                        Block currentBlock = fieldValueBlocks[i];
+                        if (currentBlock.isNull(pos) || currentBlock.getValueCount(pos) < 1) {
+                            continue;
+                        }
+                        hasNullsOnly = false;
+                        toXContents[i].positionToXContent(xContentBuilder.field(fieldName), ToXContent.EMPTY_PARAMS, pos);
+                    }
+                    xContentBuilder.endObject().flush();
+
+                    if (hasNullsOnly) {
+                        outputBlockBuilder.appendNull();
+                    } else {
+                        outputBlockBuilder.appendBytesRef(outputStream.get());
+                        outputStream.reset();
                     }
-                    toXContents[i].positionToXContent(xContentBuilder.field(fieldName), ToXContent.EMPTY_PARAMS, pos);
                 }
-                xContentBuilder.endObject().flush();
-                outputBlockBuilder.appendBytesRef(outputStream.get());
-                outputStream.reset();
             }
 
             return outputBlockBuilder.build();
@@ -125,7 +135,7 @@ public static class Factory implements ExpressionEvaluator.Factory {
         private final XContentType xContentType;
         private final Map<ColumnInfoImpl, ExpressionEvaluator.Factory> fieldsEvaluatorFactories;
 
-        private Factory(XContentType xContentType, Map<ColumnInfoImpl, ExpressionEvaluator.Factory> fieldsEvaluatorFactories) {
+        Factory(XContentType xContentType, Map<ColumnInfoImpl, ExpressionEvaluator.Factory> fieldsEvaluatorFactories) {
             this.xContentType = xContentType;
             this.fieldsEvaluatorFactories = fieldsEvaluatorFactories;
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/XContentRowEncoder.java

@@ -51,6 +51,7 @@ public InferenceAction.Request next() {
         if (hasNext() == false) {
             throw new NoSuchElementException();
         }
+
         return inferenceRequest(promptReader.readPrompt(currentPos++));
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/completion/CompletionOperatorRequestIterator.java

@@ -80,10 +80,16 @@ public Page buildOutput() {
      */
     @Override
     public void addInferenceResponse(InferenceAction.Response inferenceResponse) {
+        if (inferenceResponse == null) {
+            scoreBlockBuilder.appendNull();
+            return;
+        }
+
         Iterator<RankedDocsResults.RankedDoc> sortedRankedDocIterator = inferenceResults(inferenceResponse).getRankedDocs()
             .stream()
             .sorted(Comparator.comparingInt(RankedDocsResults.RankedDoc::index))
             .iterator();
+
         while (sortedRankedDocIterator.hasNext()) {
             scoreBlockBuilder.appendDouble(sortedRankedDocIterator.next().relevanceScore());
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/rerank/RerankOperatorOutputBuilder.java

@@ -51,22 +51,28 @@ public InferenceAction.Request next() {
             throw new NoSuchElementException();
         }
 
-        final int inputSize = Math.min(remainingPositions, batchSize);
-        final List<String> inputs = new ArrayList<>(inputSize);
+        final int maxInputSize = Math.min(remainingPositions, batchSize);
+        final List<String> inputs = new ArrayList<>(maxInputSize);
         BytesRef scratch = new BytesRef();
 
         int startIndex = inputBlock.getPositionCount() - remainingPositions;
-        for (int i = 0; i < inputSize; i++) {
+
+        if (inputBlock.isNull(startIndex)) {
+            remainingPositions -= 1;
+            return null;
+        }
+
+        for (int i = 0; i < maxInputSize; i++) {
             int pos = startIndex + i;
             if (inputBlock.isNull(pos)) {
-                inputs.add("");
+                break;
             } else {
                 scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(pos), scratch);
                 inputs.add(BytesRefs.toString(scratch));
             }
         }
 
-        remainingPositions -= inputSize;
+        remainingPositions -= inputs.size();
         return inferenceRequest(inputs);
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/rerank/RerankOperatorRequestIterator.java

@@ -758,18 +758,6 @@ public PlanFactory visitRerankCommand(EsqlBaseParser.RerankCommandContext ctx) {
         Expression queryText = expression(ctx.queryText);
         Attribute scoreAttribute = visitQualifiedName(ctx.targetField, new UnresolvedAttribute(source, MetadataAttribute.SCORE));
 
-        if (queryText instanceof Literal queryTextLiteral && DataType.isString(queryText.dataType())) {
-            if (queryTextLiteral.value() == null) {
-                throw new ParsingException(source(ctx.queryText), "Query cannot be null or undefined in RERANK", ctx.queryText.getText());
-            }
-        } else {
-            throw new ParsingException(
-                source(ctx.queryText),
-                "Query must be a valid string in RERANK, found [{}]",
-                ctx.queryText.getText()
-            );
-        }
-
         return p -> {
             checkForRemoteClusters(p, source, "RERANK");
             return applyRerankOptions(new Rerank(source, p, queryText, rerankFields, scoreAttribute), ctx.commandNamedParameters());