Esempio n. 1
0
 /**
  * Returns a new TopSuggestDocs, containing topN results across the provided TopSuggestDocs,
  * sorting by score. Each {@link TopSuggestDocs} instance must be sorted. Analogous to {@link
  * org.apache.lucene.search.TopDocs#merge(int, org.apache.lucene.search.TopDocs[])} for {@link
  * TopSuggestDocs}
  *
  * <p>NOTE: assumes every <code>shardHit</code> is already sorted by score
  */
 public static TopSuggestDocs merge(int topN, TopSuggestDocs[] shardHits) {
   SuggestScoreDocPriorityQueue priorityQueue = new SuggestScoreDocPriorityQueue(topN);
   for (TopSuggestDocs shardHit : shardHits) {
     for (SuggestScoreDoc scoreDoc : shardHit.scoreLookupDocs()) {
       if (scoreDoc == priorityQueue.insertWithOverflow(scoreDoc)) {
         break;
       }
     }
   }
   SuggestScoreDoc[] topNResults = priorityQueue.getResults();
   if (topNResults.length > 0) {
     return new TopSuggestDocs(topNResults.length, topNResults, topNResults[0].score);
   } else {
     return TopSuggestDocs.EMPTY;
   }
 }
  @Test
  public void testScoring() throws Exception {
    Analyzer analyzer = new MockAnalyzer(random());
    RandomIndexWriter iw =
        new RandomIndexWriter(random(), dir, iwcWithSuggestField(analyzer, "suggest_field"));

    int num = Math.min(1000, atLeast(100));
    String[] prefixes = {"abc", "bac", "cab"};
    Map<String, Integer> mappings = new HashMap<>();
    for (int i = 0; i < num; i++) {
      Document document = new Document();
      String suggest =
          prefixes[i % 3] + TestUtil.randomSimpleString(random(), 10) + "_" + String.valueOf(i);
      int weight = Math.abs(random().nextInt());
      document.add(new SuggestField("suggest_field", suggest, weight));
      mappings.put(suggest, weight);
      iw.addDocument(document);

      if (usually()) {
        iw.commit();
      }
    }

    DirectoryReader reader = iw.getReader();
    SuggestIndexSearcher indexSearcher = new SuggestIndexSearcher(reader);
    for (String prefix : prefixes) {
      PrefixCompletionQuery query =
          new PrefixCompletionQuery(analyzer, new Term("suggest_field", prefix));
      TopSuggestDocs suggest = indexSearcher.suggest(query, num);
      assertTrue(suggest.totalHits > 0);
      float topScore = -1;
      for (SuggestScoreDoc scoreDoc : suggest.scoreLookupDocs()) {
        if (topScore != -1) {
          assertTrue(topScore >= scoreDoc.score);
        }
        topScore = scoreDoc.score;
        assertThat((float) mappings.get(scoreDoc.key.toString()), equalTo(scoreDoc.score));
        assertNotNull(mappings.remove(scoreDoc.key.toString()));
      }
    }

    assertThat(mappings.size(), equalTo(0));
    reader.close();
    iw.close();
  }
  @Test
  public void testReturnedDocID() throws Exception {
    Analyzer analyzer = new MockAnalyzer(random());
    RandomIndexWriter iw =
        new RandomIndexWriter(random(), dir, iwcWithSuggestField(analyzer, "suggest_field"));

    int num = Math.min(1000, atLeast(10));
    for (int i = 0; i < num; i++) {
      Document document = new Document();
      document.add(new SuggestField("suggest_field", "abc_" + i, num));
      document.add(new StoredField("int_field", i));
      iw.addDocument(document);

      if (random().nextBoolean()) {
        iw.commit();
      }
    }

    DirectoryReader reader = iw.getReader();
    SuggestIndexSearcher indexSearcher = new SuggestIndexSearcher(reader);
    PrefixCompletionQuery query =
        new PrefixCompletionQuery(analyzer, new Term("suggest_field", "abc_"));
    TopSuggestDocs suggest = indexSearcher.suggest(query, num);
    assertEquals(num, suggest.totalHits);
    for (SuggestScoreDoc suggestScoreDoc : suggest.scoreLookupDocs()) {
      String key = suggestScoreDoc.key.toString();
      assertTrue(key.startsWith("abc_"));
      String substring = key.substring(4);
      int fieldValue = Integer.parseInt(substring);
      StoredDocument doc = reader.document(suggestScoreDoc.doc);
      assertEquals(doc.getField("int_field").numericValue().intValue(), fieldValue);
    }

    reader.close();
    iw.close();
  }