From 7ff51d66530f8d8ee8f270904ca76b9984cf470e Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Mon, 31 Mar 2025 11:40:03 +0100 Subject: [PATCH 1/8] [ES|QL] Infer the score mode to use from the Lucene collector This change uses the Lucene collector to infer which score mode to use when the topN collector is used. --- .../compute/lucene/LuceneCountOperator.java | 2 +- .../compute/lucene/LuceneMaxFactory.java | 4 +- .../compute/lucene/LuceneMinFactory.java | 4 +- .../compute/lucene/LuceneOperator.java | 12 +-- .../compute/lucene/LuceneSourceOperator.java | 22 +++-- .../lucene/LuceneTopNSourceOperator.java | 94 +++++++++++-------- ...TimeSeriesSortedSourceOperatorFactory.java | 4 +- .../lucene/LuceneSourceOperatorTests.java | 2 +- .../LuceneTopNSourceOperatorScoringTests.java | 4 +- .../lucene/LuceneTopNSourceOperatorTests.java | 8 +- 10 files changed, 88 insertions(+), 68 deletions(-) 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 34c27a5c1fdff..327303c45ad4b 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 @@ -49,7 +49,7 @@ public Factory( int taskConcurrency, int limit ) { - super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, ScoreMode.COMPLETE_NO_SCORES); + super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), dataPartitioning, taskConcurrency, limit, false); } @Override 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 ba7de22b1b821..3343750562cf5 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 @@ -23,6 +23,8 @@ import java.util.List; import java.util.function.Function; +import static org.elasticsearch.compute.lucene.LuceneOperator.weightFunction; + /** * Factory that generates an operator that finds the max value of a field using the {@link LuceneMinMaxOperator}. */ @@ -121,7 +123,7 @@ public LuceneMaxFactory( NumberType numberType, int limit ) { - super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, ScoreMode.COMPLETE_NO_SCORES); + super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), dataPartitioning, taskConcurrency, limit, false); 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 e3c6c8310373d..5f0849e882813 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 @@ -23,6 +23,8 @@ import java.util.List; import java.util.function.Function; +import static org.elasticsearch.compute.lucene.LuceneOperator.weightFunction; + /** * Factory that generates an operator that finds the min value of a field using the {@link LuceneMinMaxOperator}. */ @@ -121,7 +123,7 @@ public LuceneMinFactory( NumberType numberType, int limit ) { - super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, ScoreMode.COMPLETE_NO_SCORES); + super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), dataPartitioning, taskConcurrency, limit, false); 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 7547e2da3e184..2279603432d2f 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 @@ -11,7 +11,6 @@ import org.apache.lucene.search.BulkScorer; import org.apache.lucene.search.ConstantScoreQuery; import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.LeafCollector; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreMode; @@ -84,28 +83,27 @@ 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 boolean needsScore; protected final LuceneSliceQueue sliceQueue; /** * Build the factory. * - * @param scoreMode the {@link ScoreMode} passed to {@link IndexSearcher#createWeight} + * @param needsScore Whether the score is needed. */ protected Factory( List contexts, - Function queryFunction, + Function weightFunction, DataPartitioning dataPartitioning, int taskConcurrency, int limit, - ScoreMode scoreMode + boolean needsScore ) { this.limit = limit; - this.scoreMode = scoreMode; this.dataPartitioning = dataPartitioning; - var weightFunction = weightFunction(queryFunction, scoreMode); this.sliceQueue = LuceneSliceQueue.create(contexts, weightFunction, dataPartitioning, taskConcurrency); this.taskConcurrency = Math.min(sliceQueue.totalSlices(), taskConcurrency); + this.needsScore = needsScore; } public final int taskConcurrency() { 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 76f0fb0167b86..63dbf2926275e 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 @@ -11,7 +11,6 @@ import org.apache.lucene.search.LeafCollector; import org.apache.lucene.search.Query; 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; @@ -56,9 +55,16 @@ public Factory( int taskConcurrency, int maxPageSize, int limit, - boolean scoring + boolean needsScore ) { - super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, scoring ? COMPLETE : COMPLETE_NO_SCORES); + super( + contexts, + weightFunction(queryFunction, needsScore ? COMPLETE : COMPLETE_NO_SCORES), + dataPartitioning, + taskConcurrency, + limit, + needsScore + ); this.maxPageSize = maxPageSize; // TODO: use a single limiter for multiple stage execution this.limiter = limit == NO_LIMIT ? Limiter.NO_LIMIT : new Limiter(limit); @@ -66,7 +72,7 @@ public Factory( @Override public SourceOperator get(DriverContext driverContext) { - return new LuceneSourceOperator(driverContext.blockFactory(), maxPageSize, sliceQueue, limit, limiter, scoreMode); + return new LuceneSourceOperator(driverContext.blockFactory(), maxPageSize, sliceQueue, limit, limiter, needsScore); } public int maxPageSize() { @@ -81,8 +87,8 @@ public String describe() { + maxPageSize + ", limit = " + limit - + ", scoreMode = " - + scoreMode + + ", needsScore = " + + needsScore + "]"; } } @@ -94,7 +100,7 @@ public LuceneSourceOperator( LuceneSliceQueue sliceQueue, int limit, Limiter limiter, - ScoreMode scoreMode + boolean needsScore ) { super(blockFactory, maxPageSize, sliceQueue); this.minPageSize = Math.max(1, maxPageSize / 2); @@ -104,7 +110,7 @@ public LuceneSourceOperator( boolean success = false; try { this.docsBuilder = blockFactory.newIntVectorBuilder(estimatedSize); - if (scoreMode.needsScores()) { + if (needsScore) { scoreBuilder = blockFactory.newDoubleVectorBuilder(estimatedSize); this.leafCollector = new ScoringCollector(); } else { 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 83d69393a45a2..56a024e49a8d7 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 @@ -14,12 +14,12 @@ 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; @@ -36,6 +36,7 @@ import org.elasticsearch.search.sort.SortBuilder; import java.io.IOException; +import java.io.UncheckedIOException; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -43,9 +44,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) */ @@ -63,16 +61,16 @@ public Factory( int maxPageSize, int limit, List> sorts, - boolean scoring + boolean needsScore ) { - super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, scoring ? TOP_DOCS_WITH_SCORES : TOP_DOCS); + super(contexts, weightFunction(queryFunction, sorts), dataPartitioning, taskConcurrency, limit, needsScore); this.maxPageSize = maxPageSize; this.sorts = sorts; } @Override public SourceOperator get(DriverContext driverContext) { - return new LuceneTopNSourceOperator(driverContext.blockFactory(), maxPageSize, sorts, limit, sliceQueue, scoreMode); + return new LuceneTopNSourceOperator(driverContext.blockFactory(), maxPageSize, sorts, limit, sliceQueue, needsScore); } public int maxPageSize() { @@ -88,8 +86,8 @@ public String describe() { + maxPageSize + ", limit = " + limit - + ", scoreMode = " - + scoreMode + + ", needsScore = " + + needsScore + ", sorts = [" + notPrettySorts + "]]"; @@ -108,7 +106,7 @@ public String describe() { private PerShardCollector perShardCollector; private final List> sorts; private final int limit; - private final ScoreMode scoreMode; + private final boolean needsScore; public LuceneTopNSourceOperator( BlockFactory blockFactory, @@ -116,12 +114,12 @@ public LuceneTopNSourceOperator( List> sorts, int limit, LuceneSliceQueue sliceQueue, - ScoreMode scoreMode + boolean needsScore ) { super(blockFactory, maxPageSize, sliceQueue); this.sorts = sorts; this.limit = limit; - this.scoreMode = scoreMode; + this.needsScore = needsScore; } @Override @@ -163,7 +161,7 @@ private Page collect() throws IOException { try { if (perShardCollector == null || perShardCollector.shardContext.index() != scorer.shardContext().index()) { // TODO: share the bottom between shardCollectors - perShardCollector = newPerShardCollector(scorer.shardContext(), sorts, limit); + perShardCollector = newPerShardCollector(scorer.shardContext(), sorts, needsScore, limit); } var leafCollector = perShardCollector.getLeafCollector(scorer.leafReaderContext()); scorer.scoreNextRange(leafCollector, scorer.leafReaderContext().reader().getLiveDocs(), maxPageSize); @@ -261,7 +259,7 @@ private float getScore(ScoreDoc scoreDoc) { } private DoubleVector.Builder scoreVectorOrNull(int size) { - if (scoreMode.needsScores()) { + if (needsScore) { return blockFactory.newDoubleVectorFixedBuilder(size); } else { return null; @@ -271,37 +269,11 @@ private DoubleVector.Builder scoreVectorOrNull(int size) { @Override protected void describe(StringBuilder sb) { sb.append(", limit = ").append(limit); - sb.append(", scoreMode = ").append(scoreMode); + sb.append(", needsScore = ").append(needsScore); String notPrettySorts = sorts.stream().map(Strings::toString).collect(Collectors.joining(",")); sb.append(", sorts = [").append(notPrettySorts).append("]"); } - PerShardCollector newPerShardCollector(ShardContext shardContext, List> sorts, int limit) throws IOException { - Optional sortAndFormats = shardContext.buildSort(sorts); - if (sortAndFormats.isEmpty()) { - throw new IllegalStateException("sorts must not be disabled in TopN"); - } - if (scoreMode.needsScores() == false) { - return new NonScoringPerShardCollector(shardContext, sortAndFormats.get().sort, limit); - } else { - SortField[] sortFields = sortAndFormats.get().sort.getSort(); - if (sortFields != null && sortFields.length == 1 && sortFields[0].needsScores() && sortFields[0].getReverse() == false) { - // SORT _score DESC - return new ScoringPerShardCollector(shardContext, new TopScoreDocCollectorManager(limit, null, 0).newCollector()); - } else { - // SORT ..., _score, ... - var sort = new Sort(); - if (sortFields != null) { - var l = new ArrayList<>(Arrays.asList(sortFields)); - l.add(SortField.FIELD_DOC); - l.add(SortField.FIELD_SCORE); - sort = new Sort(l.toArray(SortField[]::new)); - } - return new ScoringPerShardCollector(shardContext, new TopFieldCollectorManager(sort, limit, null, 0).newCollector()); - } - } - } - abstract static class PerShardCollector { private final ShardContext shardContext; private final TopDocsCollector collector; @@ -336,4 +308,44 @@ static final class ScoringPerShardCollector extends PerShardCollector { super(shardContext, topDocsCollector); } } + + private static Function weightFunction(Function queryFunction, List> sorts) { + 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, false, 1).collector.scoreMode(); + return searcher.createWeight(searcher.rewrite(query), scoreMode, 1); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + }; + } + + private static PerShardCollector newPerShardCollector(ShardContext context, List> sorts, boolean needsScore, int limit) + throws IOException { + Optional sortAndFormats = context.buildSort(sorts); + if (sortAndFormats.isEmpty()) { + throw new IllegalStateException("sorts must not be disabled in TopN"); + } + if (needsScore == false) { + return new NonScoringPerShardCollector(context, sortAndFormats.get().sort, limit); + } + SortField[] sortFields = sortAndFormats.get().sort.getSort(); + if (sortFields != null && sortFields.length == 1 && sortFields[0].needsScores() && sortFields[0].getReverse() == false) { + // SORT _score DESC + return new ScoringPerShardCollector(context, new TopScoreDocCollectorManager(limit, null, 0).newCollector()); + } + + // SORT ..., _score, ... + var sort = new Sort(); + if (sortFields != null) { + var l = new ArrayList<>(Arrays.asList(sortFields)); + l.add(SortField.FIELD_DOC); + l.add(SortField.FIELD_SCORE); + sort = new Sort(l.toArray(SortField[]::new)); + } + return new ScoringPerShardCollector(context, new TopFieldCollectorManager(sort, limit, null, 0).newCollector()); + } } 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 8d37feb37d8b6..35738140d90c9 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 @@ -33,6 +33,8 @@ import java.util.List; import java.util.function.Function; +import static org.elasticsearch.compute.lucene.LuceneOperator.weightFunction; + /** * Creates a source operator that takes advantage of the natural sorting of segments in a tsdb index. *

@@ -56,7 +58,7 @@ private TimeSeriesSortedSourceOperatorFactory( int maxPageSize, int limit ) { - super(contexts, queryFunction, DataPartitioning.SHARD, taskConcurrency, limit, ScoreMode.COMPLETE_NO_SCORES); + super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), DataPartitioning.SHARD, taskConcurrency, limit, false); this.maxPageSize = maxPageSize; } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java index 39c85186d6a0d..1afbe083e3ef7 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java @@ -120,7 +120,7 @@ protected Matcher expectedToStringOfSimple() { protected Matcher expectedDescriptionOfSimple() { return matchesRegex( "LuceneSourceOperator" - + "\\[dataPartitioning = (DOC|SHARD|SEGMENT), maxPageSize = \\d+, limit = 100, scoreMode = (COMPLETE|COMPLETE_NO_SCORES)]" + + "\\[dataPartitioning = (DOC|SHARD|SEGMENT), maxPageSize = \\d+, limit = 100, needsScore = (true|false)]" ); } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java index 30f3b34fd42d9..2c62517edb94e 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java @@ -111,7 +111,7 @@ public Optional buildSort(List> sorts) { protected Matcher expectedToStringOfSimple() { return matchesRegex( "LuceneTopNSourceOperator\\[shards = \\[test], " - + "maxPageSize = \\d+, limit = 100, scoreMode = TOP_DOCS_WITH_SCORES, sorts = \\[\\{.+}]]" + + "maxPageSize = \\d+, limit = 100, needsScore = true, sorts = \\[\\{.+}]]" ); } @@ -119,7 +119,7 @@ protected Matcher expectedToStringOfSimple() { protected Matcher expectedDescriptionOfSimple() { return matchesRegex( "LuceneTopNSourceOperator\\[dataPartitioning = (DOC|SHARD|SEGMENT), " - + "maxPageSize = \\d+, limit = 100, scoreMode = TOP_DOCS_WITH_SCORES, sorts = \\[\\{.+}]]" + + "maxPageSize = \\d+, limit = 100, needsScore = true, sorts = \\[\\{.+}]]" ); } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java index 9210cfda65c04..7656a48f46551 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java @@ -114,19 +114,17 @@ public Optional buildSort(List> sorts) { @Override protected Matcher expectedToStringOfSimple() { - var s = scoring ? "TOP_DOCS_WITH_SCORES" : "TOP_DOCS"; return matchesRegex( - "LuceneTopNSourceOperator\\[shards = \\[test], maxPageSize = \\d+, limit = 100, scoreMode = " + s + ", sorts = \\[\\{.+}]]" + "LuceneTopNSourceOperator\\[shards = \\[test], maxPageSize = \\d+, limit = 100, needsScore = " + scoring + ", sorts = \\[\\{.+}]]" ); } @Override protected Matcher expectedDescriptionOfSimple() { - var s = scoring ? "TOP_DOCS_WITH_SCORES" : "TOP_DOCS"; return matchesRegex( "LuceneTopNSourceOperator" - + "\\[dataPartitioning = (DOC|SHARD|SEGMENT), maxPageSize = \\d+, limit = 100, scoreMode = " - + s + + "\\[dataPartitioning = (DOC|SHARD|SEGMENT), maxPageSize = \\d+, limit = 100, needsScore = " + + scoring + ", sorts = \\[\\{.+}]]" ); } From 4de27dcd9f3d97f8509bf1c74e795565f38435b4 Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Mon, 31 Mar 2025 11:42:50 +0100 Subject: [PATCH 2/8] Update docs/changelog/125930.yaml --- docs/changelog/125930.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 docs/changelog/125930.yaml diff --git a/docs/changelog/125930.yaml b/docs/changelog/125930.yaml new file mode 100644 index 0000000000000..6f68817d2da83 --- /dev/null +++ b/docs/changelog/125930.yaml @@ -0,0 +1,5 @@ +pr: 125930 +summary: Infer the score mode to use from the Lucene collector +area: "ES|QL, Search" +type: enhancement +issues: [] From 8a7e2e3cacf52a5e6270488e6b16e1217917a771 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Mon, 31 Mar 2025 10:52:22 +0000 Subject: [PATCH 3/8] [CI] Auto commit changes from spotless --- .../compute/lucene/LuceneTopNSourceOperatorScoringTests.java | 3 +-- .../compute/lucene/LuceneTopNSourceOperatorTests.java | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java index 2c62517edb94e..786cbd76f2efd 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorScoringTests.java @@ -110,8 +110,7 @@ public Optional buildSort(List> sorts) { @Override protected Matcher expectedToStringOfSimple() { return matchesRegex( - "LuceneTopNSourceOperator\\[shards = \\[test], " - + "maxPageSize = \\d+, limit = 100, needsScore = true, sorts = \\[\\{.+}]]" + "LuceneTopNSourceOperator\\[shards = \\[test], " + "maxPageSize = \\d+, limit = 100, needsScore = true, sorts = \\[\\{.+}]]" ); } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java index 7656a48f46551..7a1ee32e7a42d 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java @@ -115,7 +115,9 @@ public Optional buildSort(List> sorts) { @Override protected Matcher expectedToStringOfSimple() { return matchesRegex( - "LuceneTopNSourceOperator\\[shards = \\[test], maxPageSize = \\d+, limit = 100, needsScore = " + scoring + ", sorts = \\[\\{.+}]]" + "LuceneTopNSourceOperator\\[shards = \\[test], maxPageSize = \\d+, limit = 100, needsScore = " + + scoring + + ", sorts = \\[\\{.+}]]" ); } From 2e18147956843cfe6d175312292836b5ccdc59aa Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Mon, 31 Mar 2025 13:46:11 +0100 Subject: [PATCH 4/8] simplify collector creation and fix tests --- .../lucene/LuceneTopNSourceOperator.java | 21 ++++++++--------- .../planner/LocalExecutionPlannerTests.java | 23 ++++++++++++++++++- 2 files changed, 31 insertions(+), 13 deletions(-) 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 56a024e49a8d7..f1e8bfd89bcac 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 @@ -63,7 +63,7 @@ public Factory( List> sorts, boolean needsScore ) { - super(contexts, weightFunction(queryFunction, sorts), dataPartitioning, taskConcurrency, limit, needsScore); + super(contexts, weightFunction(queryFunction, sorts, needsScore), dataPartitioning, taskConcurrency, limit, needsScore); this.maxPageSize = maxPageSize; this.sorts = sorts; } @@ -309,13 +309,13 @@ static final class ScoringPerShardCollector extends PerShardCollector { } } - private static Function weightFunction(Function queryFunction, List> sorts) { + private static Function weightFunction(Function queryFunction, 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, false, 1).collector.scoreMode(); + var scoreMode = newPerShardCollector(ctx, sorts, needsScore, 1).collector.scoreMode(); return searcher.createWeight(searcher.rewrite(query), scoreMode, 1); } catch (IOException e) { throw new UncheckedIOException(e); @@ -332,20 +332,17 @@ private static PerShardCollector newPerShardCollector(ShardContext context, List if (needsScore == false) { return new NonScoringPerShardCollector(context, sortAndFormats.get().sort, limit); } - SortField[] sortFields = sortAndFormats.get().sort.getSort(); - if (sortFields != null && sortFields.length == 1 && sortFields[0].needsScores() && sortFields[0].getReverse() == false) { + Sort sort = sortAndFormats.get().sort; + if (Sort.RELEVANCE.equals(sort)) { // SORT _score DESC return new ScoringPerShardCollector(context, new TopScoreDocCollectorManager(limit, null, 0).newCollector()); } // SORT ..., _score, ... - var sort = new Sort(); - if (sortFields != null) { - var l = new ArrayList<>(Arrays.asList(sortFields)); - l.add(SortField.FIELD_DOC); - l.add(SortField.FIELD_SCORE); - sort = new Sort(l.toArray(SortField[]::new)); - } + var l = new ArrayList<>(Arrays.asList(sort.getSort())); + l.add(SortField.FIELD_DOC); + l.add(SortField.FIELD_SCORE); + sort = new Sort(l.toArray(SortField[]::new)); return new ScoringPerShardCollector(context, new TopFieldCollectorManager(sort, limit, null, 0).newCollector()); } } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java index 543a603a5e6f4..1df543ece13b6 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java @@ -31,6 +31,8 @@ import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy; import org.elasticsearch.index.mapper.MapperServiceTestCase; import org.elasticsearch.node.Node; +import org.elasticsearch.plugins.ExtensiblePlugin; +import org.elasticsearch.plugins.Plugin; import org.elasticsearch.search.internal.AliasFilter; import org.elasticsearch.search.internal.ContextIndexSearcher; import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; @@ -46,11 +48,14 @@ import org.elasticsearch.xpack.esql.plugin.EsqlPlugin; import org.elasticsearch.xpack.esql.plugin.QueryPragmas; import org.elasticsearch.xpack.esql.session.Configuration; +import org.elasticsearch.xpack.spatial.SpatialPlugin; import org.hamcrest.Matcher; import org.junit.After; import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -58,6 +63,7 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo; public class LocalExecutionPlannerTests extends MapperServiceTestCase { + @ParametersFactory public static Iterable parameters() throws Exception { List params = new ArrayList<>(); @@ -78,6 +84,19 @@ public LocalExecutionPlannerTests(@Name("estimatedRowSizeIsHuge") boolean estima this.estimatedRowSizeIsHuge = estimatedRowSizeIsHuge; } + @Override + protected Collection getPlugins() { + var plugin = new SpatialPlugin(); + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + public List loadExtensions(Class extensionPointType) { + return List.of(); + } + }); + + return Collections.singletonList(plugin); + } + @After public void closeIndex() throws IOException { IOUtils.close(reader, directory, () -> Releasables.close(releasables), releasables::clear); @@ -253,7 +272,9 @@ private List createShardContexts() th shardContexts.add( new EsPhysicalOperationProviders.DefaultShardContext( i, - createSearchExecutionContext(createMapperService(mapping(b -> {})), searcher), + createSearchExecutionContext(createMapperService(mapping(b -> { + b.startObject("point").field("type", "geo_point").endObject(); + })), searcher), AliasFilter.EMPTY ) ); From 433329fb1be25fe1a63f6fb3411d0922b71ba8c1 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Mon, 31 Mar 2025 12:55:44 +0000 Subject: [PATCH 5/8] [CI] Auto commit changes from spotless --- .../compute/lucene/LuceneTopNSourceOperator.java | 6 +++++- .../xpack/esql/planner/LocalExecutionPlannerTests.java | 10 +++------- 2 files changed, 8 insertions(+), 8 deletions(-) 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 f1e8bfd89bcac..cf03adcc4286b 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 @@ -309,7 +309,11 @@ static final class ScoringPerShardCollector extends PerShardCollector { } } - private static Function weightFunction(Function queryFunction, List> sorts, boolean needsScore) { + private static Function weightFunction( + Function queryFunction, + List> sorts, + boolean needsScore + ) { return ctx -> { final var query = queryFunction.apply(ctx); final var searcher = ctx.searcher(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java index 1df543ece13b6..6a7571d6964c9 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java @@ -270,13 +270,9 @@ private List createShardContexts() th ); for (int i = 0; i < numShards; i++) { shardContexts.add( - new EsPhysicalOperationProviders.DefaultShardContext( - i, - createSearchExecutionContext(createMapperService(mapping(b -> { - b.startObject("point").field("type", "geo_point").endObject(); - })), searcher), - AliasFilter.EMPTY - ) + new EsPhysicalOperationProviders.DefaultShardContext(i, createSearchExecutionContext(createMapperService(mapping(b -> { + b.startObject("point").field("type", "geo_point").endObject(); + })), searcher), AliasFilter.EMPTY) ); } releasables.add(searcher); From 79324e30e2614fccc8997c70e009559eea7dff64 Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Mon, 31 Mar 2025 14:45:15 +0100 Subject: [PATCH 6/8] fix another tests relying on score mode logging --- .../elasticsearch/xpack/esql/action/EsqlActionTaskIT.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java index 0294df0bfe5a4..abc9068554452 100644 --- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java @@ -167,7 +167,7 @@ public void testTaskContents() throws Exception { \\_AggregationOperator[mode = INITIAL, aggs = sum of longs] \\_ExchangeSinkOperator""".replace( "sourceStatus", - "dataPartitioning = SHARD, maxPageSize = " + pageSize() + ", limit = 2147483647, scoreMode = COMPLETE_NO_SCORES" + "dataPartitioning = SHARD, maxPageSize = " + pageSize() + ", limit = 2147483647, needsScore = false" ) ) ); @@ -502,7 +502,7 @@ public void testTaskContentsForTopNQuery() throws Exception { [{"pause_me":{"order":"asc","missing":"_last","unmapped_type":"long"}}]"""; String sourceStatus = "dataPartitioning = SHARD, maxPageSize = " + pageSize() - + ", limit = 1000, scoreMode = TOP_DOCS, sorts = " + + ", limit = 1000, needsScore = false, sorts = " + sortStatus; assertThat(dataTasks(tasks).get(0).description(), equalTo(""" \\_LuceneTopNSourceOperator[sourceStatus] @@ -545,7 +545,7 @@ public void testTaskContentsForLimitQuery() throws Exception { scriptPermits.release(pageSize() - prereleasedDocs); List tasks = getTasksRunning(); assertThat(dataTasks(tasks).get(0).description(), equalTo(""" - \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = limit(), scoreMode = COMPLETE_NO_SCORES] + \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = limit(), needsScore = false] \\_ValuesSourceReaderOperator[fields = [pause_me]] \\_ProjectOperator[projection = [1]] \\_ExchangeSinkOperator""".replace("pageSize()", Integer.toString(pageSize())).replace("limit()", limit))); @@ -573,7 +573,7 @@ public void testTaskContentsForGroupingStatsQuery() throws Exception { logger.info("unblocking script"); scriptPermits.release(pageSize()); List tasks = getTasksRunning(); - String sourceStatus = "dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, scoreMode = COMPLETE_NO_SCORES" + String sourceStatus = "dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, needsScore = false" .replace("pageSize()", Integer.toString(pageSize())); assertThat( dataTasks(tasks).get(0).description(), From cd0f336aa2015f0cc4161c55c5a3a06973112791 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Mon, 31 Mar 2025 14:01:22 +0000 Subject: [PATCH 7/8] [CI] Auto commit changes from spotless --- .../elasticsearch/xpack/esql/action/EsqlActionTaskIT.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java index abc9068554452..d6ef2bf1ce6a7 100644 --- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java @@ -573,8 +573,10 @@ public void testTaskContentsForGroupingStatsQuery() throws Exception { logger.info("unblocking script"); scriptPermits.release(pageSize()); List tasks = getTasksRunning(); - String sourceStatus = "dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, needsScore = false" - .replace("pageSize()", Integer.toString(pageSize())); + String sourceStatus = "dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, needsScore = false".replace( + "pageSize()", + Integer.toString(pageSize()) + ); assertThat( dataTasks(tasks).get(0).description(), equalTo( From 0d9ce3a7ae7a0ace4c347c0a38131f3696c4a6f9 Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Mon, 31 Mar 2025 17:42:34 +0100 Subject: [PATCH 8/8] Update 125930.yaml --- docs/changelog/125930.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/changelog/125930.yaml b/docs/changelog/125930.yaml index 6f68817d2da83..9bf7d18545772 100644 --- a/docs/changelog/125930.yaml +++ b/docs/changelog/125930.yaml @@ -1,5 +1,5 @@ pr: 125930 summary: Infer the score mode to use from the Lucene collector -area: "ES|QL, Search" +area: "ES|QL" type: enhancement issues: []