iter

Author: javanna

muted-tests.yml

@@ -330,9 +330,6 @@ tests:
 - class: org.elasticsearch.search.CCSDuelIT
   method: testTerminateAfter
   issue: https://github.com/elastic/elasticsearch/issues/126085
-- class: org.elasticsearch.search.sort.GeoDistanceIT
-  method: testDistanceSortingWithUnmappedField
-  issue: https://github.com/elastic/elasticsearch/issues/126118
 - class: org.elasticsearch.search.basic.SearchWithRandomDisconnectsIT
   method: testSearchWithRandomDisconnects
   issue: https://github.com/elastic/elasticsearch/issues/122707

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

@@ -303,13 +303,19 @@ private static void consumePartialMergeResult(
         Collection<InternalAggregations> aggsList
     ) {
         if (topDocsList != null) {
-            topDocsList.add(partialResult.reducedTopDocs);
+            addTopDocsToList(partialResult, topDocsList);
         }
         if (aggsList != null) {
             addAggsToList(partialResult, aggsList);
         }
     }
 
+    private static void addTopDocsToList(MergeResult partialResult, List<TopDocs> topDocsList) {
+        if (partialResult.reducedTopDocs != null) {
+            topDocsList.add(partialResult.reducedTopDocs);
+        }
+    }
+
     private static void addAggsToList(MergeResult partialResult, Collection<InternalAggregations> aggsList) {
         var aggs = partialResult.reducedAggs;
         if (aggs != null) {
@@ -340,7 +346,7 @@ private MergeResult partialReduce(
         if (hasTopDocs) {
             topDocsList = new ArrayList<>(resultSetSize);
             if (lastMerge != null) {
-                topDocsList.add(lastMerge.reducedTopDocs);
+                addTopDocsToList(lastMerge, topDocsList);
             }
         } else {
             topDocsList = null;
@@ -358,7 +364,7 @@ private MergeResult partialReduce(
                 }
             }
             // we have to merge here in the same way we collect on a shard
-            newTopDocs = topDocsList == null ? Lucene.EMPTY_TOP_DOCS : mergeTopDocs(topDocsList, topNSize, 0);
+            newTopDocs = topDocsList == null ? null : mergeTopDocs(topDocsList, topNSize, 0);
             newAggs = hasAggs
                 ? aggregate(
                     toConsume.iterator(),
@@ -647,7 +653,7 @@ private static void releaseAggs(List<QuerySearchResult> toConsume) {
 
     record MergeResult(
         List<SearchShard> processedShards,
-        TopDocs reducedTopDocs,
+        @Nullable TopDocs reducedTopDocs,
         @Nullable InternalAggregations reducedAggs,
         long estimatedSize
     ) implements Writeable {

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

@@ -60,6 +60,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -140,42 +141,31 @@ static SortedTopDocs sortDocs(
     }
 
     static TopDocs mergeTopDocs(List<TopDocs> results, int topN, int from) {
-        if (results.isEmpty()) {
+        List<TopDocs> topDocsList = results.stream().filter(Objects::nonNull).toList();
+        if (topDocsList.isEmpty()) {
             return null;
         }
-        final TopDocs topDocs = results.getFirst();
-        final TopDocs mergedTopDocs;
-        final int numShards = results.size();
+        final TopDocs topDocs = topDocsList.getFirst();
+        final int numShards = topDocsList.size();
         if (numShards == 1 && from == 0) { // only one shard and no pagination we can just return the topDocs as we got them.
             return topDocs;
-        } else if (topDocs instanceof TopFieldGroups firstTopDocs) {
+        }
+        final TopDocs mergedTopDocs;
+        if (topDocs instanceof TopFieldGroups firstTopDocs) {
             final Sort sort = new Sort(firstTopDocs.fields);
-            assert results.stream().noneMatch(topDoc -> topDoc == Lucene.EMPTY_TOP_DOCS);
-            final TopFieldGroups[] shardTopDocs = results.toArray(TopFieldGroups[]::new);
+            TopFieldGroups[] shardTopDocs = topDocsList.toArray(TopFieldGroups[]::new);
             mergedTopDocs = TopFieldGroups.merge(sort, from, topN, shardTopDocs, false);
         } else if (topDocs instanceof TopFieldDocs firstTopDocs) {
-            final Sort sort = checkSameSortTypes(results, firstTopDocs.fields);
-            final TopFieldDocs[] shardTopDocs = removeEmptyResults(results).toArray(TopFieldDocs[]::new);
+            TopFieldDocs[] shardTopDocs = topDocsList.toArray(TopFieldDocs[]::new);
+            final Sort sort = checkSameSortTypes(topDocsList, firstTopDocs.fields);
             mergedTopDocs = TopDocs.merge(sort, from, topN, shardTopDocs);
         } else {
-            final TopDocs[] shardTopDocs = results.toArray(new TopDocs[numShards]);
+            final TopDocs[] shardTopDocs = topDocsList.toArray(new TopDocs[numShards]);
             mergedTopDocs = TopDocs.merge(from, topN, shardTopDocs);
         }
         return mergedTopDocs;
     }
 
-    private static <T extends TopDocs> List<T> removeEmptyResults(List<T> results) {
-        List<T> nonEmptyResults = new ArrayList<>();
-        for (T result : results) {
-            if (result.totalHits.value() > 0 || result.totalHits.relation() != Relation.EQUAL_TO) {
-                nonEmptyResults.add(result);
-            } else {
-                assert result.scoreDocs.length == 0;
-            }
-        }
-        return nonEmptyResults;
-    }
-
     private static Sort checkSameSortTypes(Collection<TopDocs> results, SortField[] firstSortFields) {
         Sort sort = new Sort(firstSortFields);
         if (results.size() < 2) return sort;

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

@@ -27,7 +27,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
-import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.ListenableFuture;
@@ -722,7 +721,7 @@ private static final class QueryPerNodeState {
 
         private static final QueryPhaseResultConsumer.MergeResult EMPTY_PARTIAL_MERGE_RESULT = new QueryPhaseResultConsumer.MergeResult(
             List.of(),
-            Lucene.EMPTY_TOP_DOCS,
+            null,
             null,
             0L
         );
@@ -782,10 +781,12 @@ void onShardDone() {
                 // also collect the set of indices that may be part of a subsequent fetch operation here so that we can release all other
                 // indices without a roundtrip to the coordinating node
                 final BitSet relevantShardIndices = new BitSet(searchRequest.shards.size());
-                for (ScoreDoc scoreDoc : mergeResult.reducedTopDocs().scoreDocs) {
-                    final int localIndex = scoreDoc.shardIndex;
-                    scoreDoc.shardIndex = searchRequest.shards.get(localIndex).shardIndex;
-                    relevantShardIndices.set(localIndex);
+                if (mergeResult.reducedTopDocs() != null) {
+                    for (ScoreDoc scoreDoc : mergeResult.reducedTopDocs().scoreDocs) {
+                        final int localIndex = scoreDoc.shardIndex;
+                        scoreDoc.shardIndex = searchRequest.shards.get(localIndex).shardIndex;
+                        relevantShardIndices.set(localIndex);
+                    }
                 }
                 final Object[] results = new Object[queryPhaseResultConsumer.getNumShards()];
                 for (int i = 0; i < results.length; i++) {

server/src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -384,7 +384,9 @@ public static void writeTotalHits(StreamOutput out, TotalHits totalHits) throws
      * by shard for sorting purposes.
      */
     public static void writeTopDocsIncludingShardIndex(StreamOutput out, TopDocs topDocs) throws IOException {
-        if (topDocs instanceof TopFieldGroups topFieldGroups) {
+        if (topDocs == null) {
+            out.writeByte((byte) -1);
+        } else if (topDocs instanceof TopFieldGroups topFieldGroups) {
             out.writeByte((byte) 2);
             writeTotalHits(out, topDocs.totalHits);
             out.writeString(topFieldGroups.field);
@@ -424,7 +426,9 @@ public static void writeSortFieldArray(StreamOutput out, SortField[] sortFields)
      */
     public static TopDocs readTopDocsIncludingShardIndex(StreamInput in) throws IOException {
         byte type = in.readByte();
-        if (type == 0) {
+        if (type == -1) {
+            return null;
+        } else if (type == 0) {
             TotalHits totalHits = readTotalHits(in);
 
             final int scoreDocCount = in.readVInt();