diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneCountOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneCountOperator.java index 626f0b00f0e2c..cded3a3494738 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneCountOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneCountOperator.java @@ -58,7 +58,7 @@ public Factory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); this.shardRefCounters = contexts; } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMaxFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMaxFactory.java index 82d766349ce9e..7e0003efaf669 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMaxFactory.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMaxFactory.java @@ -129,7 +129,7 @@ public LuceneMaxFactory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); this.contexts = contexts; this.fieldName = fieldName; diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMinFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMinFactory.java index 505e5cd3f0d75..000ade1b19562 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMinFactory.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMinFactory.java @@ -130,7 +130,7 @@ public LuceneMinFactory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); this.shardRefCounters = contexts; this.fieldName = fieldName; diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java index 366715530f665..f3eec4147f237 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java @@ -112,11 +112,18 @@ protected Factory( int taskConcurrency, int limit, boolean needsScore, - ScoreMode scoreMode + Function 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; } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSliceQueue.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSliceQueue.java index ee9f217303195..1a0b349b45f3f 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSliceQueue.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSliceQueue.java @@ -112,12 +112,14 @@ public static LuceneSliceQueue create( DataPartitioning dataPartitioning, Function autoStrategy, int taskConcurrency, - ScoreMode scoreMode + Function scoreModeFunction ) { List slices = new ArrayList<>(); Map 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); /* diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java index 9fedc595641b4..5201eede502df 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java @@ -81,7 +81,7 @@ public Factory( taskConcurrency, limit, needsScore, - needsScore ? COMPLETE : COMPLETE_NO_SCORES + shardContext -> needsScore ? COMPLETE : COMPLETE_NO_SCORES ); this.contexts = contexts; this.maxPageSize = maxPageSize; diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperator.java index d93a5493a3aba..553b4319f22e9 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperator.java +++ b/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) */ @@ -75,7 +71,7 @@ public Factory( taskConcurrency, limit, needsScore, - needsScore ? TOP_DOCS_WITH_SCORES : TOP_DOCS + scoreModeFunction(sorts, needsScore) ); this.contexts = contexts; this.maxPageSize = maxPageSize; @@ -331,18 +327,11 @@ static final class ScoringPerShardCollector extends PerShardCollector { } } - private static Function weightFunction( - Function queryFunction, - List> sorts, - boolean needsScore - ) { + private static Function scoreModeFunction(List> 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); } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorFactory.java index 97286761b7bcf..bb1d889db3f85 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorFactory.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorFactory.java @@ -45,7 +45,7 @@ private TimeSeriesSourceOperatorFactory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); this.contexts = contexts; this.maxPageSize = maxPageSize;