Better handling of the YAML encoding.

Author: afoucret

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

@@ -14,10 +14,10 @@ FROM books METADATA _score
 ;
 
 book_no:keyword | title:text                                            | author:text                                 | _score:double
-5327            | War and Peace                                         | Leo Tolstoy                                 | 0.03703703731298447
+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.014925372786819935
+2776            | The Devil and Other Stories (Oxford World's Classics) | Leo Tolstoy                                 | 0.01515151560306549
 ;
 
 
@@ -32,9 +32,9 @@ FROM books METADATA _score
 ;
 
 book_no:keyword | title:text                                            | author:text                                 | _score:double
-5327            | War and Peace                                         | Leo Tolstoy                                 | 0.020408162847161293
+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.01123595517128706
+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
 ;
 
@@ -52,7 +52,7 @@ FROM books METADATA _score
 ;
 
 book_no:keyword | title:text                         | author:text                                 | _score:double
-5327            | War and Peace                      | Leo Tolstoy                                 | 0.03703703731298447
+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
 ;
@@ -70,7 +70,7 @@ FROM books METADATA _score
 ;
 
 book_no:keyword | title:text                         | author:text                                 | _score:double
-5327            | War and Peace                      | Leo Tolstoy                                 | 0.03703703731298447
+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
 ;

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

@@ -9,10 +9,9 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AsyncOperator;
@@ -21,16 +20,14 @@
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.inference.TaskType;
-import org.elasticsearch.xcontent.XContentBuilder;
-import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
 
-import java.io.IOException;
-import java.io.UncheckedIOException;
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.xpack.esql.inference.XContentRowEncoder.yamlRowEncoderFactory;
+
 public class RerankOperator extends AsyncOperator<Page> {
 
     // Move to a setting.
@@ -46,15 +43,15 @@ public record Factory(
 
         @Override
         public String describe() {
-            return "RerankOperator[inference_id="
+            return "RerankOperator[inference_id=["
                 + inferenceId
-                + " query="
+                + "], query=["
                 + queryText
-                + " rerank_fields="
+                + "], rerank_fields="
                 + fieldsEvaluatorFactories.keySet()
-                + " score_channel="
+                + ", score_channel=["
                 + scoreChannel
-                + "]";
+                + "]]";
         }
 
         @Override
@@ -64,56 +61,43 @@ public Operator get(DriverContext driverContext) {
                 inferenceService,
                 inferenceId,
                 queryText,
-                fieldNames(),
-                fieldsEvaluators(driverContext),
+                yamlRowEncoderFactory(fieldsEvaluatorFactories).get(driverContext),
                 scoreChannel
             );
         }
-
-        private String[] fieldNames() {
-            return fieldsEvaluatorFactories.keySet().toArray(String[]::new);
-        }
-
-        private ExpressionEvaluator[] fieldsEvaluators(DriverContext context) {
-            return fieldsEvaluatorFactories.values().stream().map(factory -> factory.get(context)).toArray(ExpressionEvaluator[]::new);
-        }
     }
 
     private final InferenceService inferenceService;
     private final BlockFactory blockFactory;
     private final String inferenceId;
     private final String queryText;
-    private final String[] fieldNames;
-    private final ExpressionEvaluator[] fieldsEvaluators;
+    private final ExpressionEvaluator rowEncoder;
     private final int scoreChannel;
 
     public RerankOperator(
         DriverContext driverContext,
         InferenceService inferenceService,
         String inferenceId,
         String queryText,
-        String[] fieldNames,
-        ExpressionEvaluator[] fieldsEvaluators,
+        ExpressionEvaluator rowEncoder,
         int scoreChannel
     ) {
         super(driverContext, inferenceService.getThreadContext(), MAX_INFERENCE_WORKER);
 
         assert inferenceService.getThreadContext() != null;
-        assert fieldNames.length == fieldsEvaluators.length;
 
         this.blockFactory = driverContext.blockFactory();
         this.inferenceService = inferenceService;
         this.inferenceId = inferenceId;
         this.queryText = queryText;
-        this.fieldNames = fieldNames;
-        this.fieldsEvaluators = fieldsEvaluators;
+        this.rowEncoder = rowEncoder;
         this.scoreChannel = scoreChannel;
     }
 
     @Override
     protected void performAsync(Page inputPage, ActionListener<Page> listener) {
         // Ensure input page blocks are released when the listener is called.
-        final ActionListener<Page> outputListener = ActionListener.runAfter(listener, () -> { inputPage.releaseBlocks(); });
+        final ActionListener<Page> outputListener = ActionListener.runAfter(listener, () -> { releasePageOnAnyThread(inputPage); });
 
         try {
             inferenceService.doInference(
@@ -130,7 +114,7 @@ protected void performAsync(Page inputPage, ActionListener<Page> listener) {
 
     @Override
     protected void doClose() {
-        Releasables.closeExpectNoException(this.fieldsEvaluators);
+        Releasables.closeExpectNoException(rowEncoder);
     }
 
     @Override
@@ -145,15 +129,15 @@ public Page getOutput() {
 
     @Override
     public String toString() {
-        return "RerankOperator[inference_id="
+        return "RerankOperator[inference_id=["
             + inferenceId
-            + " query="
+            + "], query=["
             + queryText
-            + " rerank_fields="
-            + List.of(fieldNames)
-            + " score_channel="
+            + "], row_encoder=["
+            + rowEncoder
+            + "], score_channel=["
             + scoreChannel
-            + "]";
+            + "]]";
     }
 
     private Page buildOutput(Page inputPage, InferenceAction.Response inferenceResponse) {
@@ -212,46 +196,21 @@ private Block buildScoreBlock(Page inputPage, RankedDocsResults rankedDocsResult
     }
 
     private InferenceAction.Request buildInferenceRequest(Page inputPage) {
-        Block[] inputBlocks = new Block[fieldsEvaluators.length];
-
-        try {
-            for (int b = 0; b < inputBlocks.length; b++) {
-                inputBlocks[b] = fieldsEvaluators[b].eval(inputPage);
-            }
-
+        try (BytesRefBlock encodedRowBlock = (BytesRefBlock) rowEncoder.eval(inputPage)) {
+            assert (encodedRowBlock.getPositionCount() == inputPage.getPositionCount());
             String[] inputs = new String[inputPage.getPositionCount()];
+            BytesRef buffer = new BytesRef();
+
             for (int pos = 0; pos < inputPage.getPositionCount(); pos++) {
-                try (XContentBuilder yamlBuilder = XContentFactory.yamlBuilder().startObject()) {
-                    for (int i = 0; i < inputBlocks.length; i++) {
-                        String fieldName = fieldNames[i];
-                        Block currentBlock = inputBlocks[i];
-                        if (currentBlock.isNull(pos)) {
-                            continue;
-                        }
-                        yamlBuilder.field(fieldName, toYaml(BlockUtils.toJavaObject(currentBlock, pos)));
-                    }
-                    inputs[pos] = Strings.toString(yamlBuilder.endObject());
-                } catch (IOException e) {
-                    throw new UncheckedIOException(e);
+                if (encodedRowBlock.isNull(pos)) {
+                    inputs[pos] = "";
+                } else {
+                    buffer = encodedRowBlock.getBytesRef(encodedRowBlock.getFirstValueIndex(pos), buffer);
+                    inputs[pos] = buffer.utf8ToString();
                 }
             }
 
             return InferenceAction.Request.builder(inferenceId, TaskType.RERANK).setInput(List.of(inputs)).setQuery(queryText).build();
-        } finally {
-            Releasables.closeExpectNoException(inputBlocks);
-        }
-    }
-
-    private Object toYaml(Object value) {
-        try {
-            return switch (value) {
-                case BytesRef b -> b.utf8ToString();
-                case List<?> l -> l.stream().map(this::toYaml).toList();
-                default -> value;
-            };
-        } catch (Error | Exception e) {
-            // Swallow errors caused by invalid byteref.
-            return "";
         }
     }
 }

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

@@ -0,0 +1,135 @@
+/*
+ * 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.inference;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.io.stream.BytesRefStreamOutput;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.compute.operator.EvalOperator;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentType;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.List;
+import java.util.Map;
+
+class XContentRowEncoder implements EvalOperator.ExpressionEvaluator {
+    private final XContentType xContentType;
+    private final BlockFactory blockFactory;
+    private final String[] fieldNames;
+    private final EvalOperator.ExpressionEvaluator[] fieldsValueEvaluators;
+
+    public static Factory yamlRowEncoderFactory(Map<String, EvalOperator.ExpressionEvaluator.Factory> fieldsEvaluatorFactories) {
+        return new Factory(XContentType.YAML, fieldsEvaluatorFactories);
+    }
+
+    private XContentRowEncoder(
+        XContentType xContentType,
+        BlockFactory blockFactory,
+        String[] fieldNames,
+        EvalOperator.ExpressionEvaluator[] fieldsValueEvaluators
+    ) {
+        assert fieldNames.length == fieldsValueEvaluators.length;
+        this.xContentType = xContentType;
+        this.blockFactory = blockFactory;
+        this.fieldNames = fieldNames;
+        this.fieldsValueEvaluators = fieldsValueEvaluators;
+    }
+
+    @Override
+    public void close() {
+        Releasables.closeExpectNoException(fieldsValueEvaluators);
+    }
+
+    @Override
+    public Block 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());
+        ) {
+            for (int b = 0; b < fieldValueBlocks.length; b++) {
+                fieldValueBlocks[b] = fieldsValueEvaluators[b].eval(page);
+            }
+
+            for (int pos = 0; pos < page.getPositionCount(); pos++) {
+                xContentBuilder.startObject();
+                for (int i = 0; i < fieldValueBlocks.length; i++) {
+                    String fieldName = fieldNames[i];
+                    Block currentBlock = fieldValueBlocks[i];
+                    if (currentBlock.isNull(pos)) {
+                        continue;
+                    }
+                    xContentBuilder.field(fieldName, toYamlValue(BlockUtils.toJavaObject(currentBlock, pos)));
+                }
+                xContentBuilder.endObject().flush();
+                outputBlockBuilder.appendBytesRef(outputStream.get());
+                outputStream.reset();
+            }
+
+            return outputBlockBuilder.build();
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        } finally {
+            Releasables.closeExpectNoException(fieldValueBlocks);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "XContentRowEncoder[content_type=[" + xContentType.toString() + "], field_names=" + List.of(fieldNames) + "]";
+    }
+
+    private Object toYamlValue(Object value) {
+        try {
+            return switch (value) {
+                case BytesRef b -> b.utf8ToString();
+                case List<?> l -> l.stream().map(this::toYamlValue).toList();
+                default -> value;
+            };
+        } catch (Error | Exception e) {
+            // Swallow errors caused by invalid byteref.
+            return "";
+        }
+    }
+
+    public static final class Factory implements EvalOperator.ExpressionEvaluator.Factory {
+        private final XContentType xContentType;
+        private final Map<String, EvalOperator.ExpressionEvaluator.Factory> fieldsEvaluatorFactories;
+
+        private Factory(XContentType xContentType, Map<String, EvalOperator.ExpressionEvaluator.Factory> fieldsEvaluatorFactories) {
+            this.xContentType = xContentType;
+            this.fieldsEvaluatorFactories = fieldsEvaluatorFactories;
+        }
+
+        @Override
+        public EvalOperator.ExpressionEvaluator get(DriverContext context) {
+            return new XContentRowEncoder(xContentType, context.blockFactory(), fieldNames(), fieldsValueEvaluators(context));
+        }
+
+        private String[] fieldNames() {
+            return fieldsEvaluatorFactories.keySet().toArray(String[]::new);
+        }
+
+        private EvalOperator.ExpressionEvaluator[] fieldsValueEvaluators(DriverContext context) {
+            return fieldsEvaluatorFactories.values()
+                .stream()
+                .map(factory -> factory.get(context))
+                .toArray(EvalOperator.ExpressionEvaluator[]::new);
+        }
+    }
+}