lucene-java-3.4.0/lucene/src/test/org/apache/lucene/search/TestTermVectors.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.util.IOUtils;
  21 import org.apache.lucene.util.LuceneTestCase;
  22 import org.apache.lucene.analysis.MockAnalyzer;
  23 import org.apache.lucene.analysis.MockTokenizer;
  24 import org.apache.lucene.document.Document;
  25 import org.apache.lucene.document.Field;
  26 import org.apache.lucene.document.Field.Index;
  27 import org.apache.lucene.document.Field.Store;
  28 import org.apache.lucene.document.Field.TermVector;
  29 import org.apache.lucene.index.*;
  30 import org.apache.lucene.index.IndexWriterConfig.OpenMode;
  31 import org.apache.lucene.store.Directory;
  32 import org.apache.lucene.util.English;
  33
  34 import java.io.IOException;
  35 import java.util.HashMap;
  36 import java.util.Map;
  37 import java.util.SortedSet;
  38
  39 public class TestTermVectors extends LuceneTestCase {
  40   private IndexSearcher searcher;
  41   private IndexReader reader;
  42   private Directory directory;
  43
  44   @Override
  45   public void setUp() throws Exception {
  46     super.setUp();
  47     directory = newDirectory();
  48     RandomIndexWriter writer = new RandomIndexWriter(random, directory, newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random, MockTokenizer.SIMPLE, true)).setMergePolicy(newLogMergePolicy()));
  49     //writer.setUseCompoundFile(true);
  50     //writer.infoStream = System.out;
  51     for (int i = 0; i < 1000; i++) {
  52       Document doc = new Document();
  53       Field.TermVector termVector;
  54       int mod3 = i % 3;
  55       int mod2 = i % 2;
  56       if (mod2 == 0 && mod3 == 0){
  57         termVector = Field.TermVector.WITH_POSITIONS_OFFSETS;
  58       }
  59       else if (mod2 == 0){
  60         termVector = Field.TermVector.WITH_POSITIONS;
  61       }
  62       else if (mod3 == 0){
  63         termVector = Field.TermVector.WITH_OFFSETS;
  64       }
  65       else {
  66         termVector = Field.TermVector.YES;
  67       }
  68       doc.add(new Field("field", English.intToEnglish(i),
  69           Field.Store.YES, Field.Index.ANALYZED, termVector));
  70       writer.addDocument(doc);
  71     }
  72     reader = writer.getReader();
  73     writer.close();
  74     searcher = newSearcher(reader);
  75   }
  76
  77   @Override
  78   public void tearDown() throws Exception {
  79     searcher.close();
  80     reader.close();
  81     directory.close();
  82     super.tearDown();
  83   }
  84
  85   public void test() {
  86     assertTrue(searcher != null);
  87   }
  88
  89   public void testTermVectors() {
  90     Query query = new TermQuery(new Term("field", "seventy"));
  91     try {
  92       ScoreDoc[] hits = searcher.search(query, null, 1000).scoreDocs;
  93       assertEquals(100, hits.length);
  94
  95       for (int i = 0; i < hits.length; i++)
  96       {
  97         TermFreqVector [] vector = searcher.reader.getTermFreqVectors(hits[i].doc);
  98         assertTrue(vector != null);
  99         assertTrue(vector.length == 1);
 100       }
 101     } catch (IOException e) {
 102       assertTrue(false);
 103     }
 104   }
 105
 106   public void testTermVectorsFieldOrder() throws IOException {
 107     Directory dir = newDirectory();
 108     RandomIndexWriter writer = new RandomIndexWriter(random, dir, new MockAnalyzer(random, MockTokenizer.SIMPLE, true));
 109     Document doc = new Document();
 110     doc.add(new Field("c", "some content here", Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 111     doc.add(new Field("a", "some content here", Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 112     doc.add(new Field("b", "some content here", Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 113     doc.add(new Field("x", "some content here", Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 114     writer.addDocument(doc);
 115     IndexReader reader = writer.getReader();
 116     writer.close();
 117     TermFreqVector[] v = reader.getTermFreqVectors(0);
 118     assertEquals(4, v.length);
 119     String[] expectedFields = new String[]{"a", "b", "c", "x"};
 120     int[] expectedPositions = new int[]{1, 2, 0};
 121     for(int i=0;i<v.length;i++) {
 122       TermPositionVector posVec = (TermPositionVector) v[i];
 123       assertEquals(expectedFields[i], posVec.getField());
 124       String[] terms = posVec.getTerms();
 125       assertEquals(3, terms.length);
 126       assertEquals("content", terms[0]);
 127       assertEquals("here", terms[1]);
 128       assertEquals("some", terms[2]);
 129       for(int j=0;j<3;j++) {
 130         int[] positions = posVec.getTermPositions(j);
 131         assertEquals(1, positions.length);
 132         assertEquals(expectedPositions[j], positions[0]);
 133       }
 134     }
 135     reader.close();
 136     dir.close();
 137   }
 138
 139   public void testTermPositionVectors() throws IOException {
 140     Query query = new TermQuery(new Term("field", "zero"));
 141     ScoreDoc[] hits = searcher.search(query, null, 1000).scoreDocs;
 142     assertEquals(1, hits.length);
 143
 144     for (int i = 0; i < hits.length; i++)
 145     {
 146       TermFreqVector [] vector = searcher.reader.getTermFreqVectors(hits[i].doc);
 147       assertTrue(vector != null);
 148       assertTrue(vector.length == 1);
 149
 150       boolean shouldBePosVector = (hits[i].doc % 2 == 0) ? true : false;
 151       assertTrue((shouldBePosVector == false) || (shouldBePosVector == true && (vector[0] instanceof TermPositionVector == true)));
 152
 153       boolean shouldBeOffVector = (hits[i].doc % 3 == 0) ? true : false;
 154       assertTrue((shouldBeOffVector == false) || (shouldBeOffVector == true && (vector[0] instanceof TermPositionVector == true)));
 155
 156       if(shouldBePosVector || shouldBeOffVector){
 157         TermPositionVector posVec = (TermPositionVector)vector[0];
 158         String [] terms = posVec.getTerms();
 159         assertTrue(terms != null && terms.length > 0);
 160
 161         for (int j = 0; j < terms.length; j++) {
 162           int [] positions = posVec.getTermPositions(j);
 163           TermVectorOffsetInfo [] offsets = posVec.getOffsets(j);
 164
 165           if(shouldBePosVector){
 166             assertTrue(positions != null);
 167             assertTrue(positions.length > 0);
 168           }
 169           else
 170             assertTrue(positions == null);
 171
 172           if(shouldBeOffVector){
 173             assertTrue(offsets != null);
 174             assertTrue(offsets.length > 0);
 175           }
 176           else
 177             assertTrue(offsets == null);
 178         }
 179       }
 180       else{
 181         try{
 182           assertTrue(false);
 183         }
 184         catch(ClassCastException ignore){
 185           TermFreqVector freqVec = vector[0];
 186           String [] terms = freqVec.getTerms();
 187           assertTrue(terms != null && terms.length > 0);
 188         }
 189
 190       }
 191
 192     }
 193   }
 194
 195   public void testTermOffsetVectors() {
 196     Query query = new TermQuery(new Term("field", "fifty"));
 197     try {
 198       ScoreDoc[] hits = searcher.search(query, null, 1000).scoreDocs;
 199       assertEquals(100, hits.length);
 200
 201       for (int i = 0; i < hits.length; i++)
 202       {
 203         TermFreqVector [] vector = searcher.reader.getTermFreqVectors(hits[i].doc);
 204         assertTrue(vector != null);
 205         assertTrue(vector.length == 1);
 206
 207         //assertTrue();
 208       }
 209     } catch (IOException e) {
 210       assertTrue(false);
 211     }
 212   }
 213
 214   public void testKnownSetOfDocuments() throws IOException {
 215     String test1 = "eating chocolate in a computer lab"; //6 terms
 216     String test2 = "computer in a computer lab"; //5 terms
 217     String test3 = "a chocolate lab grows old"; //5 terms
 218     String test4 = "eating chocolate with a chocolate lab in an old chocolate colored computer lab"; //13 terms
 219     Map<String,Integer> test4Map = new HashMap<String,Integer>();
 220     test4Map.put("chocolate", Integer.valueOf(3));
 221     test4Map.put("lab", Integer.valueOf(2));
 222     test4Map.put("eating", Integer.valueOf(1));
 223     test4Map.put("computer", Integer.valueOf(1));
 224     test4Map.put("with", Integer.valueOf(1));
 225     test4Map.put("a", Integer.valueOf(1));
 226     test4Map.put("colored", Integer.valueOf(1));
 227     test4Map.put("in", Integer.valueOf(1));
 228     test4Map.put("an", Integer.valueOf(1));
 229     test4Map.put("computer", Integer.valueOf(1));
 230     test4Map.put("old", Integer.valueOf(1));
 231
 232     Document testDoc1 = new Document();
 233     setupDoc(testDoc1, test1);
 234     Document testDoc2 = new Document();
 235     setupDoc(testDoc2, test2);
 236     Document testDoc3 = new Document();
 237     setupDoc(testDoc3, test3);
 238     Document testDoc4 = new Document();
 239     setupDoc(testDoc4, test4);
 240
 241     Directory dir = newDirectory();
 242
 243     RandomIndexWriter writer = new RandomIndexWriter(random, dir,
 244         newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random, MockTokenizer.SIMPLE, true))
 245                                                      .setOpenMode(OpenMode.CREATE).setMergePolicy(newLogMergePolicy()));
 246     writer.addDocument(testDoc1);
 247     writer.addDocument(testDoc2);
 248     writer.addDocument(testDoc3);
 249     writer.addDocument(testDoc4);
 250     IndexReader reader = writer.getReader();
 251     writer.close();
 252     IndexSearcher knownSearcher = newSearcher(reader);
 253     TermEnum termEnum = knownSearcher.reader.terms();
 254     TermDocs termDocs = knownSearcher.reader.termDocs();
 255     //System.out.println("Terms: " + termEnum.size() + " Orig Len: " + termArray.length);
 256
 257     //Similarity sim = knownSearcher.getSimilarity();
 258     while (termEnum.next() == true)
 259     {
 260       Term term = termEnum.term();
 261       //System.out.println("Term: " + term);
 262       termDocs.seek(term);
 263       while (termDocs.next())
 264       {
 265         int docId = termDocs.doc();
 266         int freq = termDocs.freq();
 267         //System.out.println("Doc Id: " + docId + " freq " + freq);
 268         TermFreqVector vector = knownSearcher.reader.getTermFreqVector(docId, "field");
 269         //float tf = sim.tf(freq);
 270         //float idf = sim.idf(knownSearcher.docFreq(term), knownSearcher.maxDoc());
 271         //float qNorm = sim.queryNorm()
 272         //This is fine since we don't have stop words
 273         //float lNorm = sim.lengthNorm("field", vector.getTerms().length);
 274         //float coord = sim.coord()
 275         //System.out.println("TF: " + tf + " IDF: " + idf + " LenNorm: " + lNorm);
 276         assertTrue(vector != null);
 277         String[] vTerms = vector.getTerms();
 278         int [] freqs = vector.getTermFrequencies();
 279         for (int i = 0; i < vTerms.length; i++)
 280         {
 281           if (term.text().equals(vTerms[i]))
 282           {
 283             assertTrue(freqs[i] == freq);
 284           }
 285         }
 286
 287       }
 288       //System.out.println("--------");
 289     }
 290     Query query = new TermQuery(new Term("field", "chocolate"));
 291     ScoreDoc[] hits = knownSearcher.search(query, null, 1000).scoreDocs;
 292     //doc 3 should be the first hit b/c it is the shortest match
 293     assertTrue(hits.length == 3);
 294     /*System.out.println("Hit 0: " + hits.id(0) + " Score: " + hits.score(0) + " String: " + hits.doc(0).toString());
 295       System.out.println("Explain: " + knownSearcher.explain(query, hits.id(0)));
 296       System.out.println("Hit 1: " + hits.id(1) + " Score: " + hits.score(1) + " String: " + hits.doc(1).toString());
 297       System.out.println("Explain: " + knownSearcher.explain(query, hits.id(1)));
 298       System.out.println("Hit 2: " + hits.id(2) + " Score: " + hits.score(2) + " String: " +  hits.doc(2).toString());
 299       System.out.println("Explain: " + knownSearcher.explain(query, hits.id(2)));*/
 300     assertTrue(hits[0].doc == 2);
 301     assertTrue(hits[1].doc == 3);
 302     assertTrue(hits[2].doc == 0);
 303     TermFreqVector vector = knownSearcher.reader.getTermFreqVector(hits[1].doc, "field");
 304     assertTrue(vector != null);
 305     //System.out.println("Vector: " + vector);
 306     String[] terms = vector.getTerms();
 307     int [] freqs = vector.getTermFrequencies();
 308     assertTrue(terms != null && terms.length == 10);
 309     for (int i = 0; i < terms.length; i++) {
 310       String term = terms[i];
 311       //System.out.println("Term: " + term);
 312       int freq = freqs[i];
 313       assertTrue(test4.indexOf(term) != -1);
 314       Integer freqInt = test4Map.get(term);
 315       assertTrue(freqInt != null);
 316       assertTrue(freqInt.intValue() == freq);
 317     }
 318     SortedTermVectorMapper mapper = new SortedTermVectorMapper(new TermVectorEntryFreqSortedComparator());
 319     knownSearcher.reader.getTermFreqVector(hits[1].doc, mapper);
 320     SortedSet<TermVectorEntry> vectorEntrySet = mapper.getTermVectorEntrySet();
 321     assertTrue("mapper.getTermVectorEntrySet() Size: " + vectorEntrySet.size() + " is not: " + 10, vectorEntrySet.size() == 10);
 322     TermVectorEntry last = null;
 323     for (final TermVectorEntry tve : vectorEntrySet) {
 324       if (tve != null && last != null)
 325       {
 326         assertTrue("terms are not properly sorted", last.getFrequency() >= tve.getFrequency());
 327         Integer expectedFreq =  test4Map.get(tve.getTerm());
 328         //we expect double the expectedFreq, since there are two fields with the exact same text and we are collapsing all fields
 329         assertTrue("Frequency is not correct:", tve.getFrequency() == 2*expectedFreq.intValue());
 330       }
 331       last = tve;
 332
 333     }
 334
 335     FieldSortedTermVectorMapper fieldMapper = new FieldSortedTermVectorMapper(new TermVectorEntryFreqSortedComparator());
 336     knownSearcher.reader.getTermFreqVector(hits[1].doc, fieldMapper);
 337     Map<String,SortedSet<TermVectorEntry>> map = fieldMapper.getFieldToTerms();
 338     assertTrue("map Size: " + map.size() + " is not: " + 2, map.size() == 2);
 339     vectorEntrySet = map.get("field");
 340     assertTrue("vectorEntrySet is null and it shouldn't be", vectorEntrySet != null);
 341     assertTrue("vectorEntrySet Size: " + vectorEntrySet.size() + " is not: " + 10, vectorEntrySet.size() == 10);
 342     knownSearcher.close();
 343     reader.close();
 344     dir.close();
 345   }
 346
 347   private void setupDoc(Document doc, String text)
 348   {
 349     doc.add(new Field("field2", text, Field.Store.YES,
 350         Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 351     doc.add(new Field("field", text, Field.Store.YES,
 352         Field.Index.ANALYZED, Field.TermVector.YES));
 353     //System.out.println("Document: " + doc);
 354   }
 355
 356   // Test only a few docs having vectors
 357   public void testRareVectors() throws IOException {
 358     RandomIndexWriter writer = new RandomIndexWriter(random, directory,
 359         newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random, MockTokenizer.SIMPLE, true))
 360         .setOpenMode(OpenMode.CREATE));
 361     writer.w.setInfoStream(VERBOSE ? System.out : null);
 362     if (VERBOSE) {
 363       System.out.println("TEST: now add non-vectors");
 364     }
 365     for (int i = 0; i < 100; i++) {
 366       Document doc = new Document();
 367       doc.add(new Field("field", English.intToEnglish(i),
 368                         Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.NO));
 369       writer.addDocument(doc);
 370     }
 371     if (VERBOSE) {
 372       System.out.println("TEST: now add vectors");
 373     }
 374     for(int i=0;i<10;i++) {
 375       Document doc = new Document();
 376       doc.add(new Field("field", English.intToEnglish(100+i),
 377                         Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 378       writer.addDocument(doc);
 379     }
 380
 381     if (VERBOSE) {
 382       System.out.println("TEST: now getReader");
 383     }
 384     IndexReader reader = writer.getReader();
 385     writer.close();
 386     searcher = newSearcher(reader);
 387
 388     Query query = new TermQuery(new Term("field", "hundred"));
 389     ScoreDoc[] hits = searcher.search(query, null, 1000).scoreDocs;
 390     assertEquals(10, hits.length);
 391     for (int i = 0; i < hits.length; i++) {
 392
 393       TermFreqVector [] vector = searcher.reader.getTermFreqVectors(hits[i].doc);
 394       assertTrue(vector != null);
 395       assertTrue(vector.length == 1);
 396     }
 397     reader.close();
 398   }
 399
 400
 401   // In a single doc, for the same field, mix the term
 402   // vectors up
 403   public void testMixedVectrosVectors() throws IOException {
 404     RandomIndexWriter writer = new RandomIndexWriter(random, directory,
 405         newIndexWriterConfig(TEST_VERSION_CURRENT,
 406         new MockAnalyzer(random, MockTokenizer.SIMPLE, true)).setOpenMode(OpenMode.CREATE));
 407     Document doc = new Document();
 408     doc.add(new Field("field", "one",
 409                       Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.NO));
 410     doc.add(new Field("field", "one",
 411                       Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.YES));
 412     doc.add(new Field("field", "one",
 413                       Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS));
 414     doc.add(new Field("field", "one",
 415                       Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_OFFSETS));
 416     doc.add(new Field("field", "one",
 417                       Field.Store.YES, Field.Index.ANALYZED, Field.TermVector.WITH_POSITIONS_OFFSETS));
 418     writer.addDocument(doc);
 419     IndexReader reader = writer.getReader();
 420     writer.close();
 421
 422     searcher = newSearcher(reader);
 423
 424     Query query = new TermQuery(new Term("field", "one"));
 425     ScoreDoc[] hits = searcher.search(query, null, 1000).scoreDocs;
 426     assertEquals(1, hits.length);
 427
 428     TermFreqVector [] vector = searcher.reader.getTermFreqVectors(hits[0].doc);
 429     assertTrue(vector != null);
 430     assertTrue(vector.length == 1);
 431     TermPositionVector tfv = (TermPositionVector) vector[0];
 432     assertTrue(tfv.getField().equals("field"));
 433     String[] terms = tfv.getTerms();
 434     assertEquals(1, terms.length);
 435     assertEquals(terms[0], "one");
 436     assertEquals(5, tfv.getTermFrequencies()[0]);
 437
 438     int[] positions = tfv.getTermPositions(0);
 439     assertEquals(5, positions.length);
 440     for(int i=0;i<5;i++)
 441       assertEquals(i, positions[i]);
 442     TermVectorOffsetInfo[] offsets = tfv.getOffsets(0);
 443     assertEquals(5, offsets.length);
 444     for(int i=0;i<5;i++) {
 445       assertEquals(4*i, offsets[i].getStartOffset());
 446       assertEquals(4*i+3, offsets[i].getEndOffset());
 447     }
 448     reader.close();
 449   }
 450
 451   private IndexWriter createWriter(Directory dir) throws IOException {
 452     return new IndexWriter(dir, newIndexWriterConfig(TEST_VERSION_CURRENT,
 453         new MockAnalyzer(random)).setMaxBufferedDocs(2));
 454   }
 455
 456   private void createDir(Directory dir) throws IOException {
 457     IndexWriter writer = createWriter(dir);
 458     writer.addDocument(createDoc());
 459     writer.close();
 460   }
 461
 462   private Document createDoc() {
 463     Document doc = new Document();
 464     doc.add(new Field("c", "aaa", Store.YES, Index.ANALYZED, TermVector.WITH_POSITIONS_OFFSETS));
 465     return doc;
 466   }
 467
 468   private void verifyIndex(Directory dir) throws IOException {
 469     IndexReader r = IndexReader.open(dir);
 470     int numDocs = r.numDocs();
 471     for (int i = 0; i < numDocs; i++) {
 472       TermFreqVector tfv = r.getTermFreqVector(i, "c");
 473       assertNotNull("term vectors should not have been null for document " + i, tfv);
 474     }
 475     r.close();
 476   }
 477
 478   public void testOptimizeAddDocs() throws Exception {
 479     Directory target = newDirectory();
 480     IndexWriter writer = createWriter(target);
 481     // with maxBufferedDocs=2, this results in two segments, so that optimize
 482     // actually does something.
 483     for (int i = 0; i < 4; i++) {
 484       writer.addDocument(createDoc());
 485     }
 486     writer.optimize();
 487     writer.close();
 488
 489     verifyIndex(target);
 490     target.close();
 491   }
 492
 493   public void testOptimizeAddIndexesDir() throws Exception {
 494     Directory[] input = new Directory[] { newDirectory(), newDirectory() };
 495     Directory target = newDirectory();
 496
 497     for (Directory dir : input) {
 498       createDir(dir);
 499     }
 500
 501     IndexWriter writer = createWriter(target);
 502     writer.addIndexes(input);
 503     writer.optimize();
 504     writer.close();
 505
 506     verifyIndex(target);
 507
 508     IOUtils.close(target, input[0], input[1]);
 509   }
 510
 511   public void testOptimizeAddIndexesReader() throws Exception {
 512     Directory[] input = new Directory[] { newDirectory(), newDirectory() };
 513     Directory target = newDirectory();
 514
 515     for (Directory dir : input) {
 516       createDir(dir);
 517     }
 518
 519     IndexWriter writer = createWriter(target);
 520     for (Directory dir : input) {
 521       IndexReader r = IndexReader.open(dir);
 522       writer.addIndexes(r);
 523       r.close();
 524     }
 525     writer.optimize();
 526     writer.close();
 527
 528     verifyIndex(target);
 529     IOUtils.close(target, input[0], input[1]);
 530   }
 531
 532 }