lucene-java-3.4.0/lucene/contrib/facet/src/java/org/apache/lucene/facet/search/ScoredDocIdCollector.java

   1 package org.apache.lucene.facet.search;
   2
   3 import java.io.IOException;
   4
   5 import org.apache.lucene.index.IndexReader;
   6 import org.apache.lucene.search.Collector;
   7 import org.apache.lucene.search.DocIdSet;
   8 import org.apache.lucene.search.DocIdSetIterator;
   9 import org.apache.lucene.search.Scorer;
  10 import org.apache.lucene.util.ArrayUtil;
  11 import org.apache.lucene.util.OpenBitSet;
  12
  13 /**
  14  * Licensed to the Apache Software Foundation (ASF) under one or more
  15  * contributor license agreements.  See the NOTICE file distributed with
  16  * this work for additional information regarding copyright ownership.
  17  * The ASF licenses this file to You under the Apache License, Version 2.0
  18  * (the "License"); you may not use this file except in compliance with
  19  * the License.  You may obtain a copy of the License at
  20  *
  21  *     http://www.apache.org/licenses/LICENSE-2.0
  22  *
  23  * Unless required by applicable law or agreed to in writing, software
  24  * distributed under the License is distributed on an "AS IS" BASIS,
  25  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  26  * See the License for the specific language governing permissions and
  27  * limitations under the License.
  28  */
  29
  30 /**
  31  * A {@link Collector} which stores all docIDs and their scores in a
  32  * {@link ScoredDocIDs} instance. If scoring is not enabled, then the default
  33  * score as set in {@link #setDefaultScore(float)} (or
  34  * {@link ScoredDocIDsIterator#DEFAULT_SCORE}) will be set for all documents.
  35  *
  36  * @lucene.experimental
  37  */
  38 public abstract class ScoredDocIdCollector extends Collector {
  39
  40   private static final class NonScoringDocIdCollector extends ScoredDocIdCollector {
  41
  42     float defaultScore = ScoredDocIDsIterator.DEFAULT_SCORE;
  43
  44     @SuppressWarnings("synthetic-access")
  45     public NonScoringDocIdCollector(int maxDoc) {
  46       super(maxDoc);
  47     }
  48
  49     @Override
  50     public boolean acceptsDocsOutOfOrder() { return true; }
  51
  52     @Override
  53     public void collect(int doc) throws IOException {
  54       docIds.fastSet(docBase + doc);
  55       ++numDocIds;
  56     }
  57
  58     @Override
  59     public float getDefaultScore() {
  60       return defaultScore;
  61     }
  62
  63     @Override
  64     public ScoredDocIDsIterator scoredDocIdsIterator() throws IOException {
  65       return new ScoredDocIDsIterator() {
  66
  67         private DocIdSetIterator docIdsIter = docIds.iterator();
  68         private int nextDoc;
  69
  70         public int getDocID() { return nextDoc; }
  71         public float getScore() { return defaultScore; }
  72
  73         public boolean next() {
  74           try {
  75             nextDoc = docIdsIter.nextDoc();
  76             return nextDoc != DocIdSetIterator.NO_MORE_DOCS;
  77           } catch (IOException e) {
  78             // This should not happen as we're iterating over an OpenBitSet. For
  79             // completeness, terminate iteration
  80             nextDoc = DocIdSetIterator.NO_MORE_DOCS;
  81             return false;
  82           }
  83         }
  84
  85       };
  86     }
  87
  88     @Override
  89     public void setDefaultScore(float defaultScore) {
  90       this.defaultScore = defaultScore;
  91     }
  92
  93     @Override
  94     public void setScorer(Scorer scorer) throws IOException {}
  95   }
  96
  97   private static final class ScoringDocIdCollector extends ScoredDocIdCollector {
  98
  99     float[] scores;
 100     private Scorer scorer;
 101
 102     @SuppressWarnings("synthetic-access")
 103     public ScoringDocIdCollector(int maxDoc) {
 104       super(maxDoc);
 105       scores = new float[maxDoc];
 106     }
 107
 108     @Override
 109     public boolean acceptsDocsOutOfOrder() { return false; }
 110
 111     @Override
 112     public void collect(int doc) throws IOException {
 113       docIds.fastSet(docBase + doc);
 114
 115       float score = this.scorer.score();
 116       if (numDocIds >= scores.length) {
 117         float[] newScores = new float[ArrayUtil.oversize(numDocIds + 1, 4)];
 118         System.arraycopy(scores, 0, newScores, 0, numDocIds);
 119         scores = newScores;
 120       }
 121       scores[numDocIds] = score;
 122       ++numDocIds;
 123     }
 124
 125     @Override
 126     public ScoredDocIDsIterator scoredDocIdsIterator() throws IOException {
 127       return new ScoredDocIDsIterator() {
 128
 129         private DocIdSetIterator docIdsIter = docIds.iterator();
 130         private int nextDoc;
 131         private int scoresIdx = -1;
 132
 133         public int getDocID() { return nextDoc; }
 134         public float getScore() { return scores[scoresIdx]; }
 135
 136         public boolean next() {
 137           try {
 138             nextDoc = docIdsIter.nextDoc();
 139             if (nextDoc == DocIdSetIterator.NO_MORE_DOCS) {
 140               return false;
 141             }
 142             ++scoresIdx;
 143             return true;
 144           } catch (IOException e) {
 145             // This should not happen as we're iterating over an OpenBitSet. For
 146             // completeness, terminate iteration
 147             nextDoc = DocIdSetIterator.NO_MORE_DOCS;
 148             return false;
 149           }
 150         }
 151
 152       };
 153     }
 154
 155     @Override
 156     public float getDefaultScore() { return ScoredDocIDsIterator.DEFAULT_SCORE; }
 157
 158     @Override
 159     public void setDefaultScore(float defaultScore) {}
 160
 161     @Override
 162     public void setScorer(Scorer scorer) throws IOException {
 163       this.scorer = scorer;
 164     }
 165   }
 166
 167   protected int numDocIds;
 168   protected int docBase;
 169   protected final OpenBitSet docIds;
 170
 171   /**
 172    * Creates a new {@link ScoredDocIdCollector} with the given parameters.
 173    *
 174    * @param maxDoc the number of documents that are expected to be collected.
 175    *        Note that if more documents are collected, unexpected exceptions may
 176    *        be thrown. Usually you should pass {@link IndexReader#maxDoc()} of
 177    *        the same IndexReader with which the search is executed.
 178    * @param enableScoring if scoring is enabled, a score will be computed for
 179    *        every matching document, which might be expensive. Therefore if you
 180    *        do not require scoring, it is better to set it to <i>false</i>.
 181    */
 182   public static ScoredDocIdCollector create(int maxDoc, boolean enableScoring) {
 183     return enableScoring   ? new ScoringDocIdCollector(maxDoc)
 184                           : new NonScoringDocIdCollector(maxDoc);
 185   }
 186
 187   private ScoredDocIdCollector(int maxDoc) {
 188     numDocIds = 0;
 189     docIds = new OpenBitSet(maxDoc);
 190   }
 191
 192   /** Returns the default score used when scoring is disabled. */
 193   public abstract float getDefaultScore();
 194
 195   /** Set the default score. Only applicable if scoring is disabled. */
 196   public abstract void setDefaultScore(float defaultScore);
 197
 198   public abstract ScoredDocIDsIterator scoredDocIdsIterator() throws IOException;
 199
 200   public ScoredDocIDs getScoredDocIDs() {
 201     return new ScoredDocIDs() {
 202
 203       public ScoredDocIDsIterator iterator() throws IOException {
 204         return scoredDocIdsIterator();
 205       }
 206
 207       public DocIdSet getDocIDs() {
 208         return docIds;
 209       }
 210
 211       public int size() {
 212         return numDocIds;
 213       }
 214
 215     };
 216   }
 217
 218   @Override
 219   public void setNextReader(IndexReader reader, int base) throws IOException {
 220     this.docBase = base;
 221   }
 222
 223 }