diff --git a/core/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java b/core/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java index 52fbf952fe408..068a005204d02 100644 --- a/core/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java +++ b/core/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java @@ -66,6 +66,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.StreamSupport; @@ -233,47 +234,19 @@ public ScoreDoc[] sortDocs(boolean ignoreFrom, AtomicArray sortedResult : results) { - TopDocs topDocs = sortedResult.value.queryResult().topDocs(); - // the 'index' field is the position in the resultsArr atomic array - shardTopDocs[sortedResult.index] = (CollapseTopFieldDocs) topDocs; - } + fillTopDocs(shardTopDocs, results, new CollapseTopFieldDocs(firstTopDocs.field, 0, new FieldDoc[0], + sort.getSort(), new Object[0], Float.NaN)); mergedTopDocs = CollapseTopFieldDocs.merge(sort, from, topN, shardTopDocs); } else if (result.queryResult().topDocs() instanceof TopFieldDocs) { TopFieldDocs firstTopDocs = (TopFieldDocs) result.queryResult().topDocs(); final Sort sort = new Sort(firstTopDocs.fields); - final TopFieldDocs[] shardTopDocs = new TopFieldDocs[resultsArr.length()]; - if (result.size() != shardTopDocs.length) { - // TopDocs#merge can't deal with null shard TopDocs - final TopFieldDocs empty = new TopFieldDocs(0, new FieldDoc[0], sort.getSort(), Float.NaN); - Arrays.fill(shardTopDocs, empty); - } - for (AtomicArray.Entry sortedResult : results) { - TopDocs topDocs = sortedResult.value.queryResult().topDocs(); - // the 'index' field is the position in the resultsArr atomic array - shardTopDocs[sortedResult.index] = (TopFieldDocs) topDocs; - } + fillTopDocs(shardTopDocs, results, new TopFieldDocs(0, new FieldDoc[0], sort.getSort(), Float.NaN)); mergedTopDocs = TopDocs.merge(sort, from, topN, shardTopDocs); } else { final TopDocs[] shardTopDocs = new TopDocs[resultsArr.length()]; - if (result.size() != shardTopDocs.length) { - // TopDocs#merge can't deal with null shard TopDocs - Arrays.fill(shardTopDocs, Lucene.EMPTY_TOP_DOCS); - } - for (AtomicArray.Entry sortedResult : results) { - TopDocs topDocs = sortedResult.value.queryResult().topDocs(); - // the 'index' field is the position in the resultsArr atomic array - shardTopDocs[sortedResult.index] = topDocs; - } + fillTopDocs(shardTopDocs, results, Lucene.EMPTY_TOP_DOCS); mergedTopDocs = TopDocs.merge(from, topN, shardTopDocs); } @@ -314,6 +287,20 @@ public ScoreDoc[] sortDocs(boolean ignoreFrom, AtomicArray void fillTopDocs(T[] shardTopDocs, + List> results, + T empytTopDocs) { + if (results.size() != shardTopDocs.length) { + // TopDocs#merge can't deal with null shard TopDocs + Arrays.fill(shardTopDocs, empytTopDocs); + } + for (AtomicArray.Entry resultProvider : results) { + final T topDocs = (T) resultProvider.value.queryResult().topDocs(); + assert topDocs != null : "top docs must not be null in a valid result"; + // the 'index' field is the position in the resultsArr atomic array + shardTopDocs[resultProvider.index] = topDocs; + } + } public ScoreDoc[] getLastEmittedDocPerShard(ReducedQueryPhase reducedQueryPhase, ScoreDoc[] sortedScoreDocs, int numShards) { ScoreDoc[] lastEmittedDocPerShard = new ScoreDoc[numShards]; diff --git a/core/src/test/java/org/elasticsearch/action/search/SearchPhaseControllerTests.java b/core/src/test/java/org/elasticsearch/action/search/SearchPhaseControllerTests.java index ee68f2f1f9808..36756aba946bd 100644 --- a/core/src/test/java/org/elasticsearch/action/search/SearchPhaseControllerTests.java +++ b/core/src/test/java/org/elasticsearch/action/search/SearchPhaseControllerTests.java @@ -21,6 +21,7 @@ import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TopDocs; +import org.elasticsearch.common.lucene.Lucene; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.text.Text; import org.elasticsearch.common.util.BigArrays; @@ -347,4 +348,31 @@ public void testNewSearchPhaseResults() { } } } + + public void testFillTopDocs() { + final int maxIters = randomIntBetween(5, 15); + for (int iters = 0; iters < maxIters; iters++) { + TopDocs[] topDocs = new TopDocs[randomIntBetween(2, 100)]; + int numShards = topDocs.length; + AtomicArray resultProviderAtomicArray = generateQueryResults(numShards, Collections.emptyList(), + 2, randomBoolean()); + if (randomBoolean()) { + int maxNull = randomIntBetween(1, topDocs.length - 1); + for (int i = 0; i < maxNull; i++) { + resultProviderAtomicArray.set(randomIntBetween(0, resultProviderAtomicArray.length() - 1), null); + } + } + SearchPhaseController.fillTopDocs(topDocs, resultProviderAtomicArray.asList(), Lucene.EMPTY_TOP_DOCS); + for (int i = 0; i < topDocs.length; i++) { + assertNotNull(topDocs[i]); + if (topDocs[i] == Lucene.EMPTY_TOP_DOCS) { + assertNull(resultProviderAtomicArray.get(i)); + } else { + assertNotNull(resultProviderAtomicArray.get(i)); + assertNotNull(resultProviderAtomicArray.get(i).queryResult()); + assertSame(resultProviderAtomicArray.get(i).queryResult().topDocs(), topDocs[i]); + } + } + } + } }