simplify collector creation and fix tests

Author: jimczi

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

@@ -63,7 +63,7 @@ public Factory(
             List<SortBuilder<?>> 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<ShardContext, Weight> weightFunction(Function<ShardContext, Query> queryFunction, List<SortBuilder<?>> sorts) {
+    private static Function<ShardContext, Weight> weightFunction(Function<ShardContext, Query> queryFunction, 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, 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());
     }
 }

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,18 +48,22 @@
 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;
 
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class LocalExecutionPlannerTests extends MapperServiceTestCase {
+
     @ParametersFactory
     public static Iterable<Object[]> parameters() throws Exception {
         List<Object[]> params = new ArrayList<>();
@@ -78,6 +84,19 @@ public LocalExecutionPlannerTests(@Name("estimatedRowSizeIsHuge") boolean estima
         this.estimatedRowSizeIsHuge = estimatedRowSizeIsHuge;
     }
 
+    @Override
+    protected Collection<Plugin> getPlugins() {
+        var plugin = new SpatialPlugin();
+        plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() {
+            @Override
+            public <T> List<T> loadExtensions(Class<T> 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<EsPhysicalOperationProviders.ShardContext> 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
                 )
             );