Merge remote-tracking branch 'tteofili/esql_metadata_score' into carlosdelest/esql-match-operator-colon-scoring-demo

# Conflicts:
#	x-pack/plugin/esql/qa/testFixtures/src/main/resources/match-operator.csv-spec

Author: carlosdelest

docs/changelog/113120.yaml

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

docs/reference/esql/metadata-fields.asciidoc

@@ -20,6 +20,9 @@ supported ones are:
   * <<mapping-ignored-field,`_ignored`>>: the ignored source document fields. The field is of the type
   <<keyword,keyword>>.
 
+  * `_score`: the score of each document with respect to the portion of the ES|QL query that can be
+  pushed down to Lucene. The field is of the type <<number,double>>.
+
 To enable the access to these fields, the <<esql-from,`FROM`>> source command needs
 to be provided with a dedicated directive:
 

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

@@ -438,7 +438,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/LuceneTopNSourceOperator.java

@@ -14,11 +14,13 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.TopFieldCollector;
+import org.apache.lucene.search.TopDocsCollector;
 import org.apache.lucene.search.TopFieldCollectorManager;
 import org.elasticsearch.common.Strings;
 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;
@@ -38,10 +40,10 @@
 /**
  * Source operator that builds Pages out of the output of a TopFieldCollector (aka TopN)
  */
-public final class LuceneTopNSourceOperator extends LuceneOperator {
-    public static final class Factory extends LuceneOperator.Factory {
-        private final int maxPageSize;
-        private final List<SortBuilder<?>> sorts;
+public class LuceneTopNSourceOperator extends LuceneOperator {
+    public static class Factory extends LuceneOperator.Factory {
+        protected final int maxPageSize;
+        protected final List<SortBuilder<?>> sorts;
 
         public Factory(
             List<? extends ShardContext> contexts,
@@ -50,9 +52,10 @@ public Factory(
             int taskConcurrency,
             int maxPageSize,
             int limit,
-            List<SortBuilder<?>> sorts
+            List<SortBuilder<?>> sorts,
+            ScoreMode scoreMode
         ) {
-            super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, ScoreMode.TOP_DOCS);
+            super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, scoreMode);
             this.maxPageSize = maxPageSize;
             this.sorts = sorts;
         }
@@ -91,8 +94,8 @@ public String describe() {
     private int offset = 0;
 
     private PerShardCollector perShardCollector;
-    private final List<SortBuilder<?>> sorts;
-    private final int limit;
+    protected final List<SortBuilder<?>> sorts;
+    protected final int limit;
 
     public LuceneTopNSourceOperator(
         BlockFactory blockFactory,
@@ -145,7 +148,7 @@ private Page collect() throws IOException {
         try {
             if (perShardCollector == null || perShardCollector.shardContext.index() != scorer.shardContext().index()) {
                 // TODO: share the bottom between shardCollectors
-                perShardCollector = new PerShardCollector(scorer.shardContext(), sorts, limit);
+                perShardCollector = newPerShardCollector(scorer.shardContext(), sorts, limit);
             }
             var leafCollector = perShardCollector.getLeafCollector(scorer.leafReaderContext());
             scorer.scoreNextRange(leafCollector, scorer.leafReaderContext().reader().getLiveDocs(), maxPageSize);
@@ -171,7 +174,7 @@ private Page emit(boolean startEmitting) {
             assert isEmitting() == false : "offset=" + offset + " score_docs=" + Arrays.toString(scoreDocs);
             offset = 0;
             if (perShardCollector != null) {
-                scoreDocs = perShardCollector.topFieldCollector.topDocs().scoreDocs;
+                scoreDocs = perShardCollector.collector.topDocs().scoreDocs;
             } else {
                 scoreDocs = new ScoreDoc[0];
             }
@@ -183,10 +186,12 @@ private Page emit(boolean startEmitting) {
         IntBlock shard = null;
         IntVector segments = null;
         IntVector docs = null;
+        DocBlock docBlock = null;
         Page page = null;
         try (
             IntVector.Builder currentSegmentBuilder = blockFactory.newIntVectorFixedBuilder(size);
-            IntVector.Builder currentDocsBuilder = blockFactory.newIntVectorFixedBuilder(size)
+            IntVector.Builder currentDocsBuilder = blockFactory.newIntVectorFixedBuilder(size);
+            DoubleVector.Builder currentScoresBuilder = scoreVectorOrNull(size);
         ) {
             int start = offset;
             offset += size;
@@ -196,12 +201,14 @@ private Page emit(boolean startEmitting) {
                 int segment = ReaderUtil.subIndex(doc, leafContexts);
                 currentSegmentBuilder.appendInt(segment);
                 currentDocsBuilder.appendInt(doc - leafContexts.get(segment).docBase); // the offset inside the segment
+                consumeScore(scoreDocs[i], currentScoresBuilder);
             }
 
             shard = blockFactory.newConstantIntBlockWith(perShardCollector.shardContext.index(), size);
             segments = currentSegmentBuilder.build();
             docs = currentDocsBuilder.build();
-            page = new Page(size, new DocVector(shard.asVector(), segments, docs, null).asBlock());
+            docBlock = new DocVector(shard.asVector(), segments, docs, null).asBlock();
+            page = maybeAppendScore(new Page(size, docBlock), currentScoresBuilder);
         } finally {
             if (page == null) {
                 Releasables.closeExpectNoException(shard, segments, docs);
@@ -211,20 +218,41 @@ private Page emit(boolean startEmitting) {
         return page;
     }
 
+    protected DoubleVector.Builder scoreVectorOrNull(int size) {
+        return null; // no scoring
+    }
+
+    protected void consumeScore(ScoreDoc scoreDoc, DoubleVector.Builder currentScoresBuilder) {
+        // no scoring
+        assert currentScoresBuilder == null;
+    }
+
+    protected Page maybeAppendScore(Page page, DoubleVector.Builder currentScoresBuilder) {
+        // no scoring
+        assert currentScoresBuilder == null;
+        return page;
+    }
+
     @Override
     protected void describe(StringBuilder sb) {
         sb.append(", limit = ").append(limit);
         String notPrettySorts = sorts.stream().map(Strings::toString).collect(Collectors.joining(","));
         sb.append(", sorts = [").append(notPrettySorts).append("]");
     }
 
-    static final class PerShardCollector {
-        private final ShardContext shardContext;
-        private final TopFieldCollector topFieldCollector;
+    PerShardCollector newPerShardCollector(ShardContext shardContext, List<SortBuilder<?>> sorts, int limit) throws IOException {
+        return new PerShardCollector(shardContext, sorts, limit);
+    }
+
+    static class PerShardCollector {
+        protected ShardContext shardContext;
+        protected TopDocsCollector<?> collector;
         private int leafIndex;
         private LeafCollector leafCollector;
         private Thread currentThread;
 
+        PerShardCollector() {}
+
         PerShardCollector(ShardContext shardContext, List<SortBuilder<?>> sorts, int limit) throws IOException {
             this.shardContext = shardContext;
             Optional<SortAndFormats> sortAndFormats = shardContext.buildSort(sorts);
@@ -233,16 +261,17 @@ static final class PerShardCollector {
             }
 
             // We don't use CollectorManager here as we don't retrieve the total hits and sort by score.
-            this.topFieldCollector = new TopFieldCollectorManager(sortAndFormats.get().sort, limit, null, 0, false).newCollector();
+            this.collector = new TopFieldCollectorManager(sortAndFormats.get().sort, limit, null, 0, false).newCollector();
         }
 
         LeafCollector getLeafCollector(LeafReaderContext leafReaderContext) throws IOException {
             if (currentThread != Thread.currentThread() || leafIndex != leafReaderContext.ord) {
-                leafCollector = topFieldCollector.getLeafCollector(leafReaderContext);
+                leafCollector = collector.getLeafCollector(leafReaderContext);
                 leafIndex = leafReaderContext.ord;
                 currentThread = Thread.currentThread();
             }
             return leafCollector;
         }
+
     }
 }