lucene-java-3.4.0/lucene/backwards/src/test/org/apache/lucene/search/TestElevationComparator.java

   1 package org.apache.lucene.search;
   2
   3 /**
   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
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  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.
  18  */
  19
  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;
  29 import java.util.Map;
  30
  31 public class TestElevationComparator extends LuceneTestCase {
  32
  33   private final Map<String,Integer> priority = new HashMap<String,Integer>();
  34
  35   //@Test
  36   public void testSorting() throws Throwable {
  37     Directory directory = newDirectory();
  38     IndexWriter writer = new IndexWriter(
  39         directory,
  40         newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random)).
  41             setMaxBufferedDocs(2).
  42             setMergePolicy(newLogMergePolicy(1000))
  43     );
  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"}));
  50
  51     IndexReader r = IndexReader.open(writer, true);
  52     writer.close();
  53
  54     IndexSearcher searcher = newSearcher(r);
  55
  56     runTest(searcher, true);
  57     runTest(searcher, false);
  58
  59     searcher.close();
  60     r.close();
  61     directory.close();
  62   }
  63
  64   private void runTest(IndexSearcher searcher, boolean reversed) throws Throwable {
  65
  66     BooleanQuery newq = new BooleanQuery(false);
  67     TermQuery query = new TermQuery(new Term("title", "ipod"));
  68
  69     newq.add(query, BooleanClause.Occur.SHOULD);
  70     newq.add(getElevatedQuery(new String[] {"id", "a", "id", "x"}), BooleanClause.Occur.SHOULD);
  71
  72     Sort sort = new Sort(
  73         new SortField("id", new ElevationComparatorSource(priority), false),
  74         new SortField(null, SortField.SCORE, reversed)
  75       );
  76
  77     TopDocsCollector<Entry> topCollector = TopFieldCollector.create(sort, 50, false, true, true, true);
  78     searcher.search(newq, null, topCollector);
  79
  80     TopDocs topDocs = topCollector.topDocs(0, 10);
  81     int nDocsReturned = topDocs.scoreDocs.length;
  82
  83     assertEquals(4, nDocsReturned);
  84
  85     // 0 & 3 were elevated
  86     assertEquals(0, topDocs.scoreDocs[0].doc);
  87     assertEquals(3, topDocs.scoreDocs[1].doc);
  88
  89     if (reversed) {
  90       assertEquals(2, topDocs.scoreDocs[2].doc);
  91       assertEquals(1, topDocs.scoreDocs[3].doc);
  92     } else {
  93       assertEquals(1, topDocs.scoreDocs[2].doc);
  94       assertEquals(2, topDocs.scoreDocs[3].doc);
  95     }
  96
  97     /*
  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]);
 106    }
 107     */
 108  }
 109
 110  private Query getElevatedQuery(String[] vals) {
 111    BooleanQuery q = new BooleanQuery(false);
 112    q.setBoost(0);
 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));
 118    }
 119    return q;
 120  }
 121
 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));
 126    }
 127    return doc;
 128  }
 129 }
 130
 131 class ElevationComparatorSource extends FieldComparatorSource {
 132   private final Map<String,Integer> priority;
 133
 134   public ElevationComparatorSource(final Map<String,Integer> boosts) {
 135    this.priority = boosts;
 136   }
 137
 138   @Override
 139   public FieldComparator newComparator(final String fieldname, final int numHits, int sortPos, boolean reversed) throws IOException {
 140    return new FieldComparator<Integer>() {
 141
 142      FieldCache.StringIndex idIndex;
 143      private final int[] values = new int[numHits];
 144      int bottomVal;
 145
 146      @Override
 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
 149      }
 150
 151      @Override
 152      public void setBottom(int slot) {
 153        bottomVal = values[slot];
 154      }
 155
 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();
 160      }
 161
 162      @Override
 163      public int compareBottom(int doc) throws IOException {
 164        return docVal(doc) - bottomVal;
 165      }
 166
 167      @Override
 168      public void copy(int slot, int doc) throws IOException {
 169        values[slot] = docVal(doc);
 170      }
 171
 172      @Override
 173      public void setNextReader(IndexReader reader, int docBase) throws IOException {
 174        idIndex = FieldCache.DEFAULT.getStringIndex(reader, fieldname);
 175      }
 176
 177      @Override
 178      public Integer value(int slot) {
 179        return Integer.valueOf(values[slot]);
 180      }
 181    };
 182  }
 183 }