lucene/highlighter/src/java/org/apache/lucene/search/vectorhighlight/FastVectorHighlighter.java

/*
 * 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.
 */
package org.apache.lucene.search.vectorhighlight;

import java.io.IOException;
import java.util.Iterator;
import java.util.Set;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.highlight.Encoder;

/**
 * Another highlighter implementation.
 *
 */
public class FastVectorHighlighter {
  public static final boolean DEFAULT_PHRASE_HIGHLIGHT = true;
  public static final boolean DEFAULT_FIELD_MATCH = true;
  protected final boolean phraseHighlight;
  protected final boolean fieldMatch;
  private final FragListBuilder fragListBuilder;
  private final FragmentsBuilder fragmentsBuilder;
  private int phraseLimit = Integer.MAX_VALUE;

  /**
   * the default constructor.
   */
  public FastVectorHighlighter(){
    this( DEFAULT_PHRASE_HIGHLIGHT, DEFAULT_FIELD_MATCH );
  }

  /**
   * a constructor. Using {@link SimpleFragListBuilder} and {@link ScoreOrderFragmentsBuilder}.
   * 
   * @param phraseHighlight true or false for phrase highlighting
   * @param fieldMatch true of false for field matching
   */
  public FastVectorHighlighter( boolean phraseHighlight, boolean fieldMatch ){
    this( phraseHighlight, fieldMatch, new SimpleFragListBuilder(), new ScoreOrderFragmentsBuilder() );
  }

  /**
   * a constructor. A {@link FragListBuilder} and a {@link FragmentsBuilder} can be specified (plugins).
   * 
   * @param phraseHighlight true of false for phrase highlighting
   * @param fieldMatch true of false for field matching
   * @param fragListBuilder an instance of {@link FragListBuilder}
   * @param fragmentsBuilder an instance of {@link FragmentsBuilder}
   */
  public FastVectorHighlighter( boolean phraseHighlight, boolean fieldMatch,
      FragListBuilder fragListBuilder, FragmentsBuilder fragmentsBuilder ){
    this.phraseHighlight = phraseHighlight;
    this.fieldMatch = fieldMatch;
    this.fragListBuilder = fragListBuilder;
    this.fragmentsBuilder = fragmentsBuilder;
  }

  /**
   * create a {@link FieldQuery} object.
   * 
   * @param query a query
   * @return the created {@link FieldQuery} object
   */
  public FieldQuery getFieldQuery( Query query ) {
    // TODO: should we deprecate this? 
    // because if there is no reader, then we cannot rewrite MTQ.
    try {
      return getFieldQuery(query, null);
    } catch (IOException e) {
      // should never be thrown when reader is null
      throw new RuntimeException (e);
    }
  }
  
  /**
   * create a {@link FieldQuery} object.
   * 
   * @param query a query
   * @return the created {@link FieldQuery} object
   */
  public FieldQuery getFieldQuery( Query query, IndexReader reader ) throws IOException {
    return new FieldQuery( query, reader, phraseHighlight, fieldMatch );
  }

  /**
   * return the best fragment.
   * 
   * @param fieldQuery {@link FieldQuery} object
   * @param reader {@link IndexReader} of the index
   * @param docId document id to be highlighted
   * @param fieldName field of the document to be highlighted
   * @param fragCharSize the length (number of chars) of a fragment
   * @return the best fragment (snippet) string
   * @throws IOException If there is a low-level I/O error
   */
  public final String getBestFragment( final FieldQuery fieldQuery, IndexReader reader, int docId,
      String fieldName, int fragCharSize ) throws IOException {
    FieldFragList fieldFragList =
      getFieldFragList( fragListBuilder, fieldQuery, reader, docId, fieldName, fragCharSize );
    return fragmentsBuilder.createFragment( reader, docId, fieldName, fieldFragList );
  }

  /**
   * return the best fragments.
   * 
   * @param fieldQuery {@link FieldQuery} object
   * @param reader {@link IndexReader} of the index
   * @param docId document id to be highlighted
   * @param fieldName field of the document to be highlighted
   * @param fragCharSize the length (number of chars) of a fragment
   * @param maxNumFragments maximum number of fragments
   * @return created fragments or null when no fragments created.
   *         size of the array can be less than maxNumFragments
   * @throws IOException If there is a low-level I/O error
   */
  public final String[] getBestFragments( final FieldQuery fieldQuery, IndexReader reader, int docId,
      String fieldName, int fragCharSize, int maxNumFragments ) throws IOException {
    FieldFragList fieldFragList =
      getFieldFragList( fragListBuilder, fieldQuery, reader, docId, fieldName, fragCharSize );
    return fragmentsBuilder.createFragments( reader, docId, fieldName, fieldFragList, maxNumFragments );
  }

  /**
   * return the best fragment.
   * 
   * @param fieldQuery {@link FieldQuery} object
   * @param reader {@link IndexReader} of the index
   * @param docId document id to be highlighted
   * @param fieldName field of the document to be highlighted
   * @param fragCharSize the length (number of chars) of a fragment
   * @param fragListBuilder {@link FragListBuilder} object
   * @param fragmentsBuilder {@link FragmentsBuilder} object
   * @param preTags pre-tags to be used to highlight terms
   * @param postTags post-tags to be used to highlight terms
   * @param encoder an encoder that generates encoded text
   * @return the best fragment (snippet) string
   * @throws IOException If there is a low-level I/O error
   */
  public final String getBestFragment( final FieldQuery fieldQuery, IndexReader reader, int docId,
      String fieldName, int fragCharSize,
      FragListBuilder fragListBuilder, FragmentsBuilder fragmentsBuilder,
      String[] preTags, String[] postTags, Encoder encoder ) throws IOException {
    FieldFragList fieldFragList = getFieldFragList( fragListBuilder, fieldQuery, reader, docId, fieldName, fragCharSize );
    return fragmentsBuilder.createFragment( reader, docId, fieldName, fieldFragList, preTags, postTags, encoder );
  }

  /**
   * return the best fragments.
   * 
   * @param fieldQuery {@link FieldQuery} object
   * @param reader {@link IndexReader} of the index
   * @param docId document id to be highlighted
   * @param fieldName field of the document to be highlighted
   * @param fragCharSize the length (number of chars) of a fragment
   * @param maxNumFragments maximum number of fragments
   * @param fragListBuilder {@link FragListBuilder} object
   * @param fragmentsBuilder {@link FragmentsBuilder} object
   * @param preTags pre-tags to be used to highlight terms
   * @param postTags post-tags to be used to highlight terms
   * @param encoder an encoder that generates encoded text
   * @return created fragments or null when no fragments created.
   *         size of the array can be less than maxNumFragments
   * @throws IOException If there is a low-level I/O error
   */
  public final String[] getBestFragments( final FieldQuery fieldQuery, IndexReader reader, int docId,
      String fieldName, int fragCharSize, int maxNumFragments,
      FragListBuilder fragListBuilder, FragmentsBuilder fragmentsBuilder,
      String[] preTags, String[] postTags, Encoder encoder ) throws IOException {
    FieldFragList fieldFragList =
      getFieldFragList( fragListBuilder, fieldQuery, reader, docId, fieldName, fragCharSize );
    return fragmentsBuilder.createFragments( reader, docId, fieldName, fieldFragList, maxNumFragments,
        preTags, postTags, encoder );
  }

  /**
   * Return the best fragments.  Matches are scanned from matchedFields and turned into fragments against
   * storedField.  The highlighting may not make sense if matchedFields has matches with offsets that don't
   * correspond features in storedField.  It will outright throw a {@code StringIndexOutOfBoundsException}
   * if matchedFields produces offsets outside of storedField.  As such it is advisable that all
   * matchedFields share the same source as storedField or are at least a prefix of it.
   * 
   * @param fieldQuery {@link FieldQuery} object
   * @param reader {@link IndexReader} of the index
   * @param docId document id to be highlighted
   * @param storedField field of the document that stores the text
   * @param matchedFields fields of the document to scan for matches
   * @param fragCharSize the length (number of chars) of a fragment
   * @param maxNumFragments maximum number of fragments
   * @param fragListBuilder {@link FragListBuilder} object
   * @param fragmentsBuilder {@link FragmentsBuilder} object
   * @param preTags pre-tags to be used to highlight terms
   * @param postTags post-tags to be used to highlight terms
   * @param encoder an encoder that generates encoded text
   * @return created fragments or null when no fragments created.
   *         size of the array can be less than maxNumFragments
   * @throws IOException If there is a low-level I/O error
   */
  public final String[] getBestFragments( final FieldQuery fieldQuery, IndexReader reader, int docId,
      String storedField, Set< String > matchedFields, int fragCharSize, int maxNumFragments,
      FragListBuilder fragListBuilder, FragmentsBuilder fragmentsBuilder,
      String[] preTags, String[] postTags, Encoder encoder ) throws IOException {
    FieldFragList fieldFragList =
      getFieldFragList( fragListBuilder, fieldQuery, reader, docId, matchedFields, fragCharSize );
    return fragmentsBuilder.createFragments( reader, docId, storedField, fieldFragList, maxNumFragments,
        preTags, postTags, encoder );
  }

  /**
   * Build a FieldFragList for one field.
   */
  private FieldFragList getFieldFragList( FragListBuilder fragListBuilder,
      final FieldQuery fieldQuery, IndexReader reader, int docId,
      String matchedField, int fragCharSize ) throws IOException {
    FieldTermStack fieldTermStack = new FieldTermStack( reader, docId, matchedField, fieldQuery );
    FieldPhraseList fieldPhraseList = new FieldPhraseList( fieldTermStack, fieldQuery, phraseLimit );
    return fragListBuilder.createFieldFragList( fieldPhraseList, fragCharSize );
  }

  /**
   * Build a FieldFragList for more than one field.
   */
  private FieldFragList getFieldFragList( FragListBuilder fragListBuilder,
      final FieldQuery fieldQuery, IndexReader reader, int docId,
      Set< String > matchedFields, int fragCharSize ) throws IOException {
    Iterator< String > matchedFieldsItr = matchedFields.iterator();
    if ( !matchedFieldsItr.hasNext() ) {
      throw new IllegalArgumentException( "matchedFields must contain at least on field name." );
    }
    FieldPhraseList[] toMerge = new FieldPhraseList[ matchedFields.size() ];
    int i = 0;
    while ( matchedFieldsItr.hasNext() ) {
      FieldTermStack stack = new FieldTermStack( reader, docId, matchedFieldsItr.next(), fieldQuery );
      toMerge[ i++ ] = new FieldPhraseList( stack, fieldQuery, phraseLimit );
    } 
    return fragListBuilder.createFieldFragList( new FieldPhraseList( toMerge ), fragCharSize );
  }

  /**
   * return whether phraseHighlight or not.
   * 
   * @return whether phraseHighlight or not
   */
  public boolean isPhraseHighlight(){ return phraseHighlight; }

  /**
   * return whether fieldMatch or not.
   * 
   * @return whether fieldMatch or not
   */
  public boolean isFieldMatch(){ return fieldMatch; }
  
  /**
   * @return the maximum number of phrases to analyze when searching for the highest-scoring phrase.
   */
  public int getPhraseLimit () { return phraseLimit; }
  
  /**
   * set the maximum number of phrases to analyze when searching for the highest-scoring phrase.
   * The default is unlimited (Integer.MAX_VALUE).
   */
  public void setPhraseLimit (int phraseLimit) { this.phraseLimit = phraseLimit; }
}