--- /dev/null
+package org.apache.lucene.search;
+
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import org.apache.lucene.analysis.MockAnalyzer;
+import java.io.*;
+import java.util.*;
+import org.apache.lucene.document.*;
+import org.apache.lucene.index.*;
+import org.apache.lucene.search.BooleanClause.Occur;
+import org.apache.lucene.search.Scorer.ScorerVisitor;
+import org.apache.lucene.store.*;
+import org.apache.lucene.util.*;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public class TestSubScorerFreqs extends LuceneTestCase {
+
+ private static Directory dir;
+ private static IndexSearcher s;
+
+ @BeforeClass
+ public static void makeIndex() throws Exception {
+ dir = new RAMDirectory();
+ RandomIndexWriter w = new RandomIndexWriter(
+ random, dir, newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random)).setMergePolicy(newLogMergePolicy()));
+ // make sure we have more than one segment occationally
+ int num = atLeast(31);
+ for (int i = 0; i < num; i++) {
+ Document doc = new Document();
+ doc.add(newField("f", "a b c d b c d c d d", Field.Store.NO,
+ Field.Index.ANALYZED));
+ w.addDocument(doc);
+
+ doc = new Document();
+ doc.add(newField("f", "a b c d", Field.Store.NO, Field.Index.ANALYZED));
+ w.addDocument(doc);
+ }
+
+ s = newSearcher(w.getReader());
+ w.close();
+ }
+
+ @AfterClass
+ public static void finish() throws Exception {
+ s.getIndexReader().close();
+ s.close();
+ s = null;
+ dir.close();
+ dir = null;
+ }
+
+ private static class CountingCollector extends Collector {
+ private final Collector other;
+ private int docBase;
+
+ public final Map<Integer, Map<Query, Float>> docCounts = new HashMap<Integer, Map<Query, Float>>();
+
+ private final Map<Query, Scorer> subScorers = new HashMap<Query, Scorer>();
+ private final ScorerVisitor<Query, Query, Scorer> visitor = new MockScorerVisitor();
+ private final EnumSet<Occur> collect;
+
+ private class MockScorerVisitor extends ScorerVisitor<Query, Query, Scorer> {
+
+ @Override
+ public void visitOptional(Query parent, Query child, Scorer scorer) {
+ if (collect.contains(Occur.SHOULD))
+ subScorers.put(child, scorer);
+ }
+
+ @Override
+ public void visitProhibited(Query parent, Query child, Scorer scorer) {
+ if (collect.contains(Occur.MUST_NOT))
+ subScorers.put(child, scorer);
+ }
+
+ @Override
+ public void visitRequired(Query parent, Query child, Scorer scorer) {
+ if (collect.contains(Occur.MUST))
+ subScorers.put(child, scorer);
+ }
+
+ }
+
+ public CountingCollector(Collector other) {
+ this(other, EnumSet.allOf(Occur.class));
+ }
+
+ public CountingCollector(Collector other, EnumSet<Occur> collect) {
+ this.other = other;
+ this.collect = collect;
+ }
+
+ @Override
+ public void setScorer(Scorer scorer) throws IOException {
+ other.setScorer(scorer);
+ scorer.visitScorers(visitor);
+ }
+
+ @Override
+ public void collect(int doc) throws IOException {
+ final Map<Query, Float> freqs = new HashMap<Query, Float>();
+ for (Map.Entry<Query, Scorer> ent : subScorers.entrySet()) {
+ Scorer value = ent.getValue();
+ int matchId = value.docID();
+ freqs.put(ent.getKey(), matchId == doc ? value.freq() : 0.0f);
+ }
+ docCounts.put(doc + docBase, freqs);
+ other.collect(doc);
+ }
+
+ @Override
+ public void setNextReader(IndexReader reader, int docBase)
+ throws IOException {
+ this.docBase = docBase;
+ other.setNextReader(reader, docBase);
+ }
+
+ @Override
+ public boolean acceptsDocsOutOfOrder() {
+ return other.acceptsDocsOutOfOrder();
+ }
+ }
+
+ private static final float FLOAT_TOLERANCE = 0.00001F;
+
+ @Test
+ public void testTermQuery() throws Exception {
+ TermQuery q = new TermQuery(new Term("f", "d"));
+ CountingCollector c = new CountingCollector(TopScoreDocCollector.create(10,
+ true));
+ s.search(q, null, c);
+ final int maxDocs = s.maxDoc();
+ assertEquals(maxDocs, c.docCounts.size());
+ for (int i = 0; i < maxDocs; i++) {
+ Map<Query, Float> doc0 = c.docCounts.get(i);
+ assertEquals(1, doc0.size());
+ assertEquals(4.0F, doc0.get(q), FLOAT_TOLERANCE);
+
+ Map<Query, Float> doc1 = c.docCounts.get(++i);
+ assertEquals(1, doc1.size());
+ assertEquals(1.0F, doc1.get(q), FLOAT_TOLERANCE);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ @Test
+ public void testBooleanQuery() throws Exception {
+ TermQuery aQuery = new TermQuery(new Term("f", "a"));
+ TermQuery dQuery = new TermQuery(new Term("f", "d"));
+ TermQuery cQuery = new TermQuery(new Term("f", "c"));
+ TermQuery yQuery = new TermQuery(new Term("f", "y"));
+
+ BooleanQuery query = new BooleanQuery();
+ BooleanQuery inner = new BooleanQuery();
+
+ inner.add(cQuery, Occur.SHOULD);
+ inner.add(yQuery, Occur.MUST_NOT);
+ query.add(inner, Occur.MUST);
+ query.add(aQuery, Occur.MUST);
+ query.add(dQuery, Occur.MUST);
+ EnumSet<Occur>[] occurList = new EnumSet[] {EnumSet.of(Occur.MUST), EnumSet.of(Occur.MUST, Occur.SHOULD)};
+ for (EnumSet<Occur> occur : occurList) {
+ CountingCollector c = new CountingCollector(TopScoreDocCollector.create(
+ 10, true), occur);
+ s.search(query, null, c);
+ final int maxDocs = s.maxDoc();
+ assertEquals(maxDocs, c.docCounts.size());
+ boolean includeOptional = occur.contains(Occur.SHOULD);
+ for (int i = 0; i < maxDocs; i++) {
+ Map<Query, Float> doc0 = c.docCounts.get(i);
+ assertEquals(includeOptional ? 5 : 4, doc0.size());
+ assertEquals(1.0F, doc0.get(aQuery), FLOAT_TOLERANCE);
+ assertEquals(4.0F, doc0.get(dQuery), FLOAT_TOLERANCE);
+ if (includeOptional)
+ assertEquals(3.0F, doc0.get(cQuery), FLOAT_TOLERANCE);
+
+ Map<Query, Float> doc1 = c.docCounts.get(++i);
+ assertEquals(includeOptional ? 5 : 4, doc1.size());
+ assertEquals(1.0F, doc1.get(aQuery), FLOAT_TOLERANCE);
+ assertEquals(1.0F, doc1.get(dQuery), FLOAT_TOLERANCE);
+ if (includeOptional)
+ assertEquals(1.0F, doc1.get(cQuery), FLOAT_TOLERANCE);
+
+ }
+ }
+ }
+
+ @Test
+ public void testPhraseQuery() throws Exception {
+ PhraseQuery q = new PhraseQuery();
+ q.add(new Term("f", "b"));
+ q.add(new Term("f", "c"));
+ CountingCollector c = new CountingCollector(TopScoreDocCollector.create(10,
+ true));
+ s.search(q, null, c);
+ final int maxDocs = s.maxDoc();
+ assertEquals(maxDocs, c.docCounts.size());
+ for (int i = 0; i < maxDocs; i++) {
+ Map<Query, Float> doc0 = c.docCounts.get(i);
+ assertEquals(1, doc0.size());
+ assertEquals(2.0F, doc0.get(q), FLOAT_TOLERANCE);
+
+ Map<Query, Float> doc1 = c.docCounts.get(++i);
+ assertEquals(1, doc1.size());
+ assertEquals(1.0F, doc1.get(q), FLOAT_TOLERANCE);
+ }
+
+ }
+}
fnp / pylucene.git / blobdiff