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 fb733e0cb0576..155e90c9f2cd3 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 @@ -56,7 +56,7 @@ public Factory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); } 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 49a6471b3e708..23852518710b4 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 @@ -128,7 +128,7 @@ public LuceneMaxFactory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); this.fieldName = fieldName; this.numberType = numberType; 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 1abb2e7f8085a..12bbcda63ea03 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 @@ -128,7 +128,7 @@ public LuceneMinFactory( taskConcurrency, limit, false, - ScoreMode.COMPLETE_NO_SCORES + shardContext -> ScoreMode.COMPLETE_NO_SCORES ); this.fieldName = fieldName; this.numberType = numberType; 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 11280746fc592..661fda846e724 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 @@ -101,11 +101,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 ee638e866d778..80b1f3e18ce96 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 @@ -107,12 +107,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 0b8f235a79b7a..9d8263024e1c5 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 @@ -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 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 16cc6d54eef0a..5ae98d86a3fbd 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) */ @@ -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 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/TimeSeriesSortedSourceOperatorFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSortedSourceOperatorFactory.java index c7184e4748452..1f6da6cc53952 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSortedSourceOperatorFactory.java +++ b/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; }