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 org.apache.lucene.analysis.MockAnalyzer;
23 import org.apache.lucene.document.*;
24 import org.apache.lucene.index.*;
25 import org.apache.lucene.search.BooleanClause.Occur;
26 import org.apache.lucene.search.Scorer.ScorerVisitor;
27 import org.apache.lucene.store.*;
28 import org.apache.lucene.util.*;
29 import org.junit.AfterClass;
30 import org.junit.BeforeClass;
31 import org.junit.Test;
33 public class TestSubScorerFreqs extends LuceneTestCase {
35 private static Directory dir;
36 private static IndexSearcher s;
39 public static void makeIndex() throws Exception {
40 dir = new RAMDirectory();
41 RandomIndexWriter w = new RandomIndexWriter(
42 random, dir, newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random)).setMergePolicy(newLogMergePolicy()));
43 // make sure we have more than one segment occationally
44 int num = atLeast(31);
45 for (int i = 0; i < num; i++) {
46 Document doc = new Document();
47 doc.add(newField("f", "a b c d b c d c d d", Field.Store.NO,
48 Field.Index.ANALYZED));
52 doc.add(newField("f", "a b c d", Field.Store.NO, Field.Index.ANALYZED));
56 s = newSearcher(w.getReader());
61 public static void finish() throws Exception {
62 s.getIndexReader().close();
69 private static class CountingCollector extends Collector {
70 private final Collector other;
73 public final Map<Integer, Map<Query, Float>> docCounts = new HashMap<Integer, Map<Query, Float>>();
75 private final Map<Query, Scorer> subScorers = new HashMap<Query, Scorer>();
76 private final ScorerVisitor<Query, Query, Scorer> visitor = new MockScorerVisitor();
77 private final EnumSet<Occur> collect;
79 private class MockScorerVisitor extends ScorerVisitor<Query, Query, Scorer> {
82 public void visitOptional(Query parent, Query child, Scorer scorer) {
83 if (collect.contains(Occur.SHOULD))
84 subScorers.put(child, scorer);
88 public void visitProhibited(Query parent, Query child, Scorer scorer) {
89 if (collect.contains(Occur.MUST_NOT))
90 subScorers.put(child, scorer);
94 public void visitRequired(Query parent, Query child, Scorer scorer) {
95 if (collect.contains(Occur.MUST))
96 subScorers.put(child, scorer);
101 public CountingCollector(Collector other) {
102 this(other, EnumSet.allOf(Occur.class));
105 public CountingCollector(Collector other, EnumSet<Occur> collect) {
107 this.collect = collect;
111 public void setScorer(Scorer scorer) throws IOException {
112 other.setScorer(scorer);
113 scorer.visitScorers(visitor);
117 public void collect(int doc) throws IOException {
118 final Map<Query, Float> freqs = new HashMap<Query, Float>();
119 for (Map.Entry<Query, Scorer> ent : subScorers.entrySet()) {
120 Scorer value = ent.getValue();
121 int matchId = value.docID();
122 freqs.put(ent.getKey(), matchId == doc ? value.freq() : 0.0f);
124 docCounts.put(doc + docBase, freqs);
129 public void setNextReader(IndexReader reader, int docBase)
131 this.docBase = docBase;
132 other.setNextReader(reader, docBase);
136 public boolean acceptsDocsOutOfOrder() {
137 return other.acceptsDocsOutOfOrder();
141 private static final float FLOAT_TOLERANCE = 0.00001F;
144 public void testTermQuery() throws Exception {
145 TermQuery q = new TermQuery(new Term("f", "d"));
146 CountingCollector c = new CountingCollector(TopScoreDocCollector.create(10,
148 s.search(q, null, c);
149 final int maxDocs = s.maxDoc();
150 assertEquals(maxDocs, c.docCounts.size());
151 for (int i = 0; i < maxDocs; i++) {
152 Map<Query, Float> doc0 = c.docCounts.get(i);
153 assertEquals(1, doc0.size());
154 assertEquals(4.0F, doc0.get(q), FLOAT_TOLERANCE);
156 Map<Query, Float> doc1 = c.docCounts.get(++i);
157 assertEquals(1, doc1.size());
158 assertEquals(1.0F, doc1.get(q), FLOAT_TOLERANCE);
162 @SuppressWarnings("unchecked")
164 public void testBooleanQuery() throws Exception {
165 TermQuery aQuery = new TermQuery(new Term("f", "a"));
166 TermQuery dQuery = new TermQuery(new Term("f", "d"));
167 TermQuery cQuery = new TermQuery(new Term("f", "c"));
168 TermQuery yQuery = new TermQuery(new Term("f", "y"));
170 BooleanQuery query = new BooleanQuery();
171 BooleanQuery inner = new BooleanQuery();
173 inner.add(cQuery, Occur.SHOULD);
174 inner.add(yQuery, Occur.MUST_NOT);
175 query.add(inner, Occur.MUST);
176 query.add(aQuery, Occur.MUST);
177 query.add(dQuery, Occur.MUST);
178 EnumSet<Occur>[] occurList = new EnumSet[] {EnumSet.of(Occur.MUST), EnumSet.of(Occur.MUST, Occur.SHOULD)};
179 for (EnumSet<Occur> occur : occurList) {
180 CountingCollector c = new CountingCollector(TopScoreDocCollector.create(
182 s.search(query, null, c);
183 final int maxDocs = s.maxDoc();
184 assertEquals(maxDocs, c.docCounts.size());
185 boolean includeOptional = occur.contains(Occur.SHOULD);
186 for (int i = 0; i < maxDocs; i++) {
187 Map<Query, Float> doc0 = c.docCounts.get(i);
188 assertEquals(includeOptional ? 5 : 4, doc0.size());
189 assertEquals(1.0F, doc0.get(aQuery), FLOAT_TOLERANCE);
190 assertEquals(4.0F, doc0.get(dQuery), FLOAT_TOLERANCE);
192 assertEquals(3.0F, doc0.get(cQuery), FLOAT_TOLERANCE);
194 Map<Query, Float> doc1 = c.docCounts.get(++i);
195 assertEquals(includeOptional ? 5 : 4, doc1.size());
196 assertEquals(1.0F, doc1.get(aQuery), FLOAT_TOLERANCE);
197 assertEquals(1.0F, doc1.get(dQuery), FLOAT_TOLERANCE);
199 assertEquals(1.0F, doc1.get(cQuery), FLOAT_TOLERANCE);
206 public void testPhraseQuery() throws Exception {
207 PhraseQuery q = new PhraseQuery();
208 q.add(new Term("f", "b"));
209 q.add(new Term("f", "c"));
210 CountingCollector c = new CountingCollector(TopScoreDocCollector.create(10,
212 s.search(q, null, c);
213 final int maxDocs = s.maxDoc();
214 assertEquals(maxDocs, c.docCounts.size());
215 for (int i = 0; i < maxDocs; i++) {
216 Map<Query, Float> doc0 = c.docCounts.get(i);
217 assertEquals(1, doc0.size());
218 assertEquals(2.0F, doc0.get(q), FLOAT_TOLERANCE);
220 Map<Query, Float> doc1 = c.docCounts.get(++i);
221 assertEquals(1, doc1.size());
222 assertEquals(1.0F, doc1.get(q), FLOAT_TOLERANCE);