1 package org.apache.lucene.facet.search;
3 import java.io.IOException;
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;
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
21 * http://www.apache.org/licenses/LICENSE-2.0
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.
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.
36 * @lucene.experimental
38 public abstract class ScoredDocIdCollector extends Collector {
40 private static final class NonScoringDocIdCollector extends ScoredDocIdCollector {
42 float defaultScore = ScoredDocIDsIterator.DEFAULT_SCORE;
44 @SuppressWarnings("synthetic-access")
45 public NonScoringDocIdCollector(int maxDoc) {
50 public boolean acceptsDocsOutOfOrder() { return true; }
53 public void collect(int doc) throws IOException {
54 docIds.fastSet(docBase + doc);
59 public float getDefaultScore() {
64 public ScoredDocIDsIterator scoredDocIdsIterator() throws IOException {
65 return new ScoredDocIDsIterator() {
67 private DocIdSetIterator docIdsIter = docIds.iterator();
70 public int getDocID() { return nextDoc; }
71 public float getScore() { return defaultScore; }
73 public boolean next() {
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;
89 public void setDefaultScore(float defaultScore) {
90 this.defaultScore = defaultScore;
94 public void setScorer(Scorer scorer) throws IOException {}
97 private static final class ScoringDocIdCollector extends ScoredDocIdCollector {
100 private Scorer scorer;
102 @SuppressWarnings("synthetic-access")
103 public ScoringDocIdCollector(int maxDoc) {
105 scores = new float[maxDoc];
109 public boolean acceptsDocsOutOfOrder() { return false; }
112 public void collect(int doc) throws IOException {
113 docIds.fastSet(docBase + doc);
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);
121 scores[numDocIds] = score;
126 public ScoredDocIDsIterator scoredDocIdsIterator() throws IOException {
127 return new ScoredDocIDsIterator() {
129 private DocIdSetIterator docIdsIter = docIds.iterator();
131 private int scoresIdx = -1;
133 public int getDocID() { return nextDoc; }
134 public float getScore() { return scores[scoresIdx]; }
136 public boolean next() {
138 nextDoc = docIdsIter.nextDoc();
139 if (nextDoc == DocIdSetIterator.NO_MORE_DOCS) {
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;
156 public float getDefaultScore() { return ScoredDocIDsIterator.DEFAULT_SCORE; }
159 public void setDefaultScore(float defaultScore) {}
162 public void setScorer(Scorer scorer) throws IOException {
163 this.scorer = scorer;
167 protected int numDocIds;
168 protected int docBase;
169 protected final OpenBitSet docIds;
172 * Creates a new {@link ScoredDocIdCollector} with the given parameters.
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>.
182 public static ScoredDocIdCollector create(int maxDoc, boolean enableScoring) {
183 return enableScoring ? new ScoringDocIdCollector(maxDoc)
184 : new NonScoringDocIdCollector(maxDoc);
187 private ScoredDocIdCollector(int maxDoc) {
189 docIds = new OpenBitSet(maxDoc);
192 /** Returns the default score used when scoring is disabled. */
193 public abstract float getDefaultScore();
195 /** Set the default score. Only applicable if scoring is disabled. */
196 public abstract void setDefaultScore(float defaultScore);
198 public abstract ScoredDocIDsIterator scoredDocIdsIterator() throws IOException;
200 public ScoredDocIDs getScoredDocIDs() {
201 return new ScoredDocIDs() {
203 public ScoredDocIDsIterator iterator() throws IOException {
204 return scoredDocIdsIterator();
207 public DocIdSet getDocIDs() {
219 public void setNextReader(IndexReader reader, int base) throws IOException {