lucene-java-3.5.0/lucene/src/java/org/apache/lucene/search/TopDocs.java

   1 package org.apache.lucene.search;
   2
   3 /**
   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
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  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.
  18  */
  19
  20 import java.io.IOException;
  21
  22 import org.apache.lucene.util.PriorityQueue;
  23
  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 {
  28
  29   /** The total number of hits for the query. */
  30   public int totalHits;
  31
  32   /** The top hits for the query. */
  33   public ScoreDoc[] scoreDocs;
  34
  35   /** Stores the maximum score value encountered, needed for normalizing. */
  36   private float maxScore;
  37
  38   /**
  39    * Returns the maximum score value encountered. Note that in case
  40    * scores are not tracked, this returns {@link Float#NaN}.
  41    */
  42   public float getMaxScore() {
  43     return maxScore;
  44   }
  45
  46   /** Sets the maximum score value encountered. */
  47   public void setMaxScore(float maxScore) {
  48     this.maxScore=maxScore;
  49   }
  50
  51   /** Constructs a TopDocs with a default maxScore=Float.NaN. */
  52   TopDocs(int totalHits, ScoreDoc[] scoreDocs) {
  53     this(totalHits, scoreDocs, Float.NaN);
  54   }
  55
  56   public TopDocs(int totalHits, ScoreDoc[] scoreDocs, float maxScore) {
  57     this.totalHits = totalHits;
  58     this.scoreDocs = scoreDocs;
  59     this.maxScore = maxScore;
  60   }
  61
  62   // Refers to one hit:
  63   private static class ShardRef {
  64     // Which shard (index into shardHits[]):
  65     final int shardIndex;
  66
  67     // Which hit within the shard:
  68     int hitIndex;
  69
  70     public ShardRef(int shardIndex) {
  71       this.shardIndex = shardIndex;
  72     }
  73
  74     @Override
  75     public String toString() {
  76       return "ShardRef(shardIndex=" + shardIndex + " hitIndex=" + hitIndex + ")";
  77     }
  78   };
  79
  80   // Specialized MergeSortQueue that just merges by
  81   // relevance score, descending:
  82   private static class ScoreMergeSortQueue extends PriorityQueue<ShardRef> {
  83     final ScoreDoc[][] shardHits;
  84
  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;
  90       }
  91     }
  92
  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;
  98
  99       if (firstScore < secondScore) {
 100         return false;
 101       } else if (firstScore > secondScore) {
 102         return true;
 103       } else {
 104         // Tie break: earlier shard wins
 105         if (first.shardIndex < second.shardIndex) {
 106           return true;
 107         } else if (first.shardIndex > second.shardIndex) {
 108           return false;
 109         } else {
 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;
 114         }
 115       }
 116     }
 117   }
 118
 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;
 124
 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);
 131         if (shard != null) {
 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)");
 138             }
 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");
 142             }
 143           }
 144         }
 145       }
 146
 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;
 154       }
 155     }
 156
 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);
 164
 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]);
 168
 169         final int cmp = reverseMul[compIDX] * comp.compareValues(firstFD.fields[compIDX], secondFD.fields[compIDX]);
 170
 171         if (cmp != 0) {
 172           //System.out.println("    return " + (cmp < 0));
 173           return cmp < 0;
 174         }
 175       }
 176
 177       // Tie break: earlier shard wins
 178       if (first.shardIndex < second.shardIndex) {
 179         //System.out.println("    return tb true");
 180         return true;
 181       } else if (first.shardIndex > second.shardIndex) {
 182         //System.out.println("    return tb false");
 183         return false;
 184       } else {
 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;
 190       }
 191     }
 192   }
 193
 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
 199    *  passed to {@link
 200    *  TopFieldCollector#create}.
 201    *
 202    * <p>Pass sort=null to merge sort by score descending.
 203    *
 204    * @lucene.experimental */
 205   public static TopDocs merge(Sort sort, int topN, TopDocs[] shardHits) throws IOException {
 206
 207     final PriorityQueue<ShardRef> queue;
 208     if (sort == null) {
 209       queue = new ScoreMergeSortQueue(shardHits);
 210     } else {
 211       queue = new MergeSortQueue(sort, shardHits);
 212     }
 213
 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());
 225       }
 226     }
 227
 228     final ScoreDoc[] hits = new ScoreDoc[Math.min(topN, availHitCount)];
 229
 230     int hitUpto = 0;
 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;
 236       hits[hitUpto] = hit;
 237
 238       //System.out.println("  hitUpto=" + hitUpto);
 239       //System.out.println("    doc=" + hits[hitUpto].doc + " score=" + hits[hitUpto].score);
 240
 241       hitUpto++;
 242
 243       if (ref.hitIndex < shardHits[ref.shardIndex].scoreDocs.length) {
 244         // Not done with this these TopDocs yet:
 245         queue.add(ref);
 246       }
 247     }
 248
 249     if (sort == null) {
 250       return new TopDocs(totalHitCount, hits, maxScore);
 251     } else {
 252       return new TopFieldDocs(totalHitCount, hits, sort.getSort(), maxScore);
 253     }
 254   }
 255 }