diff --git a/docs/changelog/125930.yaml b/docs/changelog/125930.yaml new file mode 100644 index 0000000000000..9bf7d18545772 --- /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" +type: enhancement +issues: [] 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 2f72c309b5f21..2b5b7c06dec92 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; @@ -83,28 +82,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 d25cb3a870da7..77b33bd030aff 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.COMPLETE; -import static org.apache.lucene.search.ScoreMode.TOP_DOCS; - /** * Source operator that builds Pages out of the output of a TopFieldCollector (aka TopN) */ @@ -62,16 +60,16 @@ public Factory( int maxPageSize, int limit, List> sorts, - boolean scoring + boolean needsScore ) { - super(contexts, queryFunction, dataPartitioning, taskConcurrency, limit, scoring ? COMPLETE : TOP_DOCS); + super(contexts, weightFunction(queryFunction, sorts, needsScore), 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() { @@ -87,8 +85,8 @@ public String describe() { + maxPageSize + ", limit = " + limit - + ", scoreMode = " - + scoreMode + + ", needsScore = " + + needsScore + ", sorts = [" + notPrettySorts + "]]"; @@ -107,7 +105,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, @@ -115,12 +113,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 @@ -162,7 +160,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); @@ -260,7 +258,7 @@ private float getScore(ScoreDoc scoreDoc) { } private DoubleVector.Builder scoreVectorOrNull(int size) { - if (scoreMode.needsScores()) { + if (needsScore) { return blockFactory.newDoubleVectorFixedBuilder(size); } else { return null; @@ -270,43 +268,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, limit, false).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, limit, false).newCollector() - ); - } - } - } - abstract static class PerShardCollector { private final ShardContext shardContext; private final TopDocsCollector collector; @@ -341,4 +307,45 @@ static final class ScoringPerShardCollector extends PerShardCollector { super(shardContext, topDocsCollector); } } + + 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, needsScore, 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); + } + 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 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/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 2322ced00bc25..a69d7a90894b6 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 @@ -119,7 +119,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 20af40bcc6840..82d578e2a2d49 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 @@ -109,14 +109,14 @@ public Optional buildSort(List> sorts) { @Override protected Matcher expectedToStringOfSimple() { - return matchesRegex("LuceneTopNSourceOperator\\[maxPageSize = \\d+, limit = 100, scoreMode = COMPLETE, sorts = \\[\\{.+}]]"); + return matchesRegex("LuceneTopNSourceOperator\\[maxPageSize = \\d+, limit = 100, needsScore = true, sorts = \\[\\{.+}]]"); } @Override protected Matcher expectedDescriptionOfSimple() { return matchesRegex( "LuceneTopNSourceOperator" - + "\\[dataPartitioning = (DOC|SHARD|SEGMENT), maxPageSize = \\d+, limit = 100, scoreMode = COMPLETE, sorts = \\[\\{.+}]]" + + "\\[dataPartitioning = (DOC|SHARD|SEGMENT), 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 a6d652d499d84..c235bf5ae9883 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,17 +114,17 @@ public Optional buildSort(List> sorts) { @Override protected Matcher expectedToStringOfSimple() { - var s = scoring ? "COMPLETE" : "TOP_DOCS"; - return matchesRegex("LuceneTopNSourceOperator\\[maxPageSize = \\d+, limit = 100, scoreMode = " + s + ", sorts = \\[\\{.+}]]"); + return matchesRegex( + "LuceneTopNSourceOperator\\[maxPageSize = \\d+, limit = 100, needsScore = " + scoring + ", sorts = \\[\\{.+}]]" + ); } @Override protected Matcher expectedDescriptionOfSimple() { - var s = scoring ? "COMPLETE" : "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 = \\[\\{.+}]]" ); } 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 85c03ce7860d3..f973c3fe2347a 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 @@ -90,7 +90,7 @@ public void setup() { assumeTrue("requires query pragmas", canUseQueryPragmas()); nodeLevelReduction = randomBoolean(); READ_DESCRIPTION = """ - \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, scoreMode = COMPLETE_NO_SCORES] + \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, needsScore = false] \\_ValuesSourceReaderOperator[fields = [pause_me]] \\_AggregationOperator[mode = INITIAL, aggs = sum of longs] \\_ExchangeSinkOperator""".replace("pageSize()", Integer.toString(pageSize())); @@ -478,7 +478,7 @@ protected void doRun() throws Exception { public void testTaskContentsForTopNQuery() throws Exception { READ_DESCRIPTION = ("\\_LuceneTopNSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 1000, " - + "scoreMode = TOP_DOCS, " + + "needsScore = false, " + "sorts = [{\"pause_me\":{\"order\":\"asc\",\"missing\":\"_last\",\"unmapped_type\":\"long\"}}]]\n" + "\\_ValuesSourceReaderOperator[fields = [pause_me]]\n" + "\\_ProjectOperator[projection = [1]]\n" @@ -513,7 +513,7 @@ public void testTaskContentsForTopNQuery() throws Exception { public void testTaskContentsForLimitQuery() throws Exception { String limit = Integer.toString(randomIntBetween(pageSize() + 1, 2 * numberOfDocs())); READ_DESCRIPTION = """ - \\_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); @@ -542,7 +542,7 @@ public void testTaskContentsForLimitQuery() throws Exception { public void testTaskContentsForGroupingStatsQuery() throws Exception { READ_DESCRIPTION = """ - \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, scoreMode = COMPLETE_NO_SCORES] + \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 2147483647, needsScore = false] \\_ValuesSourceReaderOperator[fields = [foo]] \\_OrdinalsGroupingOperator(aggs = max of longs) \\_ExchangeSinkOperator""".replace("pageSize()", Integer.toString(pageSize())); 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 68810cc7d5367..a8bb8ef140be3 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 @@ -29,6 +29,8 @@ import org.elasticsearch.index.IndexMode; import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy; import org.elasticsearch.index.mapper.MapperServiceTestCase; +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; @@ -44,11 +46,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; @@ -56,6 +61,7 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo; public class LocalExecutionPlannerTests extends MapperServiceTestCase { + @ParametersFactory public static Iterable parameters() throws Exception { List params = new ArrayList<>(); @@ -76,6 +82,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); @@ -220,11 +239,9 @@ private List createShardContexts() th ); for (int i = 0; i < numShards; i++) { shardContexts.add( - new EsPhysicalOperationProviders.DefaultShardContext( - i, - createSearchExecutionContext(createMapperService(mapping(b -> {})), 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);