lucene-java-3.4.0/lucene/src/java/org/apache/lucene/search/TopScoreDocCollector.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.index.IndexReader;
  23
  24 /**
  25  * A {@link Collector} implementation that collects the top-scoring hits,
  26  * returning them as a {@link TopDocs}. This is used by {@link IndexSearcher} to
  27  * implement {@link TopDocs}-based search. Hits are sorted by score descending
  28  * and then (when the scores are tied) docID ascending. When you create an
  29  * instance of this collector you should know in advance whether documents are
  30  * going to be collected in doc Id order or not.
  31  *
  32  * <p><b>NOTE</b>: The values {@link Float#NaN} and
  33  * {Float#NEGATIVE_INFINITY} are not valid scores.  This
  34  * collector will not properly collect hits with such
  35  * scores.
  36  */
  37 public abstract class TopScoreDocCollector extends TopDocsCollector<ScoreDoc> {
  38
  39   // Assumes docs are scored in order.
  40   private static class InOrderTopScoreDocCollector extends TopScoreDocCollector {
  41     private InOrderTopScoreDocCollector(int numHits) {
  42       super(numHits);
  43     }
  44
  45     @Override
  46     public void collect(int doc) throws IOException {
  47       float score = scorer.score();
  48
  49       // This collector cannot handle these scores:
  50       assert score != Float.NEGATIVE_INFINITY;
  51       assert !Float.isNaN(score);
  52
  53       totalHits++;
  54       if (score <= pqTop.score) {
  55         // Since docs are returned in-order (i.e., increasing doc Id), a document
  56         // with equal score to pqTop.score cannot compete since HitQueue favors
  57         // documents with lower doc Ids. Therefore reject those docs too.
  58         return;
  59       }
  60       pqTop.doc = doc + docBase;
  61       pqTop.score = score;
  62       pqTop = pq.updateTop();
  63     }
  64
  65     @Override
  66     public boolean acceptsDocsOutOfOrder() {
  67       return false;
  68     }
  69   }
  70
  71   // Assumes docs are scored out of order.
  72   private static class OutOfOrderTopScoreDocCollector extends TopScoreDocCollector {
  73     private OutOfOrderTopScoreDocCollector(int numHits) {
  74       super(numHits);
  75     }
  76
  77     @Override
  78     public void collect(int doc) throws IOException {
  79       float score = scorer.score();
  80
  81       // This collector cannot handle NaN
  82       assert !Float.isNaN(score);
  83
  84       totalHits++;
  85       if (score < pqTop.score) {
  86         // Doesn't compete w/ bottom entry in queue
  87         return;
  88       }
  89       doc += docBase;
  90       if (score == pqTop.score && doc > pqTop.doc) {
  91         // Break tie in score by doc ID:
  92         return;
  93       }
  94       pqTop.doc = doc;
  95       pqTop.score = score;
  96       pqTop = pq.updateTop();
  97     }
  98
  99     @Override
 100     public boolean acceptsDocsOutOfOrder() {
 101       return true;
 102     }
 103   }
 104
 105   /**
 106    * Creates a new {@link TopScoreDocCollector} given the number of hits to
 107    * collect and whether documents are scored in order by the input
 108    * {@link Scorer} to {@link #setScorer(Scorer)}.
 109    *
 110    * <p><b>NOTE</b>: The instances returned by this method
 111    * pre-allocate a full array of length
 112    * <code>numHits</code>, and fill the array with sentinel
 113    * objects.
 114    */
 115   public static TopScoreDocCollector create(int numHits, boolean docsScoredInOrder) {
 116
 117     if (numHits <= 0) {
 118       throw new IllegalArgumentException("numHits must be > 0; please use TotalHitCountCollector if you just need the total hit count");
 119     }
 120
 121     if (docsScoredInOrder) {
 122       return new InOrderTopScoreDocCollector(numHits);
 123     } else {
 124       return new OutOfOrderTopScoreDocCollector(numHits);
 125     }
 126
 127   }
 128
 129   ScoreDoc pqTop;
 130   int docBase = 0;
 131   Scorer scorer;
 132
 133   // prevents instantiation
 134   private TopScoreDocCollector(int numHits) {
 135     super(new HitQueue(numHits, true));
 136     // HitQueue implements getSentinelObject to return a ScoreDoc, so we know
 137     // that at this point top() is already initialized.
 138     pqTop = pq.top();
 139   }
 140
 141   @Override
 142   protected TopDocs newTopDocs(ScoreDoc[] results, int start) {
 143     if (results == null) {
 144       return EMPTY_TOPDOCS;
 145     }
 146
 147     // We need to compute maxScore in order to set it in TopDocs. If start == 0,
 148     // it means the largest element is already in results, use its score as
 149     // maxScore. Otherwise pop everything else, until the largest element is
 150     // extracted and use its score as maxScore.
 151     float maxScore = Float.NaN;
 152     if (start == 0) {
 153       maxScore = results[0].score;
 154     } else {
 155       for (int i = pq.size(); i > 1; i--) { pq.pop(); }
 156       maxScore = pq.pop().score;
 157     }
 158
 159     return new TopDocs(totalHits, results, maxScore);
 160   }
 161
 162   @Override
 163   public void setNextReader(IndexReader reader, int base) {
 164     docBase = base;
 165   }
 166
 167   @Override
 168   public void setScorer(Scorer scorer) throws IOException {
 169     this.scorer = scorer;
 170   }
 171 }