ESQL - enabling scoring with METADATA _score (#113120)

* ESQL - enabling scoring with METADATA _score
Co-authored-by: ChrisHegarty <[email protected]>

Author: tteofili

docs/changelog/113120.yaml

@@ -0,0 +1,5 @@
+pr: 113120
+summary: ESQL - enabling scoring with METADATA `_score`
+area: ES|QL
+type: enhancement
+issues: []

muted-tests.yml

@@ -224,6 +224,12 @@ tests:
   issue: https://github.com/elastic/elasticsearch/issues/117591
 - class: org.elasticsearch.repositories.s3.RepositoryS3ClientYamlTestSuiteIT
   issue: https://github.com/elastic/elasticsearch/issues/117596
+- class: "org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT"
+  method: "test {scoring.*}"
+  issue: https://github.com/elastic/elasticsearch/issues/117641
+- class: "org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecIT"
+  method: "test {scoring.*}"
+  issue: https://github.com/elastic/elasticsearch/issues/117641
 
 # Examples:
 #

server/src/main/java/org/elasticsearch/search/sort/SortBuilder.java

@@ -158,6 +158,11 @@ private static void parseCompoundSortField(XContentParser parser, List<SortBuild
     }
 
     public static Optional<SortAndFormats> buildSort(List<SortBuilder<?>> sortBuilders, SearchExecutionContext context) throws IOException {
+        return buildSort(sortBuilders, context, true);
+    }
+
+    public static Optional<SortAndFormats> buildSort(List<SortBuilder<?>> sortBuilders, SearchExecutionContext context, boolean optimize)
+        throws IOException {
         List<SortField> sortFields = new ArrayList<>(sortBuilders.size());
         List<DocValueFormat> sortFormats = new ArrayList<>(sortBuilders.size());
         for (SortBuilder<?> builder : sortBuilders) {
@@ -172,9 +177,13 @@ public static Optional<SortAndFormats> buildSort(List<SortBuilder<?>> sortBuilde
             if (sortFields.size() > 1) {
                 sort = true;
             } else {
-                SortField sortField = sortFields.get(0);
-                if (sortField.getType() == SortField.Type.SCORE && sortField.getReverse() == false) {
-                    sort = false;
+                if (optimize) {
+                    SortField sortField = sortFields.get(0);
+                    if (sortField.getType() == SortField.Type.SCORE && sortField.getReverse() == false) {
+                        sort = false;
+                    } else {
+                        sort = true;
+                    }
                 } else {
                     sort = true;
                 }

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

@@ -31,6 +31,7 @@
 public class MetadataAttribute extends TypedAttribute {
     public static final String TIMESTAMP_FIELD = "@timestamp";
     public static final String TSID_FIELD = "_tsid";
+    public static final String SCORE = "_score";
 
     static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Attribute.class,
@@ -50,7 +51,9 @@ public class MetadataAttribute extends TypedAttribute {
         SourceFieldMapper.NAME,
         tuple(DataType.SOURCE, false),
         IndexModeFieldMapper.NAME,
-        tuple(DataType.KEYWORD, true)
+        tuple(DataType.KEYWORD, true),
+        SCORE,
+        tuple(DataType.DOUBLE, false)
     );
 
     private final boolean searchable;

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java

@@ -79,6 +79,7 @@ public abstract static class Factory implements SourceOperator.SourceOperatorFac
         protected final DataPartitioning dataPartitioning;
         protected final int taskConcurrency;
         protected final int limit;
+        protected final ScoreMode scoreMode;
         protected final LuceneSliceQueue sliceQueue;
 
         /**
@@ -95,6 +96,7 @@ protected Factory(
             ScoreMode scoreMode
         ) {
             this.limit = limit;
+            this.scoreMode = scoreMode;
             this.dataPartitioning = dataPartitioning;
             var weightFunction = weightFunction(queryFunction, scoreMode);
             this.sliceQueue = LuceneSliceQueue.create(contexts, weightFunction, dataPartitioning, taskConcurrency);
@@ -438,7 +440,8 @@ static Function<ShardContext, Weight> weightFunction(Function<ShardContext, Quer
             final var query = queryFunction.apply(ctx);
             final var searcher = ctx.searcher();
             try {
-                return searcher.createWeight(searcher.rewrite(new ConstantScoreQuery(query)), scoreMode, 1);
+                Query actualQuery = scoreMode.needsScores() ? query : new ConstantScoreQuery(query);
+                return searcher.createWeight(searcher.rewrite(actualQuery), scoreMode, 1);
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
             }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java

@@ -13,7 +13,9 @@
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.DocBlock;
 import org.elasticsearch.compute.data.DocVector;
+import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
@@ -25,6 +27,9 @@
 import java.util.List;
 import java.util.function.Function;
 
+import static org.apache.lucene.search.ScoreMode.COMPLETE;
+import static org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES;
+
 /**
  * Source operator that incrementally runs Lucene searches
  */
@@ -34,6 +39,7 @@ public class LuceneSourceOperator extends LuceneOperator {
     private int remainingDocs;
 
     private IntVector.Builder docsBuilder;
+    private DoubleVector.Builder scoreBuilder;
     private final LeafCollector leafCollector;
     private final int minPageSize;
 
@@ -47,15 +53,16 @@ public Factory(
             DataPartitioning dataPartitioning,
             int taskConcurrency,
             int maxPageSize,
-            int limit
+            int limit,
+            boolean scoring
         ) {
-            super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, ScoreMode.COMPLETE_NO_SCORES);
+            super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, scoring ? COMPLETE : COMPLETE_NO_SCORES);
             this.maxPageSize = maxPageSize;
         }
 
         @Override
         public SourceOperator get(DriverContext driverContext) {
-            return new LuceneSourceOperator(driverContext.blockFactory(), maxPageSize, sliceQueue, limit);
+            return new LuceneSourceOperator(driverContext.blockFactory(), maxPageSize, sliceQueue, limit, scoreMode);
         }
 
         public int maxPageSize() {
@@ -70,32 +77,65 @@ public String describe() {
                 + maxPageSize
                 + ", limit = "
                 + limit
+                + ", scoreMode = "
+                + scoreMode
                 + "]";
         }
     }
 
-    public LuceneSourceOperator(BlockFactory blockFactory, int maxPageSize, LuceneSliceQueue sliceQueue, int limit) {
+    @SuppressWarnings("this-escape")
+    public LuceneSourceOperator(BlockFactory blockFactory, int maxPageSize, LuceneSliceQueue sliceQueue, int limit, ScoreMode scoreMode) {
         super(blockFactory, maxPageSize, sliceQueue);
         this.minPageSize = Math.max(1, maxPageSize / 2);
         this.remainingDocs = limit;
-        this.docsBuilder = blockFactory.newIntVectorBuilder(Math.min(limit, maxPageSize));
-        this.leafCollector = new LeafCollector() {
-            @Override
-            public void setScorer(Scorable scorer) {
-
+        int estimatedSize = Math.min(limit, maxPageSize);
+        boolean success = false;
+        try {
+            this.docsBuilder = blockFactory.newIntVectorBuilder(estimatedSize);
+            if (scoreMode.needsScores()) {
+                scoreBuilder = blockFactory.newDoubleVectorBuilder(estimatedSize);
+                this.leafCollector = new ScoringCollector();
+            } else {
+                scoreBuilder = null;
+                this.leafCollector = new LimitingCollector();
             }
+            success = true;
+        } finally {
+            if (success == false) {
+                close();
+            }
+        }
+    }
 
-            @Override
-            public void collect(int doc) {
-                if (remainingDocs > 0) {
-                    --remainingDocs;
-                    docsBuilder.appendInt(doc);
-                    currentPagePos++;
-                } else {
-                    throw new CollectionTerminatedException();
-                }
+    class LimitingCollector implements LeafCollector {
+        @Override
+        public void setScorer(Scorable scorer) {}
+
+        @Override
+        public void collect(int doc) throws IOException {
+            if (remainingDocs > 0) {
+                --remainingDocs;
+                docsBuilder.appendInt(doc);
+                currentPagePos++;
+            } else {
+                throw new CollectionTerminatedException();
             }
-        };
+        }
+    }
+
+    final class ScoringCollector extends LuceneSourceOperator.LimitingCollector {
+        private Scorable scorable;
+
+        @Override
+        public void setScorer(Scorable scorer) {
+            this.scorable = scorer;
+        }
+
+        @Override
+        public void collect(int doc) throws IOException {
+            super.collect(doc);
+            scoreBuilder.appendDouble(scorable.score());
+        }
     }
 
     @Override
@@ -139,15 +179,27 @@ public Page getCheckedOutput() throws IOException {
                 IntBlock shard = null;
                 IntBlock leaf = null;
                 IntVector docs = null;
+                DoubleVector scores = null;
+                DocBlock docBlock = null;
                 try {
                     shard = blockFactory.newConstantIntBlockWith(scorer.shardContext().index(), currentPagePos);
                     leaf = blockFactory.newConstantIntBlockWith(scorer.leafReaderContext().ord, currentPagePos);
                     docs = docsBuilder.build();
                     docsBuilder = blockFactory.newIntVectorBuilder(Math.min(remainingDocs, maxPageSize));
-                    page = new Page(currentPagePos, new DocVector(shard.asVector(), leaf.asVector(), docs, true).asBlock());
+                    docBlock = new DocVector(shard.asVector(), leaf.asVector(), docs, true).asBlock();
+                    shard = null;
+                    leaf = null;
+                    docs = null;
+                    if (scoreBuilder == null) {
+                        page = new Page(currentPagePos, docBlock);
+                    } else {
+                        scores = scoreBuilder.build();
+                        scoreBuilder = blockFactory.newDoubleVectorBuilder(Math.min(remainingDocs, maxPageSize));
+                        page = new Page(currentPagePos, docBlock, scores.asBlock());
+                    }
                 } finally {
                     if (page == null) {
-                        Releasables.closeExpectNoException(shard, leaf, docs);
+                        Releasables.closeExpectNoException(shard, leaf, docs, docBlock, scores);
                     }
                 }
                 currentPagePos = 0;
@@ -160,7 +212,7 @@ public Page getCheckedOutput() throws IOException {
 
     @Override
     public void close() {
-        docsBuilder.close();
+        Releasables.close(docsBuilder, scoreBuilder);
     }
 
     @Override