| /* |
| * 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.suggest.analyzing; |
| |
| import java.io.Closeable; |
| import java.io.IOException; |
| import java.io.StringReader; |
| import java.nio.file.Path; |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| |
| import org.apache.lucene.analysis.Analyzer; |
| import org.apache.lucene.analysis.AnalyzerWrapper; |
| import org.apache.lucene.analysis.TokenFilter; |
| import org.apache.lucene.analysis.TokenStream; |
| import org.apache.lucene.analysis.ngram.EdgeNGramTokenFilter; |
| import org.apache.lucene.analysis.tokenattributes.CharTermAttribute; |
| import org.apache.lucene.analysis.tokenattributes.OffsetAttribute; |
| import org.apache.lucene.document.BinaryDocValuesField; |
| import org.apache.lucene.document.Document; |
| import org.apache.lucene.document.Field; |
| import org.apache.lucene.document.FieldType; |
| import org.apache.lucene.document.NumericDocValuesField; |
| import org.apache.lucene.document.SortedSetDocValuesField; |
| import org.apache.lucene.document.StringField; |
| import org.apache.lucene.document.TextField; |
| import org.apache.lucene.index.BinaryDocValues; |
| import org.apache.lucene.index.DirectoryReader; |
| import org.apache.lucene.index.FilterLeafReader; |
| import org.apache.lucene.index.IndexOptions; |
| import org.apache.lucene.index.IndexWriter; |
| import org.apache.lucene.index.IndexWriterConfig; |
| import org.apache.lucene.index.LeafReader; |
| import org.apache.lucene.index.LeafReaderContext; |
| import org.apache.lucene.index.MultiDocValues; |
| import org.apache.lucene.index.ReaderUtil; |
| import org.apache.lucene.index.SegmentReader; |
| import org.apache.lucene.index.SortedSetDocValues; |
| import org.apache.lucene.index.Term; |
| import org.apache.lucene.search.BooleanClause.Occur; |
| import org.apache.lucene.search.BooleanClause; |
| import org.apache.lucene.search.BooleanQuery; |
| import org.apache.lucene.search.FieldDoc; |
| import org.apache.lucene.search.IndexSearcher; |
| import org.apache.lucene.search.PrefixQuery; |
| import org.apache.lucene.search.Query; |
| import org.apache.lucene.search.SearcherManager; |
| import org.apache.lucene.search.Sort; |
| import org.apache.lucene.search.SortField; |
| import org.apache.lucene.search.TermQuery; |
| import org.apache.lucene.search.TopFieldCollector; |
| import org.apache.lucene.search.TopFieldDocs; |
| import org.apache.lucene.search.suggest.InputIterator; |
| import org.apache.lucene.search.suggest.Lookup; |
| import org.apache.lucene.store.DataInput; |
| import org.apache.lucene.store.DataOutput; |
| import org.apache.lucene.store.Directory; |
| import org.apache.lucene.store.FSDirectory; |
| import org.apache.lucene.util.Accountable; |
| import org.apache.lucene.util.Accountables; |
| import org.apache.lucene.util.BytesRef; |
| import org.apache.lucene.util.RamUsageEstimator; |
| |
| // TODO: |
| // - a PostingsFormat that stores super-high-freq terms as |
| // a bitset should be a win for the prefix terms? |
| // (LUCENE-5052) |
| // - we could offer a better integration with |
| // DocumentDictionary and NRT? so that your suggester |
| // "automatically" keeps in sync w/ your index |
| |
| /** Analyzes the input text and then suggests matches based |
| * on prefix matches to any tokens in the indexed text. |
| * This also highlights the tokens that match. |
| * |
| * <p>This suggester supports payloads. Matches are sorted only |
| * by the suggest weight; it would be nice to support |
| * blended score + weight sort in the future. This means |
| * this suggester best applies when there is a strong |
| * a-priori ranking of all the suggestions. |
| * |
| * <p>This suggester supports contexts, including arbitrary binary |
| * terms. |
| * |
| * @lucene.experimental */ |
| |
| public class AnalyzingInfixSuggester extends Lookup implements Closeable { |
| |
| /** edgegrams for searching short prefixes without Prefix Query |
| * that's controlled by {@linkplain #minPrefixChars} */ |
| protected final static String TEXTGRAMS_FIELD_NAME = "textgrams"; |
| |
| /** Field name used for the indexed text. */ |
| protected final static String TEXT_FIELD_NAME = "text"; |
| |
| /** Field name used for the indexed text, as a |
| * StringField, for exact lookup. */ |
| protected final static String EXACT_TEXT_FIELD_NAME = "exacttext"; |
| |
| /** Field name used for the indexed context, as a |
| * StringField and a SortedSetDVField, for filtering. */ |
| protected final static String CONTEXTS_FIELD_NAME = "contexts"; |
| |
| /** Analyzer used at search time */ |
| protected final Analyzer queryAnalyzer; |
| /** Analyzer used at index time */ |
| protected final Analyzer indexAnalyzer; |
| private final Directory dir; |
| final int minPrefixChars; |
| |
| private final boolean allTermsRequired; |
| private final boolean highlight; |
| |
| private final boolean commitOnBuild; |
| private final boolean closeIndexWriterOnBuild; |
| |
| /** Used for ongoing NRT additions/updates. */ |
| protected IndexWriter writer; |
| |
| /** {@link IndexSearcher} used for lookups. */ |
| protected SearcherManager searcherMgr; |
| |
| /** Used to manage concurrent access to searcherMgr */ |
| protected final Object searcherMgrLock = new Object(); |
| |
| /** Default minimum number of leading characters before |
| * PrefixQuery is used (4). */ |
| public static final int DEFAULT_MIN_PREFIX_CHARS = 4; |
| |
| /** Default boolean clause option for multiple terms matching (all terms required). */ |
| public static final boolean DEFAULT_ALL_TERMS_REQUIRED = true; |
| |
| /** Default higlighting option. */ |
| public static final boolean DEFAULT_HIGHLIGHT = true; |
| |
| /** Default option to close the IndexWriter once the index has been built. */ |
| protected final static boolean DEFAULT_CLOSE_INDEXWRITER_ON_BUILD = true; |
| |
| /** How we sort the postings and search results. */ |
| private static final Sort SORT = new Sort(new SortField("weight", SortField.Type.LONG, true)); |
| |
| /** Create a new instance, loading from a previously built |
| * AnalyzingInfixSuggester directory, if it exists. This directory must be |
| * private to the infix suggester (i.e., not an external |
| * Lucene index). Note that {@link #close} |
| * will also close the provided directory. */ |
| public AnalyzingInfixSuggester(Directory dir, Analyzer analyzer) throws IOException { |
| this(dir, analyzer, analyzer, DEFAULT_MIN_PREFIX_CHARS, false, DEFAULT_ALL_TERMS_REQUIRED, DEFAULT_HIGHLIGHT); |
| } |
| |
| /** Create a new instance, loading from a previously built |
| * AnalyzingInfixSuggester directory, if it exists. This directory must be |
| * private to the infix suggester (i.e., not an external |
| * Lucene index). Note that {@link #close} |
| * will also close the provided directory. |
| * |
| * @param minPrefixChars Minimum number of leading characters |
| * before PrefixQuery is used (default 4). |
| * Prefixes shorter than this are indexed as character |
| * ngrams (increasing index size but making lookups |
| * faster). |
| * |
| * @param commitOnBuild Call commit after the index has finished building. This would persist the |
| * suggester index to disk and future instances of this suggester can use this pre-built dictionary. |
| */ |
| public AnalyzingInfixSuggester(Directory dir, Analyzer indexAnalyzer, Analyzer queryAnalyzer, int minPrefixChars, |
| boolean commitOnBuild) throws IOException { |
| this(dir, indexAnalyzer, queryAnalyzer, minPrefixChars, commitOnBuild, DEFAULT_ALL_TERMS_REQUIRED, DEFAULT_HIGHLIGHT); |
| } |
| |
| /** Create a new instance, loading from a previously built |
| * AnalyzingInfixSuggester directory, if it exists. This directory must be |
| * private to the infix suggester (i.e., not an external |
| * Lucene index). Note that {@link #close} |
| * will also close the provided directory. |
| * |
| * @param minPrefixChars Minimum number of leading characters |
| * before PrefixQuery is used (default 4). |
| * Prefixes shorter than this are indexed as character |
| * ngrams (increasing index size but making lookups |
| * faster). |
| * |
| * @param commitOnBuild Call commit after the index has finished building. This would persist the |
| * suggester index to disk and future instances of this suggester can use this pre-built dictionary. |
| * |
| * @param allTermsRequired All terms in the suggest query must be matched. |
| * @param highlight Highlight suggest query in suggestions. |
| * |
| */ |
| public AnalyzingInfixSuggester(Directory dir, Analyzer indexAnalyzer, Analyzer queryAnalyzer, int minPrefixChars, |
| boolean commitOnBuild, |
| boolean allTermsRequired, boolean highlight) throws IOException { |
| this(dir, indexAnalyzer, queryAnalyzer, minPrefixChars, commitOnBuild, allTermsRequired, highlight, |
| DEFAULT_CLOSE_INDEXWRITER_ON_BUILD); |
| } |
| |
| /** Create a new instance, loading from a previously built |
| * AnalyzingInfixSuggester directory, if it exists. This directory must be |
| * private to the infix suggester (i.e., not an external |
| * Lucene index). Note that {@link #close} |
| * will also close the provided directory. |
| * |
| * @param minPrefixChars Minimum number of leading characters |
| * before PrefixQuery is used (default 4). |
| * Prefixes shorter than this are indexed as character |
| * ngrams (increasing index size but making lookups |
| * faster). |
| * |
| * @param commitOnBuild Call commit after the index has finished building. This would persist the |
| * suggester index to disk and future instances of this suggester can use this pre-built dictionary. |
| * |
| * @param allTermsRequired All terms in the suggest query must be matched. |
| * @param highlight Highlight suggest query in suggestions. |
| * @param closeIndexWriterOnBuild If true, the IndexWriter will be closed after the index has finished building. |
| */ |
| public AnalyzingInfixSuggester(Directory dir, Analyzer indexAnalyzer, Analyzer queryAnalyzer, int minPrefixChars, |
| boolean commitOnBuild, boolean allTermsRequired, |
| boolean highlight, boolean closeIndexWriterOnBuild) throws IOException { |
| |
| if (minPrefixChars < 0) { |
| throw new IllegalArgumentException("minPrefixChars must be >= 0; got: " + minPrefixChars); |
| } |
| |
| this.queryAnalyzer = queryAnalyzer; |
| this.indexAnalyzer = indexAnalyzer; |
| this.dir = dir; |
| this.minPrefixChars = minPrefixChars; |
| this.commitOnBuild = commitOnBuild; |
| this.allTermsRequired = allTermsRequired; |
| this.highlight = highlight; |
| this.closeIndexWriterOnBuild = closeIndexWriterOnBuild; |
| |
| if (DirectoryReader.indexExists(dir)) { |
| // Already built; open it: |
| searcherMgr = new SearcherManager(dir, null); |
| } |
| } |
| |
| /** Override this to customize index settings, e.g. which |
| * codec to use. */ |
| protected IndexWriterConfig getIndexWriterConfig(Analyzer indexAnalyzer, IndexWriterConfig.OpenMode openMode) { |
| IndexWriterConfig iwc = new IndexWriterConfig(indexAnalyzer); |
| iwc.setOpenMode(openMode); |
| |
| // This way all merged segments will be sorted at |
| // merge time, allow for per-segment early termination |
| // when those segments are searched: |
| iwc.setIndexSort(SORT); |
| |
| return iwc; |
| } |
| |
| /** Subclass can override to choose a specific {@link |
| * Directory} implementation. */ |
| protected Directory getDirectory(Path path) throws IOException { |
| return FSDirectory.open(path); |
| } |
| |
| @Override |
| public void build(InputIterator iter) throws IOException { |
| |
| synchronized (searcherMgrLock) { |
| if (searcherMgr != null) { |
| searcherMgr.close(); |
| searcherMgr = null; |
| } |
| |
| if (writer != null) { |
| writer.close(); |
| writer = null; |
| } |
| |
| boolean success = false; |
| try { |
| // First pass: build a temporary normal Lucene index, |
| // just indexing the suggestions as they iterate: |
| writer = new IndexWriter(dir, |
| getIndexWriterConfig(getGramAnalyzer(), IndexWriterConfig.OpenMode.CREATE)); |
| //long t0 = System.nanoTime(); |
| |
| // TODO: use threads? |
| BytesRef text; |
| while ((text = iter.next()) != null) { |
| BytesRef payload; |
| if (iter.hasPayloads()) { |
| payload = iter.payload(); |
| } else { |
| payload = null; |
| } |
| |
| add(text, iter.contexts(), iter.weight(), payload); |
| } |
| |
| //System.out.println("initial indexing time: " + ((System.nanoTime()-t0)/1000000) + " msec"); |
| if (commitOnBuild || closeIndexWriterOnBuild) { |
| commit(); |
| } |
| searcherMgr = new SearcherManager(writer, null); |
| success = true; |
| } finally { |
| if (success) { |
| if (closeIndexWriterOnBuild) { |
| writer.close(); |
| writer = null; |
| } |
| } else { // failure |
| if (writer != null) { |
| writer.rollback(); |
| writer = null; |
| } |
| } |
| } |
| } |
| } |
| |
| /** Commits all pending changes made to this suggester to disk. |
| * |
| * @see IndexWriter#commit */ |
| public void commit() throws IOException { |
| if (writer == null) { |
| if (searcherMgr == null || closeIndexWriterOnBuild == false) { |
| throw new IllegalStateException("Cannot commit on an closed writer. Add documents first"); |
| } |
| // else no-op: writer was committed and closed after the index was built, so commit is unnecessary |
| } else { |
| writer.commit(); |
| } |
| } |
| |
| private Analyzer getGramAnalyzer() { |
| return new AnalyzerWrapper(Analyzer.PER_FIELD_REUSE_STRATEGY) { |
| @Override |
| protected Analyzer getWrappedAnalyzer(String fieldName) { |
| return indexAnalyzer; |
| } |
| |
| @Override |
| protected TokenStreamComponents wrapComponents(String fieldName, TokenStreamComponents components) { |
| assert !(fieldName.equals(TEXTGRAMS_FIELD_NAME) && minPrefixChars == 0) |
| : "no need \"textgrams\" when minPrefixChars="+minPrefixChars; |
| if (fieldName.equals(TEXTGRAMS_FIELD_NAME) && minPrefixChars > 0) { |
| // TODO: should use an EdgeNGramTokenFilterFactory here |
| TokenFilter filter = new EdgeNGramTokenFilter(components.getTokenStream(), 1, minPrefixChars, false); |
| return new TokenStreamComponents(components.getSource(), filter); |
| } else { |
| return components; |
| } |
| } |
| }; |
| } |
| |
| private void ensureOpen() throws IOException { |
| synchronized (searcherMgrLock) { |
| if (writer == null) { |
| if (DirectoryReader.indexExists(dir)) { |
| // Already built; open it: |
| writer = new IndexWriter(dir, getIndexWriterConfig(getGramAnalyzer(), IndexWriterConfig.OpenMode.APPEND)); |
| } else { |
| writer = new IndexWriter(dir, getIndexWriterConfig(getGramAnalyzer(), IndexWriterConfig.OpenMode.CREATE)); |
| } |
| |
| SearcherManager oldSearcherMgr = searcherMgr; |
| searcherMgr = new SearcherManager(writer, null); |
| if (oldSearcherMgr != null) { |
| oldSearcherMgr.close(); |
| } |
| } |
| } |
| } |
| |
| /** Adds a new suggestion. Be sure to use {@link #update} |
| * instead if you want to replace a previous suggestion. |
| * After adding or updating a batch of new suggestions, |
| * you must call {@link #refresh} in the end in order to |
| * see the suggestions in {@link #lookup} */ |
| public void add(BytesRef text, Set<BytesRef> contexts, long weight, BytesRef payload) throws IOException { |
| ensureOpen(); |
| writer.addDocument(buildDocument(text, contexts, weight, payload)); |
| } |
| |
| /** Updates a previous suggestion, matching the exact same |
| * text as before. Use this to change the weight or |
| * payload of an already added suggestion. If you know |
| * this text is not already present you can use {@link |
| * #add} instead. After adding or updating a batch of |
| * new suggestions, you must call {@link #refresh} in the |
| * end in order to see the suggestions in {@link #lookup} */ |
| public void update(BytesRef text, Set<BytesRef> contexts, long weight, BytesRef payload) throws IOException { |
| ensureOpen(); |
| writer.updateDocument(new Term(EXACT_TEXT_FIELD_NAME, text.utf8ToString()), |
| buildDocument(text, contexts, weight, payload)); |
| } |
| |
| private Document buildDocument(BytesRef text, Set<BytesRef> contexts, long weight, BytesRef payload) throws IOException { |
| String textString = text.utf8ToString(); |
| Document doc = new Document(); |
| FieldType ft = getTextFieldType(); |
| doc.add(new Field(TEXT_FIELD_NAME, textString, ft)); |
| if (minPrefixChars>0) { |
| doc.add(new Field(TEXTGRAMS_FIELD_NAME, textString, ft)); |
| } |
| doc.add(new StringField(EXACT_TEXT_FIELD_NAME, textString, Field.Store.NO)); |
| doc.add(new BinaryDocValuesField(TEXT_FIELD_NAME, text)); |
| doc.add(new NumericDocValuesField("weight", weight)); |
| if (payload != null) { |
| doc.add(new BinaryDocValuesField("payloads", payload)); |
| } |
| if (contexts != null) { |
| for(BytesRef context : contexts) { |
| doc.add(new StringField(CONTEXTS_FIELD_NAME, context, Field.Store.NO)); |
| doc.add(new SortedSetDocValuesField(CONTEXTS_FIELD_NAME, context)); |
| } |
| } |
| return doc; |
| } |
| |
| /** Reopens the underlying searcher; it's best to "batch |
| * up" many additions/updates, and then call refresh |
| * once in the end. */ |
| public void refresh() throws IOException { |
| if (searcherMgr == null) { |
| throw new IllegalStateException("suggester was not built"); |
| } |
| if (writer != null) { |
| searcherMgr.maybeRefreshBlocking(); |
| } |
| // else no-op: writer was committed and closed after the index was built |
| // and before searchMgr was constructed, so refresh is unnecessary |
| } |
| |
| /** |
| * Subclass can override this method to change the field type of the text field |
| * e.g. to change the index options |
| */ |
| protected FieldType getTextFieldType(){ |
| FieldType ft = new FieldType(TextField.TYPE_NOT_STORED); |
| ft.setIndexOptions(IndexOptions.DOCS); |
| ft.setOmitNorms(true); |
| |
| return ft; |
| } |
| |
| @Override |
| public List<LookupResult> lookup(CharSequence key, Set<BytesRef> contexts, boolean onlyMorePopular, int num) throws IOException { |
| return lookup(key, contexts, num, allTermsRequired, highlight); |
| } |
| |
| /** Lookup, without any context. */ |
| public List<LookupResult> lookup(CharSequence key, int num, boolean allTermsRequired, boolean doHighlight) throws IOException { |
| return lookup(key, (BooleanQuery)null, num, allTermsRequired, doHighlight); |
| } |
| |
| /** Lookup, with context but without booleans. Context booleans default to SHOULD, |
| * so each suggestion must have at least one of the contexts. */ |
| public List<LookupResult> lookup(CharSequence key, Set<BytesRef> contexts, int num, boolean allTermsRequired, boolean doHighlight) throws IOException { |
| return lookup(key, toQuery(contexts), num, allTermsRequired, doHighlight); |
| } |
| |
| /** This is called if the last token isn't ended |
| * (e.g. user did not type a space after it). Return an |
| * appropriate Query clause to add to the BooleanQuery. */ |
| protected Query getLastTokenQuery(String token) throws IOException { |
| if (token.length() < minPrefixChars) { |
| // The leading ngram was directly indexed: |
| return new TermQuery(new Term(TEXTGRAMS_FIELD_NAME, token)); |
| } |
| |
| return new PrefixQuery(new Term(TEXT_FIELD_NAME, token)); |
| } |
| |
| /** Retrieve suggestions, specifying whether all terms |
| * must match ({@code allTermsRequired}) and whether the hits |
| * should be highlighted ({@code doHighlight}). */ |
| public List<LookupResult> lookup(CharSequence key, Map<BytesRef, BooleanClause.Occur> contextInfo, int num, boolean allTermsRequired, boolean doHighlight) throws IOException { |
| return lookup(key, toQuery(contextInfo), num, allTermsRequired, doHighlight); |
| } |
| |
| private BooleanQuery toQuery(Map<BytesRef,BooleanClause.Occur> contextInfo) { |
| if (contextInfo == null || contextInfo.isEmpty()) { |
| return null; |
| } |
| |
| BooleanQuery.Builder contextFilter = new BooleanQuery.Builder(); |
| for (Map.Entry<BytesRef,BooleanClause.Occur> entry : contextInfo.entrySet()) { |
| addContextToQuery(contextFilter, entry.getKey(), entry.getValue()); |
| } |
| |
| return contextFilter.build(); |
| } |
| |
| private BooleanQuery toQuery(Set<BytesRef> contextInfo) { |
| if (contextInfo == null || contextInfo.isEmpty()) { |
| return null; |
| } |
| |
| BooleanQuery.Builder contextFilter = new BooleanQuery.Builder(); |
| for (BytesRef context : contextInfo) { |
| addContextToQuery(contextFilter, context, BooleanClause.Occur.SHOULD); |
| } |
| return contextFilter.build(); |
| } |
| |
| |
| /** |
| * This method is handy as we do not need access to internal fields such as CONTEXTS_FIELD_NAME in order to build queries |
| * However, here may not be its best location. |
| * |
| * @param query an instance of @See {@link BooleanQuery} |
| * @param context the context |
| * @param clause one of {@link Occur} |
| */ |
| public void addContextToQuery(BooleanQuery.Builder query, BytesRef context, BooleanClause.Occur clause) { |
| // NOTE: we "should" wrap this in |
| // ConstantScoreQuery, or maybe send this as a |
| // Filter instead to search. |
| |
| // TODO: if we had a BinaryTermField we could fix |
| // this "must be valid ut8f" limitation: |
| query.add(new TermQuery(new Term(CONTEXTS_FIELD_NAME, context)), clause); |
| } |
| |
| /** |
| * This is an advanced method providing the capability to send down to the suggester any |
| * arbitrary lucene query to be used to filter the result of the suggester |
| * |
| * @param key the keyword being looked for |
| * @param contextQuery an arbitrary Lucene query to be used to filter the result of the suggester. {@link #addContextToQuery} could be used to build this contextQuery. |
| * @param num number of items to return |
| * @param allTermsRequired all searched terms must match or not |
| * @param doHighlight if true, the matching term will be highlighted in the search result |
| * @return the result of the suggester |
| * @throws IOException f the is IO exception while reading data from the index |
| */ |
| public List<LookupResult> lookup(CharSequence key, BooleanQuery contextQuery, int num, boolean allTermsRequired, boolean doHighlight) throws IOException { |
| |
| if (searcherMgr == null) { |
| throw new IllegalStateException("suggester was not built"); |
| } |
| |
| final BooleanClause.Occur occur; |
| if (allTermsRequired) { |
| occur = BooleanClause.Occur.MUST; |
| } else { |
| occur = BooleanClause.Occur.SHOULD; |
| } |
| |
| BooleanQuery.Builder query; |
| Set<String> matchedTokens; |
| String prefixToken = null; |
| |
| try (TokenStream ts = queryAnalyzer.tokenStream("", new StringReader(key.toString()))) { |
| //long t0 = System.currentTimeMillis(); |
| ts.reset(); |
| final CharTermAttribute termAtt = ts.addAttribute(CharTermAttribute.class); |
| final OffsetAttribute offsetAtt = ts.addAttribute(OffsetAttribute.class); |
| String lastToken = null; |
| query = new BooleanQuery.Builder(); |
| int maxEndOffset = -1; |
| matchedTokens = new HashSet<>(); |
| while (ts.incrementToken()) { |
| if (lastToken != null) { |
| matchedTokens.add(lastToken); |
| query.add(new TermQuery(new Term(TEXT_FIELD_NAME, lastToken)), occur); |
| } |
| lastToken = termAtt.toString(); |
| if (lastToken != null) { |
| maxEndOffset = Math.max(maxEndOffset, offsetAtt.endOffset()); |
| } |
| } |
| ts.end(); |
| |
| if (lastToken != null) { |
| Query lastQuery; |
| if (maxEndOffset == offsetAtt.endOffset()) { |
| // Use PrefixQuery (or the ngram equivalent) when |
| // there was no trailing discarded chars in the |
| // string (e.g. whitespace), so that if query does |
| // not end with a space we show prefix matches for |
| // that token: |
| lastQuery = getLastTokenQuery(lastToken); |
| prefixToken = lastToken; |
| } else { |
| // Use TermQuery for an exact match if there were |
| // trailing discarded chars (e.g. whitespace), so |
| // that if query ends with a space we only show |
| // exact matches for that term: |
| matchedTokens.add(lastToken); |
| lastQuery = new TermQuery(new Term(TEXT_FIELD_NAME, lastToken)); |
| } |
| |
| if (lastQuery != null) { |
| query.add(lastQuery, occur); |
| } |
| } |
| |
| if (contextQuery != null) { |
| boolean allMustNot = true; |
| for (BooleanClause clause : contextQuery.clauses()) { |
| if (clause.getOccur() != BooleanClause.Occur.MUST_NOT) { |
| allMustNot = false; |
| break; |
| } |
| } |
| |
| if (allMustNot) { |
| // All are MUST_NOT: add the contextQuery to the main query instead (not as sub-query) |
| for (BooleanClause clause : contextQuery.clauses()) { |
| query.add(clause); |
| } |
| } else if (allTermsRequired == false) { |
| // We must carefully upgrade the query clauses to MUST: |
| BooleanQuery.Builder newQuery = new BooleanQuery.Builder(); |
| newQuery.add(query.build(), BooleanClause.Occur.MUST); |
| newQuery.add(contextQuery, BooleanClause.Occur.MUST); |
| query = newQuery; |
| } else { |
| // Add contextQuery as sub-query |
| query.add(contextQuery, BooleanClause.Occur.MUST); |
| } |
| } |
| } |
| |
| // TODO: we could allow blended sort here, combining |
| // weight w/ score. Now we ignore score and sort only |
| // by weight: |
| |
| Query finalQuery = finishQuery(query, allTermsRequired); |
| |
| //System.out.println("finalQuery=" + finalQuery); |
| |
| // Sort by weight, descending: |
| TopFieldCollector c = TopFieldCollector.create(SORT, num, 1); |
| List<LookupResult> results = null; |
| SearcherManager mgr; |
| IndexSearcher searcher; |
| synchronized (searcherMgrLock) { |
| mgr = searcherMgr; // acquire & release on same SearcherManager, via local reference |
| searcher = mgr.acquire(); |
| } |
| try { |
| //System.out.println("got searcher=" + searcher); |
| searcher.search(finalQuery, c); |
| |
| TopFieldDocs hits = c.topDocs(); |
| |
| // Slower way if postings are not pre-sorted by weight: |
| // hits = searcher.search(query, null, num, SORT); |
| results = createResults(searcher, hits, num, key, doHighlight, matchedTokens, prefixToken); |
| } finally { |
| mgr.release(searcher); |
| } |
| |
| //System.out.println((System.currentTimeMillis() - t0) + " msec for infix suggest"); |
| //System.out.println(results); |
| |
| return results; |
| } |
| |
| /** |
| * Create the results based on the search hits. |
| * Can be overridden by subclass to add particular behavior (e.g. weight transformation). |
| * Note that there is no prefix token (the {@code prefixToken} argument will |
| * be null) whenever the final token in the incoming request was in fact finished |
| * (had trailing characters, such as white-space). |
| * |
| * @throws IOException If there are problems reading fields from the underlying Lucene index. |
| */ |
| protected List<LookupResult> createResults(IndexSearcher searcher, TopFieldDocs hits, int num, |
| CharSequence charSequence, |
| boolean doHighlight, Set<String> matchedTokens, String prefixToken) |
| throws IOException { |
| |
| List<LeafReaderContext> leaves = searcher.getIndexReader().leaves(); |
| List<LookupResult> results = new ArrayList<>(); |
| for (int i=0;i<hits.scoreDocs.length;i++) { |
| FieldDoc fd = (FieldDoc) hits.scoreDocs[i]; |
| BinaryDocValues textDV = MultiDocValues.getBinaryValues(searcher.getIndexReader(), TEXT_FIELD_NAME); |
| textDV.advance(fd.doc); |
| BytesRef term = textDV.binaryValue(); |
| String text = term.utf8ToString(); |
| long score = (Long) fd.fields[0]; |
| |
| // This will just be null if app didn't pass payloads to build(): |
| // TODO: maybe just stored fields? they compress... |
| BinaryDocValues payloadsDV = MultiDocValues.getBinaryValues(searcher.getIndexReader(), "payloads"); |
| |
| BytesRef payload; |
| if (payloadsDV != null) { |
| if (payloadsDV.advance(fd.doc) == fd.doc) { |
| payload = BytesRef.deepCopyOf(payloadsDV.binaryValue()); |
| } else { |
| payload = new BytesRef(BytesRef.EMPTY_BYTES); |
| } |
| } else { |
| payload = null; |
| } |
| |
| // Must look up sorted-set by segment: |
| int segment = ReaderUtil.subIndex(fd.doc, leaves); |
| SortedSetDocValues contextsDV = leaves.get(segment).reader().getSortedSetDocValues(CONTEXTS_FIELD_NAME); |
| Set<BytesRef> contexts; |
| if (contextsDV != null) { |
| contexts = new HashSet<BytesRef>(); |
| int targetDocID = fd.doc - leaves.get(segment).docBase; |
| if (contextsDV.advance(targetDocID) == targetDocID) { |
| long ord; |
| while ((ord = contextsDV.nextOrd()) != SortedSetDocValues.NO_MORE_ORDS) { |
| BytesRef context = BytesRef.deepCopyOf(contextsDV.lookupOrd(ord)); |
| contexts.add(context); |
| } |
| } |
| } else { |
| contexts = null; |
| } |
| |
| LookupResult result; |
| |
| if (doHighlight) { |
| result = new LookupResult(text, highlight(text, matchedTokens, prefixToken), score, payload, contexts); |
| } else { |
| result = new LookupResult(text, score, payload, contexts); |
| } |
| |
| results.add(result); |
| } |
| |
| return results; |
| } |
| |
| /** Subclass can override this to tweak the Query before |
| * searching. */ |
| protected Query finishQuery(BooleanQuery.Builder in, boolean allTermsRequired) { |
| return in.build(); |
| } |
| |
| /** Override this method to customize the Object |
| * representing a single highlighted suggestions; the |
| * result is set on each {@link |
| * org.apache.lucene.search.suggest.Lookup.LookupResult#highlightKey} member. */ |
| protected Object highlight(String text, Set<String> matchedTokens, String prefixToken) throws IOException { |
| try (TokenStream ts = queryAnalyzer.tokenStream("text", new StringReader(text))) { |
| CharTermAttribute termAtt = ts.addAttribute(CharTermAttribute.class); |
| OffsetAttribute offsetAtt = ts.addAttribute(OffsetAttribute.class); |
| ts.reset(); |
| StringBuilder sb = new StringBuilder(); |
| int upto = 0; |
| while (ts.incrementToken()) { |
| String token = termAtt.toString(); |
| int startOffset = offsetAtt.startOffset(); |
| int endOffset = offsetAtt.endOffset(); |
| if (upto < startOffset) { |
| addNonMatch(sb, text.substring(upto, startOffset)); |
| upto = startOffset; |
| } else if (upto > startOffset) { |
| continue; |
| } |
| |
| if (matchedTokens.contains(token)) { |
| // Token matches. |
| addWholeMatch(sb, text.substring(startOffset, endOffset), token); |
| upto = endOffset; |
| } else if (prefixToken != null && token.startsWith(prefixToken)) { |
| addPrefixMatch(sb, text.substring(startOffset, endOffset), token, prefixToken); |
| upto = endOffset; |
| } |
| } |
| ts.end(); |
| int endOffset = offsetAtt.endOffset(); |
| if (upto < endOffset) { |
| addNonMatch(sb, text.substring(upto)); |
| } |
| return sb.toString(); |
| } |
| } |
| |
| /** Called while highlighting a single result, to append a |
| * non-matching chunk of text from the suggestion to the |
| * provided fragments list. |
| * @param sb The {@code StringBuilder} to append to |
| * @param text The text chunk to add |
| */ |
| protected void addNonMatch(StringBuilder sb, String text) { |
| sb.append(text); |
| } |
| |
| /** Called while highlighting a single result, to append |
| * the whole matched token to the provided fragments list. |
| * @param sb The {@code StringBuilder} to append to |
| * @param surface The surface form (original) text |
| * @param analyzed The analyzed token corresponding to the surface form text |
| */ |
| protected void addWholeMatch(StringBuilder sb, String surface, String analyzed) { |
| sb.append("<b>"); |
| sb.append(surface); |
| sb.append("</b>"); |
| } |
| |
| /** Called while highlighting a single result, to append a |
| * matched prefix token, to the provided fragments list. |
| * @param sb The {@code StringBuilder} to append to |
| * @param surface The fragment of the surface form |
| * (indexed during {@link #build}, corresponding to |
| * this match |
| * @param analyzed The analyzed token that matched |
| * @param prefixToken The prefix of the token that matched |
| */ |
| protected void addPrefixMatch(StringBuilder sb, String surface, String analyzed, String prefixToken) { |
| // TODO: apps can try to invert their analysis logic |
| // here, e.g. downcase the two before checking prefix: |
| if (prefixToken.length() >= surface.length()) { |
| addWholeMatch(sb, surface, analyzed); |
| return; |
| } |
| sb.append("<b>"); |
| sb.append(surface.substring(0, prefixToken.length())); |
| sb.append("</b>"); |
| sb.append(surface.substring(prefixToken.length())); |
| } |
| |
| @Override |
| public boolean store(DataOutput in) throws IOException { |
| return false; |
| } |
| |
| @Override |
| public boolean load(DataInput out) throws IOException { |
| return false; |
| } |
| |
| @Override |
| public void close() throws IOException { |
| if (searcherMgr != null) { |
| searcherMgr.close(); |
| searcherMgr = null; |
| } |
| if (writer != null) { |
| writer.close(); |
| writer = null; |
| } |
| if (dir != null) { |
| dir.close(); |
| } |
| } |
| |
| @Override |
| public long ramBytesUsed() { |
| long mem = RamUsageEstimator.shallowSizeOf(this); |
| try { |
| if (searcherMgr != null) { |
| SearcherManager mgr; |
| IndexSearcher searcher; |
| synchronized (searcherMgrLock) { |
| mgr = searcherMgr; // acquire & release on same SearcherManager, via local reference |
| searcher = mgr.acquire(); |
| } |
| try { |
| for (LeafReaderContext context : searcher.getIndexReader().leaves()) { |
| LeafReader reader = FilterLeafReader.unwrap(context.reader()); |
| if (reader instanceof SegmentReader) { |
| mem += ((SegmentReader) context.reader()).ramBytesUsed(); |
| } |
| } |
| } finally { |
| mgr.release(searcher); |
| } |
| } |
| return mem; |
| } catch (IOException ioe) { |
| throw new RuntimeException(ioe); |
| } |
| } |
| |
| @Override |
| public Collection<Accountable> getChildResources() { |
| List<Accountable> resources = new ArrayList<>(); |
| try { |
| if (searcherMgr != null) { |
| SearcherManager mgr; |
| IndexSearcher searcher; |
| synchronized (searcherMgrLock) { |
| mgr = searcherMgr; // acquire & release on same SearcherManager, via local reference |
| searcher = mgr.acquire(); |
| } |
| try { |
| for (LeafReaderContext context : searcher.getIndexReader().leaves()) { |
| LeafReader reader = FilterLeafReader.unwrap(context.reader()); |
| if (reader instanceof SegmentReader) { |
| resources.add(Accountables.namedAccountable("segment", (SegmentReader)reader)); |
| } |
| } |
| } finally { |
| mgr.release(searcher); |
| } |
| } |
| return Collections.unmodifiableList(resources); |
| } catch (IOException ioe) { |
| throw new RuntimeException(ioe); |
| } |
| } |
| |
| @Override |
| public long getCount() throws IOException { |
| if (searcherMgr == null) { |
| return 0; |
| } |
| SearcherManager mgr; |
| IndexSearcher searcher; |
| synchronized (searcherMgrLock) { |
| mgr = searcherMgr; // acquire & release on same SearcherManager, via local reference |
| searcher = mgr.acquire(); |
| } |
| try { |
| return searcher.getIndexReader().numDocs(); |
| } finally { |
| mgr.release(searcher); |
| } |
| } |
| } |
