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;
21 import org.apache.lucene.document.Document;
22 import org.apache.lucene.document.Field;
23 import org.apache.lucene.index.*;
24 import org.apache.lucene.search.FieldValueHitQueue.Entry;
25 import org.apache.lucene.store.*;
26 import org.apache.lucene.util.LuceneTestCase;
27 import java.io.IOException;
28 import java.util.HashMap;
31 public class TestElevationComparator extends LuceneTestCase {
33 private final Map<String,Integer> priority = new HashMap<String,Integer>();
36 public void testSorting() throws Throwable {
37 Directory directory = newDirectory();
38 IndexWriter writer = new IndexWriter(
40 newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random)).
41 setMaxBufferedDocs(2).
42 setMergePolicy(newLogMergePolicy(1000))
44 writer.addDocument(adoc(new String[] {"id", "a", "title", "ipod", "str_s", "a"}));
45 writer.addDocument(adoc(new String[] {"id", "b", "title", "ipod ipod", "str_s", "b"}));
46 writer.addDocument(adoc(new String[] {"id", "c", "title", "ipod ipod ipod", "str_s","c"}));
47 writer.addDocument(adoc(new String[] {"id", "x", "title", "boosted", "str_s", "x"}));
48 writer.addDocument(adoc(new String[] {"id", "y", "title", "boosted boosted", "str_s","y"}));
49 writer.addDocument(adoc(new String[] {"id", "z", "title", "boosted boosted boosted","str_s", "z"}));
51 IndexReader r = IndexReader.open(writer, true);
54 IndexSearcher searcher = newSearcher(r);
56 runTest(searcher, true);
57 runTest(searcher, false);
64 private void runTest(IndexSearcher searcher, boolean reversed) throws Throwable {
66 BooleanQuery newq = new BooleanQuery(false);
67 TermQuery query = new TermQuery(new Term("title", "ipod"));
69 newq.add(query, BooleanClause.Occur.SHOULD);
70 newq.add(getElevatedQuery(new String[] {"id", "a", "id", "x"}), BooleanClause.Occur.SHOULD);
73 new SortField("id", new ElevationComparatorSource(priority), false),
74 new SortField(null, SortField.SCORE, reversed)
77 TopDocsCollector<Entry> topCollector = TopFieldCollector.create(sort, 50, false, true, true, true);
78 searcher.search(newq, null, topCollector);
80 TopDocs topDocs = topCollector.topDocs(0, 10);
81 int nDocsReturned = topDocs.scoreDocs.length;
83 assertEquals(4, nDocsReturned);
85 // 0 & 3 were elevated
86 assertEquals(0, topDocs.scoreDocs[0].doc);
87 assertEquals(3, topDocs.scoreDocs[1].doc);
90 assertEquals(2, topDocs.scoreDocs[2].doc);
91 assertEquals(1, topDocs.scoreDocs[3].doc);
93 assertEquals(1, topDocs.scoreDocs[2].doc);
94 assertEquals(2, topDocs.scoreDocs[3].doc);
98 for (int i = 0; i < nDocsReturned; i++) {
99 ScoreDoc scoreDoc = topDocs.scoreDocs[i];
100 ids[i] = scoreDoc.doc;
101 scores[i] = scoreDoc.score;
102 documents[i] = searcher.doc(ids[i]);
103 System.out.println("ids[i] = " + ids[i]);
104 System.out.println("documents[i] = " + documents[i]);
105 System.out.println("scores[i] = " + scores[i]);
110 private Query getElevatedQuery(String[] vals) {
111 BooleanQuery q = new BooleanQuery(false);
113 int max = (vals.length / 2) + 5;
114 for (int i = 0; i < vals.length - 1; i += 2) {
115 q.add(new TermQuery(new Term(vals[i], vals[i + 1])), BooleanClause.Occur.SHOULD);
116 priority.put(vals[i + 1], Integer.valueOf(max--));
117 // System.out.println(" pri doc=" + vals[i+1] + " pri=" + (1+max));
122 private Document adoc(String[] vals) {
123 Document doc = new Document();
124 for (int i = 0; i < vals.length - 2; i += 2) {
125 doc.add(newField(vals[i], vals[i + 1], Field.Store.YES, Field.Index.ANALYZED));
131 class ElevationComparatorSource extends FieldComparatorSource {
132 private final Map<String,Integer> priority;
134 public ElevationComparatorSource(final Map<String,Integer> boosts) {
135 this.priority = boosts;
139 public FieldComparator newComparator(final String fieldname, final int numHits, int sortPos, boolean reversed) throws IOException {
140 return new FieldComparator<Integer>() {
142 FieldCache.StringIndex idIndex;
143 private final int[] values = new int[numHits];
147 public int compare(int slot1, int slot2) {
148 return values[slot2] - values[slot1]; // values will be small enough that there is no overflow concern
152 public void setBottom(int slot) {
153 bottomVal = values[slot];
156 private int docVal(int doc) throws IOException {
157 String id = idIndex.lookup[idIndex.order[doc]];
158 Integer prio = priority.get(id);
159 return prio == null ? 0 : prio.intValue();
163 public int compareBottom(int doc) throws IOException {
164 return docVal(doc) - bottomVal;
168 public void copy(int slot, int doc) throws IOException {
169 values[slot] = docVal(doc);
173 public void setNextReader(IndexReader reader, int docBase) throws IOException {
174 idIndex = FieldCache.DEFAULT.getStringIndex(reader, fieldname);
178 public Integer value(int slot) {
179 return Integer.valueOf(values[slot]);