ES|QL: Infer score mode from collector (elastic#131006) (elastic#131083)

Author: ioanatia

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

@@ -56,7 +56,7 @@ public Factory(
                 taskConcurrency,
                 limit,
                 false,
-                ScoreMode.COMPLETE_NO_SCORES
+                shardContext -> ScoreMode.COMPLETE_NO_SCORES
             );
         }
 

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

@@ -128,7 +128,7 @@ public LuceneMaxFactory(
             taskConcurrency,
             limit,
             false,
-            ScoreMode.COMPLETE_NO_SCORES
+            shardContext -> ScoreMode.COMPLETE_NO_SCORES
         );
         this.fieldName = fieldName;
         this.numberType = numberType;

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

@@ -128,7 +128,7 @@ public LuceneMinFactory(
             taskConcurrency,
             limit,
             false,
-            ScoreMode.COMPLETE_NO_SCORES
+            shardContext -> ScoreMode.COMPLETE_NO_SCORES
         );
         this.fieldName = fieldName;
         this.numberType = numberType;

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

@@ -101,11 +101,18 @@ protected Factory(
             int taskConcurrency,
             int limit,
             boolean needsScore,
-            ScoreMode scoreMode
+            Function<ShardContext, ScoreMode> scoreModeFunction
         ) {
             this.limit = limit;
             this.dataPartitioning = dataPartitioning;
-            this.sliceQueue = LuceneSliceQueue.create(contexts, queryFunction, dataPartitioning, autoStrategy, taskConcurrency, scoreMode);
+            this.sliceQueue = LuceneSliceQueue.create(
+                contexts,
+                queryFunction,
+                dataPartitioning,
+                autoStrategy,
+                taskConcurrency,
+                scoreModeFunction
+            );
             this.taskConcurrency = Math.min(sliceQueue.totalSlices(), taskConcurrency);
             this.needsScore = needsScore;
         }

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

@@ -107,12 +107,14 @@ public static LuceneSliceQueue create(
         DataPartitioning dataPartitioning,
         Function<Query, PartitioningStrategy> autoStrategy,
         int taskConcurrency,
-        ScoreMode scoreMode
+        Function<ShardContext, ScoreMode> scoreModeFunction
     ) {
         List<LuceneSlice> slices = new ArrayList<>();
         Map<String, PartitioningStrategy> partitioningStrategies = new HashMap<>(contexts.size());
+
         for (ShardContext ctx : contexts) {
             for (QueryAndTags queryAndExtra : queryFunction.apply(ctx)) {
+                var scoreMode = scoreModeFunction.apply(ctx);
                 Query query = queryAndExtra.query;
                 query = scoreMode.needsScores() ? query : new ConstantScoreQuery(query);
                 /*

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

@@ -80,7 +80,7 @@ public Factory(
                 taskConcurrency,
                 limit,
                 needsScore,
-                needsScore ? COMPLETE : COMPLETE_NO_SCORES
+                shardContext -> needsScore ? COMPLETE : COMPLETE_NO_SCORES
             );
             this.maxPageSize = maxPageSize;
             // TODO: use a single limiter for multiple stage execution

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

@@ -12,14 +12,13 @@
 import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.LeafCollector;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.TopDocsCollector;
 import org.apache.lucene.search.TopFieldCollectorManager;
 import org.apache.lucene.search.TopScoreDocCollectorManager;
-import org.apache.lucene.search.Weight;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DocBlock;
@@ -44,9 +43,6 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-import static org.apache.lucene.search.ScoreMode.TOP_DOCS;
-import static org.apache.lucene.search.ScoreMode.TOP_DOCS_WITH_SCORES;
-
 /**
  * Source operator that builds Pages out of the output of a TopFieldCollector (aka TopN)
  */
@@ -73,7 +69,7 @@ public Factory(
                 taskConcurrency,
                 limit,
                 needsScore,
-                needsScore ? TOP_DOCS_WITH_SCORES : TOP_DOCS
+                scoreModeFunction(sorts, needsScore)
             );
             this.maxPageSize = maxPageSize;
             this.sorts = sorts;
@@ -323,18 +319,11 @@ static final class ScoringPerShardCollector extends PerShardCollector {
         }
     }
 
-    private static Function<ShardContext, Weight> weightFunction(
-        Function<ShardContext, Query> queryFunction,
-        List<SortBuilder<?>> sorts,
-        boolean needsScore
-    ) {
+    private static Function<ShardContext, ScoreMode> scoreModeFunction(List<SortBuilder<?>> sorts, boolean needsScore) {
         return ctx -> {
-            final var query = queryFunction.apply(ctx);
-            final var searcher = ctx.searcher();
             try {
                 // we create a collector with a limit of 1 to determine the appropriate score mode to use.
-                var scoreMode = newPerShardCollector(ctx, sorts, needsScore, 1).collector.scoreMode();
-                return searcher.createWeight(searcher.rewrite(query), scoreMode, 1);
+                return newPerShardCollector(ctx, sorts, needsScore, 1).collector.scoreMode();
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
             }

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

@@ -63,7 +63,7 @@ private TimeSeriesSortedSourceOperatorFactory(
             taskConcurrency,
             limit,
             false,
-            ScoreMode.COMPLETE_NO_SCORES
+            shardContext -> ScoreMode.COMPLETE_NO_SCORES
         );
         this.maxPageSize = maxPageSize;
     }