fix test

Author: original-brownbear

server/src/main/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhase.java

@@ -22,11 +22,11 @@
 import org.elasticsearch.search.CanMatchShardResponse;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.SearchShardTarget;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.sort.MinAndMax;
-import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.Transport;
 
@@ -474,21 +474,23 @@ private synchronized List<SearchShardIterator> getIterator(List<SearchShardItera
         if (shouldSortShards(minAndMaxes) == false) {
             return shardsIts;
         }
-        int[] indexTranslation = sortShards(shardsIts, minAndMaxes, FieldSortBuilder.getPrimaryFieldSortOrNull(request.source()).order());
+        int[] indexTranslation = sortShards(shardsIts, minAndMaxes, request.source());
         List<SearchShardIterator> list = new ArrayList<>(indexTranslation.length);
         for (int in : indexTranslation) {
             list.add(shardsIts.get(in));
         }
         return list;
     }
 
-    public static <T extends Comparable<T>> int[] sortShards(List<T> shardsIts, MinAndMax<?>[] minAndMaxes, SortOrder order) {
+    public static <T extends Comparable<T>> int[] sortShards(List<T> shardsIts, MinAndMax<?>[] minAndMaxes, SearchSourceBuilder source) {
         int bound = shardsIts.size();
         List<Integer> toSort = new ArrayList<>(bound);
         for (int i = 0; i < bound; i++) {
             toSort.add(i);
         }
-        Comparator<? super MinAndMax<?>> keyComparator = forciblyCast(MinAndMax.getComparator(order));
+        Comparator<? super MinAndMax<?>> keyComparator = forciblyCast(
+            MinAndMax.getComparator(FieldSortBuilder.getPrimaryFieldSortOrNull(source).order())
+        );
         toSort.sort((idx1, idx2) -> {
             int res = keyComparator.compare(minAndMaxes[idx1], minAndMaxes[idx2]);
             if (res != 0) {

server/src/main/java/org/elasticsearch/action/search/SearchQueryThenFetchAsyncAction.java

@@ -48,7 +48,6 @@
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.query.QuerySearchResult;
-import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.sort.MinAndMax;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
@@ -480,11 +479,7 @@ protected void doRun() {
                         }
 
                         try {
-                            final int[] indexes = CanMatchPreFilterSearchPhase.sortShards(
-                                shards,
-                                minAndMax,
-                                FieldSortBuilder.getPrimaryFieldSortOrNull(request.source()).order()
-                            );
+                            final int[] indexes = CanMatchPreFilterSearchPhase.sortShards(shards, minAndMax, request.source());
                             final ShardToQuery[] orig = shards.toArray(new ShardToQuery[0]);
                             for (int i = 0; i < indexes.length; i++) {
                                 shards.set(i, orig[indexes[i]]);
@@ -652,31 +647,29 @@ static void registerNodeSearchAction(
             NodeQueryRequest::new,
             (request, channel, task) -> {
                 final SearchRequest searchRequest = request.searchRequest;
-                final IntUnaryOperator shards;
-                final ShardSearchRequest[] shardSearchRequests;
+                ShardSearchRequest[] shardSearchRequests = null;
+                IntUnaryOperator shards = IntUnaryOperator.identity();
                 final int shardCount = request.shards.size();
                 if (shardCount > 1 && hasPrimaryFieldSort(searchRequest.source())) {
-                    shardSearchRequests = new ShardSearchRequest[shardCount];
-                    @SuppressWarnings("rawtypes")
-                    final MinAndMax[] minAndMax = new MinAndMax[shardCount];
-                    for (int i = 0; i < minAndMax.length; i++) {
-                        ShardSearchRequest r = buildShardSearchRequestForLocal(request, request.shards.get(i));
-                        shardSearchRequests[i] = r;
-                        var canMatch = searchService.canMatch(r);
-                        if (canMatch.canMatch()) {
-                            r.setRunCanMatchInQueryPhase(false);
-                            minAndMax[i] = canMatch.estimatedMinAndMax();
+                    try {
+                        shardSearchRequests = new ShardSearchRequest[shardCount];
+                        @SuppressWarnings("rawtypes")
+                        final MinAndMax[] minAndMax = new MinAndMax[shardCount];
+                        for (int i = 0; i < minAndMax.length; i++) {
+                            ShardSearchRequest r = buildShardSearchRequestForLocal(request, request.shards.get(i));
+                            shardSearchRequests[i] = r;
+                            var canMatch = searchService.canMatch(r);
+                            if (canMatch.canMatch()) {
+                                r.setRunCanMatchInQueryPhase(false);
+                                minAndMax[i] = canMatch.estimatedMinAndMax();
+                            }
                         }
+                        int[] indexes = CanMatchPreFilterSearchPhase.sortShards(request.shards, minAndMax, searchRequest.source());
+                        shards = pos -> indexes[pos];
+                    } catch (Exception e) {
+                        // TODO: ignored for now but we'll be guaranteed to fail the query phase at this point, fix things to fail here
+                        // already
                     }
-                    int[] indexes = CanMatchPreFilterSearchPhase.sortShards(
-                        request.shards,
-                        minAndMax,
-                        FieldSortBuilder.getPrimaryFieldSortOrNull(searchRequest.source()).order()
-                    );
-                    shards = pos -> indexes[pos];
-                } else {
-                    shardSearchRequests = null;
-                    shards = IntUnaryOperator.identity();
                 }
                 final CancellableTask cancellableTask = (CancellableTask) task;
                 int workers = Math.min(searchRequest.getMaxConcurrentShardRequests(), Math.min(shardCount, searchPoolMax));