1 package org.apache.lucene.search;
4 * Licensed to the Apache Software Foundation (ASF) under one or more
5 * contributor license agreements. See the NOTICE file distributed with
6 * this work for additional information regarding copyright ownership.
7 * The ASF licenses this file to You under the Apache License, Version 2.0
8 * (the "License"); you may not use this file except in compliance with
9 * the License. You may obtain a copy of the License at
11 * http://www.apache.org/licenses/LICENSE-2.0
13 * Unless required by applicable law or agreed to in writing, software
14 * distributed under the License is distributed on an "AS IS" BASIS,
15 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 * See the License for the specific language governing permissions and
17 * limitations under the License.
20 import java.io.IOException;
22 import org.apache.lucene.util.PriorityQueue;
24 /** Represents hits returned by {@link
25 * Searcher#search(Query,Filter,int)} and {@link
26 * Searcher#search(Query,int)}. */
27 public class TopDocs implements java.io.Serializable {
29 /** The total number of hits for the query. */
32 /** The top hits for the query. */
33 public ScoreDoc[] scoreDocs;
35 /** Stores the maximum score value encountered, needed for normalizing. */
36 private float maxScore;
39 * Returns the maximum score value encountered. Note that in case
40 * scores are not tracked, this returns {@link Float#NaN}.
42 public float getMaxScore() {
46 /** Sets the maximum score value encountered. */
47 public void setMaxScore(float maxScore) {
48 this.maxScore=maxScore;
51 /** Constructs a TopDocs with a default maxScore=Float.NaN. */
52 TopDocs(int totalHits, ScoreDoc[] scoreDocs) {
53 this(totalHits, scoreDocs, Float.NaN);
56 public TopDocs(int totalHits, ScoreDoc[] scoreDocs, float maxScore) {
57 this.totalHits = totalHits;
58 this.scoreDocs = scoreDocs;
59 this.maxScore = maxScore;
63 private static class ShardRef {
64 // Which shard (index into shardHits[]):
67 // Which hit within the shard:
70 public ShardRef(int shardIndex) {
71 this.shardIndex = shardIndex;
75 public String toString() {
76 return "ShardRef(shardIndex=" + shardIndex + " hitIndex=" + hitIndex + ")";
80 // Specialized MergeSortQueue that just merges by
81 // relevance score, descending:
82 private static class ScoreMergeSortQueue extends PriorityQueue<ShardRef> {
83 final ScoreDoc[][] shardHits;
85 public ScoreMergeSortQueue(TopDocs[] shardHits) {
86 initialize(shardHits.length);
87 this.shardHits = new ScoreDoc[shardHits.length][];
88 for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
89 this.shardHits[shardIDX] = shardHits[shardIDX].scoreDocs;
93 // Returns true if first is < second
94 public boolean lessThan(ShardRef first, ShardRef second) {
95 assert first != second;
96 final float firstScore = shardHits[first.shardIndex][first.hitIndex].score;
97 final float secondScore = shardHits[second.shardIndex][second.hitIndex].score;
99 if (firstScore < secondScore) {
101 } else if (firstScore > secondScore) {
104 // Tie break: earlier shard wins
105 if (first.shardIndex < second.shardIndex) {
107 } else if (first.shardIndex > second.shardIndex) {
110 // Tie break in same shard: resolve however the
111 // shard had resolved it:
112 assert first.hitIndex != second.hitIndex;
113 return first.hitIndex < second.hitIndex;
119 private static class MergeSortQueue extends PriorityQueue<ShardRef> {
120 // These are really FieldDoc instances:
121 final ScoreDoc[][] shardHits;
122 final FieldComparator[] comparators;
123 final int[] reverseMul;
125 public MergeSortQueue(Sort sort, TopDocs[] shardHits) throws IOException {
126 initialize(shardHits.length);
127 this.shardHits = new ScoreDoc[shardHits.length][];
128 for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
129 final ScoreDoc[] shard = shardHits[shardIDX].scoreDocs;
130 //System.out.println(" init shardIdx=" + shardIDX + " hits=" + shard);
132 this.shardHits[shardIDX] = shard;
133 // Fail gracefully if API is misused:
134 for(int hitIDX=0;hitIDX<shard.length;hitIDX++) {
135 final ScoreDoc sd = shard[hitIDX];
136 if (!(sd instanceof FieldDoc)) {
137 throw new IllegalArgumentException("shard " + shardIDX + " was not sorted by the provided Sort (expected FieldDoc but got ScoreDoc)");
139 final FieldDoc fd = (FieldDoc) sd;
140 if (fd.fields == null) {
141 throw new IllegalArgumentException("shard " + shardIDX + " did not set sort field values (FieldDoc.fields is null); you must pass fillFields=true to IndexSearcher.search on each shard");
147 final SortField[] sortFields = sort.getSort();
148 comparators = new FieldComparator[sortFields.length];
149 reverseMul = new int[sortFields.length];
150 for(int compIDX=0;compIDX<sortFields.length;compIDX++) {
151 final SortField sortField = sortFields[compIDX];
152 comparators[compIDX] = sortField.getComparator(1, compIDX);
153 reverseMul[compIDX] = sortField.getReverse() ? -1 : 1;
157 // Returns true if first is < second
158 @SuppressWarnings("unchecked")
159 public boolean lessThan(ShardRef first, ShardRef second) {
160 assert first != second;
161 final FieldDoc firstFD = (FieldDoc) shardHits[first.shardIndex][first.hitIndex];
162 final FieldDoc secondFD = (FieldDoc) shardHits[second.shardIndex][second.hitIndex];
163 //System.out.println(" lessThan:\n first=" + first + " doc=" + firstFD.doc + " score=" + firstFD.score + "\n second=" + second + " doc=" + secondFD.doc + " score=" + secondFD.score);
165 for(int compIDX=0;compIDX<comparators.length;compIDX++) {
166 final FieldComparator comp = comparators[compIDX];
167 //System.out.println(" cmp idx=" + compIDX + " cmp1=" + firstFD.fields[compIDX] + " cmp2=" + secondFD.fields[compIDX] + " reverse=" + reverseMul[compIDX]);
169 final int cmp = reverseMul[compIDX] * comp.compareValues(firstFD.fields[compIDX], secondFD.fields[compIDX]);
172 //System.out.println(" return " + (cmp < 0));
177 // Tie break: earlier shard wins
178 if (first.shardIndex < second.shardIndex) {
179 //System.out.println(" return tb true");
181 } else if (first.shardIndex > second.shardIndex) {
182 //System.out.println(" return tb false");
185 // Tie break in same shard: resolve however the
186 // shard had resolved it:
187 //System.out.println(" return tb " + (first.hitIndex < second.hitIndex));
188 assert first.hitIndex != second.hitIndex;
189 return first.hitIndex < second.hitIndex;
194 /** Returns a new TopDocs, containing topN results across
195 * the provided TopDocs, sorting by the specified {@link
196 * Sort}. Each of the TopDocs must have been sorted by
197 * the same Sort, and sort field values must have been
198 * filled (ie, <code>fillFields=true</code> must be
200 * TopFieldCollector#create}.
202 * <p>Pass sort=null to merge sort by score descending.
204 * @lucene.experimental */
205 public static TopDocs merge(Sort sort, int topN, TopDocs[] shardHits) throws IOException {
207 final PriorityQueue<ShardRef> queue;
209 queue = new ScoreMergeSortQueue(shardHits);
211 queue = new MergeSortQueue(sort, shardHits);
214 int totalHitCount = 0;
215 int availHitCount = 0;
216 float maxScore = Float.MIN_VALUE;
217 for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
218 final TopDocs shard = shardHits[shardIDX];
219 if (shard.scoreDocs != null && shard.scoreDocs.length > 0) {
220 totalHitCount += shard.totalHits;
221 availHitCount += shard.scoreDocs.length;
222 queue.add(new ShardRef(shardIDX));
223 maxScore = Math.max(maxScore, shard.getMaxScore());
224 //System.out.println(" maxScore now " + maxScore + " vs " + shard.getMaxScore());
228 final ScoreDoc[] hits = new ScoreDoc[Math.min(topN, availHitCount)];
231 while(hitUpto < hits.length) {
232 assert queue.size() > 0;
233 ShardRef ref = queue.pop();
234 final ScoreDoc hit = shardHits[ref.shardIndex].scoreDocs[ref.hitIndex++];
235 hit.shardIndex = ref.shardIndex;
238 //System.out.println(" hitUpto=" + hitUpto);
239 //System.out.println(" doc=" + hits[hitUpto].doc + " score=" + hits[hitUpto].score);
243 if (ref.hitIndex < shardHits[ref.shardIndex].scoreDocs.length) {
244 // Not done with this these TopDocs yet:
250 return new TopDocs(totalHitCount, hits, maxScore);
252 return new TopFieldDocs(totalHitCount, hits, sort.getSort(), maxScore);