| /* |
| * 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.IOException; |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.Comparator; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Set; |
| |
| import org.apache.lucene.analysis.Analyzer; |
| import org.apache.lucene.analysis.TokenStream; |
| import org.apache.lucene.analysis.TokenStreamToAutomaton; |
| import org.apache.lucene.codecs.CodecUtil; |
| import org.apache.lucene.search.suggest.InputIterator; |
| import org.apache.lucene.search.suggest.Lookup; |
| import org.apache.lucene.store.ByteArrayDataInput; |
| import org.apache.lucene.store.ByteArrayDataOutput; |
| import org.apache.lucene.store.DataInput; |
| import org.apache.lucene.store.DataOutput; |
| import org.apache.lucene.store.Directory; |
| import org.apache.lucene.store.IOContext; |
| import org.apache.lucene.store.IndexOutput; |
| import org.apache.lucene.util.Accountable; |
| import org.apache.lucene.util.Accountables; |
| import org.apache.lucene.util.ArrayUtil; |
| import org.apache.lucene.util.BytesRef; |
| import org.apache.lucene.util.BytesRefBuilder; |
| import org.apache.lucene.util.CharsRefBuilder; |
| import org.apache.lucene.util.IOUtils; |
| import org.apache.lucene.util.IntsRef; |
| import org.apache.lucene.util.IntsRefBuilder; |
| import org.apache.lucene.util.OfflineSorter; |
| import org.apache.lucene.util.automaton.Automaton; |
| import org.apache.lucene.util.automaton.LimitedFiniteStringsIterator; |
| import org.apache.lucene.util.automaton.Operations; |
| import org.apache.lucene.util.automaton.Transition; |
| import org.apache.lucene.util.fst.Builder; |
| import org.apache.lucene.util.fst.ByteSequenceOutputs; |
| import org.apache.lucene.util.fst.FST.BytesReader; |
| import org.apache.lucene.util.fst.FST; |
| import org.apache.lucene.util.fst.PairOutputs.Pair; |
| import org.apache.lucene.util.fst.PairOutputs; |
| import org.apache.lucene.util.fst.PositiveIntOutputs; |
| import org.apache.lucene.util.fst.Util.Result; |
| import org.apache.lucene.util.fst.Util.TopResults; |
| import org.apache.lucene.util.fst.Util; |
| |
| import static org.apache.lucene.util.automaton.Operations.DEFAULT_DETERMINIZE_WORK_LIMIT; |
| |
| /** |
| * Suggester that first analyzes the surface form, adds the |
| * analyzed form to a weighted FST, and then does the same |
| * thing at lookup time. This means lookup is based on the |
| * analyzed form while suggestions are still the surface |
| * form(s). |
| * |
| * <p> |
| * This can result in powerful suggester functionality. For |
| * example, if you use an analyzer removing stop words, |
| * then the partial text "ghost chr..." could see the |
| * suggestion "The Ghost of Christmas Past". Note that |
| * position increments MUST NOT be preserved for this example |
| * to work, so you should call the constructor with |
| * <code>preservePositionIncrements</code> parameter set to |
| * false |
| * |
| * <p> |
| * If SynonymFilter is used to map wifi and wireless network to |
| * hotspot then the partial text "wirele..." could suggest |
| * "wifi router". Token normalization like stemmers, accent |
| * removal, etc., would allow suggestions to ignore such |
| * variations. |
| * |
| * <p> |
| * When two matching suggestions have the same weight, they |
| * are tie-broken by the analyzed form. If their analyzed |
| * form is the same then the order is undefined. |
| * |
| * <p> |
| * There are some limitations: |
| * <ul> |
| * |
| * <li> A lookup from a query like "net" in English won't |
| * be any different than "net " (ie, user added a |
| * trailing space) because analyzers don't reflect |
| * when they've seen a token separator and when they |
| * haven't. |
| * |
| * <li> If you're using {@code StopFilter}, and the user will |
| * type "fast apple", but so far all they've typed is |
| * "fast a", again because the analyzer doesn't convey whether |
| * it's seen a token separator after the "a", |
| * {@code StopFilter} will remove that "a" causing |
| * far more matches than you'd expect. |
| * |
| * <li> Lookups with the empty string return no results |
| * instead of all results. |
| * </ul> |
| * |
| * @lucene.experimental |
| */ |
| // redundant 'implements Accountable' to workaround javadocs bugs |
| public class AnalyzingSuggester extends Lookup implements Accountable { |
| |
| /** |
| * FST<Weight,Surface>: |
| * input is the analyzed form, with a null byte between terms |
| * weights are encoded as costs: (Integer.MAX_VALUE-weight) |
| * surface is the original, unanalyzed form. |
| */ |
| private FST<Pair<Long,BytesRef>> fst = null; |
| |
| /** |
| * Analyzer that will be used for analyzing suggestions at |
| * index time. |
| */ |
| private final Analyzer indexAnalyzer; |
| |
| /** |
| * Analyzer that will be used for analyzing suggestions at |
| * query time. |
| */ |
| private final Analyzer queryAnalyzer; |
| |
| /** |
| * True if exact match suggestions should always be returned first. |
| */ |
| private final boolean exactFirst; |
| |
| /** |
| * True if separator between tokens should be preserved. |
| */ |
| private final boolean preserveSep; |
| |
| /** Include this flag in the options parameter to {@link |
| * #AnalyzingSuggester(Directory,String,Analyzer,Analyzer,int,int,int,boolean)} to always |
| * return the exact match first, regardless of score. This |
| * has no performance impact but could result in |
| * low-quality suggestions. */ |
| public static final int EXACT_FIRST = 1; |
| |
| /** Include this flag in the options parameter to {@link |
| * #AnalyzingSuggester(Directory,String,Analyzer,Analyzer,int,int,int,boolean)} to preserve |
| * token separators when matching. */ |
| public static final int PRESERVE_SEP = 2; |
| |
| /** Represents the separation between tokens, if |
| * PRESERVE_SEP was specified */ |
| private static final int SEP_LABEL = '\u001F'; |
| |
| /** Marks end of the analyzed input and start of dedup |
| * byte. */ |
| private static final int END_BYTE = 0x0; |
| |
| /** Maximum number of dup surface forms (different surface |
| * forms for the same analyzed form). */ |
| private final int maxSurfaceFormsPerAnalyzedForm; |
| |
| /** Maximum graph paths to index for a single analyzed |
| * surface form. This only matters if your analyzer |
| * makes lots of alternate paths (e.g. contains |
| * SynonymFilter). */ |
| private final int maxGraphExpansions; |
| |
| private final Directory tempDir; |
| private final String tempFileNamePrefix; |
| |
| /** Highest number of analyzed paths we saw for any single |
| * input surface form. For analyzers that never create |
| * graphs this will always be 1. */ |
| private int maxAnalyzedPathsForOneInput; |
| |
| private boolean hasPayloads; |
| |
| private static final int PAYLOAD_SEP = '\u001f'; |
| |
| /** Whether position holes should appear in the automaton. */ |
| private boolean preservePositionIncrements; |
| |
| /** Number of entries the lookup was built with */ |
| private long count = 0; |
| |
| /** |
| * Calls {@link #AnalyzingSuggester(Directory,String,Analyzer,Analyzer,int,int,int,boolean) |
| * AnalyzingSuggester(analyzer, analyzer, EXACT_FIRST | |
| * PRESERVE_SEP, 256, -1, true)} |
| */ |
| public AnalyzingSuggester(Directory tempDir, String tempFileNamePrefix, Analyzer analyzer) { |
| this(tempDir, tempFileNamePrefix, analyzer, analyzer, EXACT_FIRST | PRESERVE_SEP, 256, -1, true); |
| } |
| |
| /** |
| * Calls {@link #AnalyzingSuggester(Directory,String,Analyzer,Analyzer,int,int,int,boolean) |
| * AnalyzingSuggester(indexAnalyzer, queryAnalyzer, EXACT_FIRST | |
| * PRESERVE_SEP, 256, -1, true)} |
| */ |
| public AnalyzingSuggester(Directory tempDir, String tempFileNamePrefix, Analyzer indexAnalyzer, Analyzer queryAnalyzer) { |
| this(tempDir, tempFileNamePrefix, indexAnalyzer, queryAnalyzer, EXACT_FIRST | PRESERVE_SEP, 256, -1, true); |
| } |
| |
| /** |
| * Creates a new suggester. |
| * |
| * @param indexAnalyzer Analyzer that will be used for |
| * analyzing suggestions while building the index. |
| * @param queryAnalyzer Analyzer that will be used for |
| * analyzing query text during lookup |
| * @param options see {@link #EXACT_FIRST}, {@link #PRESERVE_SEP} |
| * @param maxSurfaceFormsPerAnalyzedForm Maximum number of |
| * surface forms to keep for a single analyzed form. |
| * When there are too many surface forms we discard the |
| * lowest weighted ones. |
| * @param maxGraphExpansions Maximum number of graph paths |
| * to expand from the analyzed form. Set this to -1 for |
| * no limit. |
| * @param preservePositionIncrements Whether position holes |
| * should appear in the automata |
| */ |
| public AnalyzingSuggester(Directory tempDir, String tempFileNamePrefix, Analyzer indexAnalyzer, Analyzer queryAnalyzer, int options, int maxSurfaceFormsPerAnalyzedForm, int maxGraphExpansions, |
| boolean preservePositionIncrements) { |
| this.indexAnalyzer = indexAnalyzer; |
| this.queryAnalyzer = queryAnalyzer; |
| if ((options & ~(EXACT_FIRST | PRESERVE_SEP)) != 0) { |
| throw new IllegalArgumentException("options should only contain EXACT_FIRST and PRESERVE_SEP; got " + options); |
| } |
| this.exactFirst = (options & EXACT_FIRST) != 0; |
| this.preserveSep = (options & PRESERVE_SEP) != 0; |
| |
| // NOTE: this is just an implementation limitation; if |
| // somehow this is a problem we could fix it by using |
| // more than one byte to disambiguate ... but 256 seems |
| // like it should be way more then enough. |
| if (maxSurfaceFormsPerAnalyzedForm <= 0 || maxSurfaceFormsPerAnalyzedForm > 256) { |
| throw new IllegalArgumentException("maxSurfaceFormsPerAnalyzedForm must be > 0 and < 256 (got: " + maxSurfaceFormsPerAnalyzedForm + ")"); |
| } |
| this.maxSurfaceFormsPerAnalyzedForm = maxSurfaceFormsPerAnalyzedForm; |
| |
| if (maxGraphExpansions < 1 && maxGraphExpansions != -1) { |
| throw new IllegalArgumentException("maxGraphExpansions must -1 (no limit) or > 0 (got: " + maxGraphExpansions + ")"); |
| } |
| this.maxGraphExpansions = maxGraphExpansions; |
| this.preservePositionIncrements = preservePositionIncrements; |
| this.tempDir = tempDir; |
| this.tempFileNamePrefix = tempFileNamePrefix; |
| } |
| |
| /** Returns byte size of the underlying FST. */ |
| @Override |
| public long ramBytesUsed() { |
| return fst == null ? 0 : fst.ramBytesUsed(); |
| } |
| |
| @Override |
| public Collection<Accountable> getChildResources() { |
| if (fst == null) { |
| return Collections.emptyList(); |
| } else { |
| return Collections.singletonList(Accountables.namedAccountable("fst", fst)); |
| } |
| } |
| |
| // Replaces SEP with epsilon or remaps them if |
| // we were asked to preserve them: |
| private Automaton replaceSep(Automaton a) { |
| |
| int numStates = a.getNumStates(); |
| Automaton.Builder result = new Automaton.Builder(numStates, a.getNumTransitions()); |
| // Copy all states over |
| result.copyStates(a); |
| |
| // Go in reverse topo sort so we know we only have to |
| // make one pass: |
| Transition t = new Transition(); |
| int[] topoSortStates = Operations.topoSortStates(a); |
| for(int i=0;i<topoSortStates.length;i++) { |
| int state = topoSortStates[topoSortStates.length-1-i]; |
| int count = a.initTransition(state, t); |
| for(int j=0;j<count;j++) { |
| a.getNextTransition(t); |
| if (t.min == TokenStreamToAutomaton.POS_SEP) { |
| assert t.max == TokenStreamToAutomaton.POS_SEP; |
| if (preserveSep) { |
| // Remap to SEP_LABEL: |
| result.addTransition(state, t.dest, SEP_LABEL); |
| } else { |
| result.addEpsilon(state, t.dest); |
| } |
| } else if (t.min == TokenStreamToAutomaton.HOLE) { |
| assert t.max == TokenStreamToAutomaton.HOLE; |
| |
| // Just remove the hole: there will then be two |
| // SEP tokens next to each other, which will only |
| // match another hole at search time. Note that |
| // it will also match an empty-string token ... if |
| // that's somehow a problem we can always map HOLE |
| // to a dedicated byte (and escape it in the |
| // input). |
| result.addEpsilon(state, t.dest); |
| } else { |
| result.addTransition(state, t.dest, t.min, t.max); |
| } |
| } |
| } |
| |
| return result.finish(); |
| } |
| |
| /** Used by subclass to change the lookup automaton, if |
| * necessary. */ |
| protected Automaton convertAutomaton(Automaton a) { |
| return a; |
| } |
| |
| TokenStreamToAutomaton getTokenStreamToAutomaton() { |
| final TokenStreamToAutomaton tsta = new TokenStreamToAutomaton(); |
| tsta.setPreservePositionIncrements(preservePositionIncrements); |
| tsta.setFinalOffsetGapAsHole(true); |
| return tsta; |
| } |
| |
| private static class AnalyzingComparator implements Comparator<BytesRef> { |
| |
| private final boolean hasPayloads; |
| |
| public AnalyzingComparator(boolean hasPayloads) { |
| this.hasPayloads = hasPayloads; |
| } |
| |
| private final ByteArrayDataInput readerA = new ByteArrayDataInput(); |
| private final ByteArrayDataInput readerB = new ByteArrayDataInput(); |
| private final BytesRef scratchA = new BytesRef(); |
| private final BytesRef scratchB = new BytesRef(); |
| |
| @Override |
| public int compare(BytesRef a, BytesRef b) { |
| |
| // First by analyzed form: |
| readerA.reset(a.bytes, a.offset, a.length); |
| scratchA.length = readerA.readShort(); |
| scratchA.bytes = a.bytes; |
| scratchA.offset = readerA.getPosition(); |
| |
| readerB.reset(b.bytes, b.offset, b.length); |
| scratchB.bytes = b.bytes; |
| scratchB.length = readerB.readShort(); |
| scratchB.offset = readerB.getPosition(); |
| |
| int cmp = scratchA.compareTo(scratchB); |
| if (cmp != 0) { |
| return cmp; |
| } |
| readerA.skipBytes(scratchA.length); |
| readerB.skipBytes(scratchB.length); |
| |
| // Next by cost: |
| long aCost = readerA.readInt(); |
| long bCost = readerB.readInt(); |
| assert decodeWeight(aCost) >= 0; |
| assert decodeWeight(bCost) >= 0; |
| if (aCost < bCost) { |
| return -1; |
| } else if (aCost > bCost) { |
| return 1; |
| } |
| |
| // Finally by surface form: |
| if (hasPayloads) { |
| scratchA.length = readerA.readShort(); |
| scratchB.length = readerB.readShort(); |
| scratchA.offset = readerA.getPosition(); |
| scratchB.offset = readerB.getPosition(); |
| } else { |
| scratchA.offset = readerA.getPosition(); |
| scratchB.offset = readerB.getPosition(); |
| scratchA.length = readerA.length() - readerA.getPosition(); |
| scratchB.length = readerB.length() - readerB.getPosition(); |
| } |
| assert scratchA.isValid(); |
| assert scratchB.isValid(); |
| |
| return scratchA.compareTo(scratchB); |
| } |
| } |
| |
| @Override |
| public void build(InputIterator iterator) throws IOException { |
| if (iterator.hasContexts()) { |
| throw new IllegalArgumentException("this suggester doesn't support contexts"); |
| } |
| |
| hasPayloads = iterator.hasPayloads(); |
| |
| OfflineSorter sorter = new OfflineSorter(tempDir, tempFileNamePrefix, new AnalyzingComparator(hasPayloads)); |
| |
| IndexOutput tempInput = tempDir.createTempOutput(tempFileNamePrefix, "input", IOContext.DEFAULT); |
| |
| OfflineSorter.ByteSequencesWriter writer = new OfflineSorter.ByteSequencesWriter(tempInput); |
| OfflineSorter.ByteSequencesReader reader = null; |
| BytesRefBuilder scratch = new BytesRefBuilder(); |
| |
| TokenStreamToAutomaton ts2a = getTokenStreamToAutomaton(); |
| |
| String tempSortedFileName = null; |
| |
| count = 0; |
| byte buffer[] = new byte[8]; |
| try { |
| ByteArrayDataOutput output = new ByteArrayDataOutput(buffer); |
| |
| for (BytesRef surfaceForm; (surfaceForm = iterator.next()) != null;) { |
| LimitedFiniteStringsIterator finiteStrings = |
| new LimitedFiniteStringsIterator(toAutomaton(surfaceForm, ts2a), maxGraphExpansions); |
| |
| for (IntsRef string; (string = finiteStrings.next()) != null; count++) { |
| Util.toBytesRef(string, scratch); |
| |
| // length of the analyzed text (FST input) |
| if (scratch.length() > Short.MAX_VALUE-2) { |
| throw new IllegalArgumentException("cannot handle analyzed forms > " + (Short.MAX_VALUE-2) + " in length (got " + scratch.length() + ")"); |
| } |
| short analyzedLength = (short) scratch.length(); |
| |
| // compute the required length: |
| // analyzed sequence + weight (4) + surface + analyzedLength (short) |
| int requiredLength = analyzedLength + 4 + surfaceForm.length + 2; |
| |
| BytesRef payload; |
| |
| if (hasPayloads) { |
| if (surfaceForm.length > (Short.MAX_VALUE-2)) { |
| throw new IllegalArgumentException("cannot handle surface form > " + (Short.MAX_VALUE-2) + " in length (got " + surfaceForm.length + ")"); |
| } |
| payload = iterator.payload(); |
| // payload + surfaceLength (short) |
| requiredLength += payload.length + 2; |
| } else { |
| payload = null; |
| } |
| |
| buffer = ArrayUtil.grow(buffer, requiredLength); |
| |
| output.reset(buffer); |
| |
| output.writeShort(analyzedLength); |
| |
| output.writeBytes(scratch.bytes(), 0, scratch.length()); |
| |
| output.writeInt(encodeWeight(iterator.weight())); |
| |
| if (hasPayloads) { |
| for(int i=0;i<surfaceForm.length;i++) { |
| if (surfaceForm.bytes[i] == PAYLOAD_SEP) { |
| throw new IllegalArgumentException("surface form cannot contain unit separator character U+001F; this character is reserved"); |
| } |
| } |
| output.writeShort((short) surfaceForm.length); |
| output.writeBytes(surfaceForm.bytes, surfaceForm.offset, surfaceForm.length); |
| output.writeBytes(payload.bytes, payload.offset, payload.length); |
| } else { |
| output.writeBytes(surfaceForm.bytes, surfaceForm.offset, surfaceForm.length); |
| } |
| |
| assert output.getPosition() == requiredLength: output.getPosition() + " vs " + requiredLength; |
| writer.write(buffer, 0, output.getPosition()); |
| } |
| |
| maxAnalyzedPathsForOneInput = Math.max(maxAnalyzedPathsForOneInput, finiteStrings.size()); |
| } |
| CodecUtil.writeFooter(tempInput); |
| writer.close(); |
| |
| // Sort all input/output pairs (required by FST.Builder): |
| tempSortedFileName = sorter.sort(tempInput.getName()); |
| |
| // Free disk space: |
| tempDir.deleteFile(tempInput.getName()); |
| |
| reader = new OfflineSorter.ByteSequencesReader(tempDir.openChecksumInput(tempSortedFileName, IOContext.READONCE), tempSortedFileName); |
| |
| PairOutputs<Long,BytesRef> outputs = new PairOutputs<>(PositiveIntOutputs.getSingleton(), ByteSequenceOutputs.getSingleton()); |
| Builder<Pair<Long,BytesRef>> builder = new Builder<>(FST.INPUT_TYPE.BYTE1, outputs); |
| |
| // Build FST: |
| BytesRefBuilder previousAnalyzed = null; |
| BytesRefBuilder analyzed = new BytesRefBuilder(); |
| BytesRef surface = new BytesRef(); |
| IntsRefBuilder scratchInts = new IntsRefBuilder(); |
| ByteArrayDataInput input = new ByteArrayDataInput(); |
| |
| // Used to remove duplicate surface forms (but we |
| // still index the hightest-weight one). We clear |
| // this when we see a new analyzed form, so it cannot |
| // grow unbounded (at most 256 entries): |
| Set<BytesRef> seenSurfaceForms = new HashSet<>(); |
| |
| int dedup = 0; |
| while (true) { |
| BytesRef bytes = reader.next(); |
| if (bytes == null) { |
| break; |
| } |
| input.reset(bytes.bytes, bytes.offset, bytes.length); |
| short analyzedLength = input.readShort(); |
| analyzed.grow(analyzedLength+2); |
| input.readBytes(analyzed.bytes(), 0, analyzedLength); |
| analyzed.setLength(analyzedLength); |
| |
| long cost = input.readInt(); |
| |
| surface.bytes = bytes.bytes; |
| if (hasPayloads) { |
| surface.length = input.readShort(); |
| surface.offset = input.getPosition(); |
| } else { |
| surface.offset = input.getPosition(); |
| surface.length = bytes.length - surface.offset; |
| } |
| |
| if (previousAnalyzed == null) { |
| previousAnalyzed = new BytesRefBuilder(); |
| previousAnalyzed.copyBytes(analyzed.get()); |
| seenSurfaceForms.add(BytesRef.deepCopyOf(surface)); |
| } else if (analyzed.get().equals(previousAnalyzed.get())) { |
| dedup++; |
| if (dedup >= maxSurfaceFormsPerAnalyzedForm) { |
| // More than maxSurfaceFormsPerAnalyzedForm |
| // dups: skip the rest: |
| continue; |
| } |
| if (seenSurfaceForms.contains(surface)) { |
| continue; |
| } |
| seenSurfaceForms.add(BytesRef.deepCopyOf(surface)); |
| } else { |
| dedup = 0; |
| previousAnalyzed.copyBytes(analyzed); |
| seenSurfaceForms.clear(); |
| seenSurfaceForms.add(BytesRef.deepCopyOf(surface)); |
| } |
| |
| // TODO: I think we can avoid the extra 2 bytes when |
| // there is no dup (dedup==0), but we'd have to fix |
| // the exactFirst logic ... which would be sort of |
| // hairy because we'd need to special case the two |
| // (dup/not dup)... |
| |
| // NOTE: must be byte 0 so we sort before whatever |
| // is next |
| analyzed.append((byte) 0); |
| analyzed.append((byte) dedup); |
| |
| Util.toIntsRef(analyzed.get(), scratchInts); |
| //System.out.println("ADD: " + scratchInts + " -> " + cost + ": " + surface.utf8ToString()); |
| if (!hasPayloads) { |
| builder.add(scratchInts.get(), outputs.newPair(cost, BytesRef.deepCopyOf(surface))); |
| } else { |
| int payloadOffset = input.getPosition() + surface.length; |
| int payloadLength = bytes.length - payloadOffset; |
| BytesRef br = new BytesRef(surface.length + 1 + payloadLength); |
| System.arraycopy(surface.bytes, surface.offset, br.bytes, 0, surface.length); |
| br.bytes[surface.length] = PAYLOAD_SEP; |
| System.arraycopy(bytes.bytes, payloadOffset, br.bytes, surface.length+1, payloadLength); |
| br.length = br.bytes.length; |
| builder.add(scratchInts.get(), outputs.newPair(cost, br)); |
| } |
| } |
| fst = builder.finish(); |
| |
| //Util.dotToFile(fst, "/tmp/suggest.dot"); |
| } finally { |
| IOUtils.closeWhileHandlingException(reader, writer); |
| IOUtils.deleteFilesIgnoringExceptions(tempDir, tempInput.getName(), tempSortedFileName); |
| } |
| } |
| |
| @Override |
| public boolean store(DataOutput output) throws IOException { |
| output.writeVLong(count); |
| if (fst == null) { |
| return false; |
| } |
| |
| fst.save(output, output); |
| output.writeVInt(maxAnalyzedPathsForOneInput); |
| output.writeByte((byte) (hasPayloads ? 1 : 0)); |
| return true; |
| } |
| |
| @Override |
| public boolean load(DataInput input) throws IOException { |
| count = input.readVLong(); |
| this.fst = new FST<>(input, input, new PairOutputs<>(PositiveIntOutputs.getSingleton(), ByteSequenceOutputs.getSingleton())); |
| maxAnalyzedPathsForOneInput = input.readVInt(); |
| hasPayloads = input.readByte() == 1; |
| return true; |
| } |
| |
| private LookupResult getLookupResult(Long output1, BytesRef output2, CharsRefBuilder spare) { |
| LookupResult result; |
| if (hasPayloads) { |
| int sepIndex = -1; |
| for(int i=0;i<output2.length;i++) { |
| if (output2.bytes[output2.offset+i] == PAYLOAD_SEP) { |
| sepIndex = i; |
| break; |
| } |
| } |
| assert sepIndex != -1; |
| spare.grow(sepIndex); |
| final int payloadLen = output2.length - sepIndex - 1; |
| spare.copyUTF8Bytes(output2.bytes, output2.offset, sepIndex); |
| BytesRef payload = new BytesRef(payloadLen); |
| System.arraycopy(output2.bytes, sepIndex+1, payload.bytes, 0, payloadLen); |
| payload.length = payloadLen; |
| result = new LookupResult(spare.toString(), decodeWeight(output1), payload); |
| } else { |
| spare.grow(output2.length); |
| spare.copyUTF8Bytes(output2); |
| result = new LookupResult(spare.toString(), decodeWeight(output1)); |
| } |
| |
| return result; |
| } |
| |
| private boolean sameSurfaceForm(BytesRef key, BytesRef output2) { |
| if (hasPayloads) { |
| // output2 has at least PAYLOAD_SEP byte: |
| if (key.length >= output2.length) { |
| return false; |
| } |
| for(int i=0;i<key.length;i++) { |
| if (key.bytes[key.offset+i] != output2.bytes[output2.offset+i]) { |
| return false; |
| } |
| } |
| return output2.bytes[output2.offset + key.length] == PAYLOAD_SEP; |
| } else { |
| return key.bytesEquals(output2); |
| } |
| } |
| |
| @Override |
| public List<LookupResult> lookup(final CharSequence key, Set<BytesRef> contexts, boolean onlyMorePopular, int num) { |
| assert num > 0; |
| |
| if (onlyMorePopular) { |
| throw new IllegalArgumentException("this suggester only works with onlyMorePopular=false"); |
| } |
| if (contexts != null) { |
| throw new IllegalArgumentException("this suggester doesn't support contexts"); |
| } |
| if (fst == null) { |
| return Collections.emptyList(); |
| } |
| |
| //System.out.println("lookup key=" + key + " num=" + num); |
| for (int i = 0; i < key.length(); i++) { |
| if (key.charAt(i) == 0x1E) { |
| throw new IllegalArgumentException("lookup key cannot contain HOLE character U+001E; this character is reserved"); |
| } |
| if (key.charAt(i) == 0x1F) { |
| throw new IllegalArgumentException("lookup key cannot contain unit separator character U+001F; this character is reserved"); |
| } |
| } |
| final BytesRef utf8Key = new BytesRef(key); |
| try { |
| Automaton lookupAutomaton = toLookupAutomaton(key); |
| |
| final CharsRefBuilder spare = new CharsRefBuilder(); |
| |
| //System.out.println(" now intersect exactFirst=" + exactFirst); |
| |
| // Intersect automaton w/ suggest wFST and get all |
| // prefix starting nodes & their outputs: |
| //final PathIntersector intersector = getPathIntersector(lookupAutomaton, fst); |
| |
| //System.out.println(" prefixPaths: " + prefixPaths.size()); |
| |
| BytesReader bytesReader = fst.getBytesReader(); |
| |
| FST.Arc<Pair<Long,BytesRef>> scratchArc = new FST.Arc<>(); |
| |
| final List<LookupResult> results = new ArrayList<>(); |
| |
| List<FSTUtil.Path<Pair<Long,BytesRef>>> prefixPaths = FSTUtil.intersectPrefixPaths(convertAutomaton(lookupAutomaton), fst); |
| |
| if (exactFirst) { |
| |
| int count = 0; |
| for (FSTUtil.Path<Pair<Long,BytesRef>> path : prefixPaths) { |
| if (fst.findTargetArc(END_BYTE, path.fstNode, scratchArc, bytesReader) != null) { |
| // This node has END_BYTE arc leaving, meaning it's an |
| // "exact" match: |
| count++; |
| } |
| } |
| |
| // Searcher just to find the single exact only |
| // match, if present: |
| Util.TopNSearcher<Pair<Long,BytesRef>> searcher; |
| searcher = new Util.TopNSearcher<>(fst, count * maxSurfaceFormsPerAnalyzedForm, count * maxSurfaceFormsPerAnalyzedForm, weightComparator); |
| |
| // NOTE: we could almost get away with only using |
| // the first start node. The only catch is if |
| // maxSurfaceFormsPerAnalyzedForm had kicked in and |
| // pruned our exact match from one of these nodes |
| // ...: |
| for (FSTUtil.Path<Pair<Long,BytesRef>> path : prefixPaths) { |
| if (fst.findTargetArc(END_BYTE, path.fstNode, scratchArc, bytesReader) != null) { |
| // This node has END_BYTE arc leaving, meaning it's an |
| // "exact" match: |
| searcher.addStartPaths(scratchArc, fst.outputs.add(path.output, scratchArc.output()), false, path.input); |
| } |
| } |
| |
| TopResults<Pair<Long,BytesRef>> completions = searcher.search(); |
| assert completions.isComplete; |
| |
| // NOTE: this is rather inefficient: we enumerate |
| // every matching "exactly the same analyzed form" |
| // path, and then do linear scan to see if one of |
| // these exactly matches the input. It should be |
| // possible (though hairy) to do something similar |
| // to getByOutput, since the surface form is encoded |
| // into the FST output, so we more efficiently hone |
| // in on the exact surface-form match. Still, I |
| // suspect very little time is spent in this linear |
| // seach: it's bounded by how many prefix start |
| // nodes we have and the |
| // maxSurfaceFormsPerAnalyzedForm: |
| for(Result<Pair<Long,BytesRef>> completion : completions) { |
| BytesRef output2 = completion.output.output2; |
| if (sameSurfaceForm(utf8Key, output2)) { |
| results.add(getLookupResult(completion.output.output1, output2, spare)); |
| break; |
| } |
| } |
| |
| if (results.size() == num) { |
| // That was quick: |
| return results; |
| } |
| } |
| |
| Util.TopNSearcher<Pair<Long,BytesRef>> searcher; |
| searcher = new Util.TopNSearcher<Pair<Long,BytesRef>>(fst, |
| num - results.size(), |
| num * maxAnalyzedPathsForOneInput, |
| weightComparator) { |
| private final Set<BytesRef> seen = new HashSet<>(); |
| |
| @Override |
| protected boolean acceptResult(IntsRef input, Pair<Long,BytesRef> output) { |
| |
| // Dedup: when the input analyzes to a graph we |
| // can get duplicate surface forms: |
| if (seen.contains(output.output2)) { |
| return false; |
| } |
| seen.add(output.output2); |
| |
| if (!exactFirst) { |
| return true; |
| } else { |
| // In exactFirst mode, don't accept any paths |
| // matching the surface form since that will |
| // create duplicate results: |
| if (sameSurfaceForm(utf8Key, output.output2)) { |
| // We found exact match, which means we should |
| // have already found it in the first search: |
| assert results.size() == 1; |
| return false; |
| } else { |
| return true; |
| } |
| } |
| } |
| }; |
| |
| prefixPaths = getFullPrefixPaths(prefixPaths, lookupAutomaton, fst); |
| |
| for (FSTUtil.Path<Pair<Long,BytesRef>> path : prefixPaths) { |
| searcher.addStartPaths(path.fstNode, path.output, true, path.input); |
| } |
| |
| TopResults<Pair<Long,BytesRef>> completions = searcher.search(); |
| assert completions.isComplete; |
| |
| for(Result<Pair<Long,BytesRef>> completion : completions) { |
| |
| LookupResult result = getLookupResult(completion.output.output1, completion.output.output2, spare); |
| |
| // TODO: for fuzzy case would be nice to return |
| // how many edits were required |
| |
| //System.out.println(" result=" + result); |
| results.add(result); |
| |
| if (results.size() == num) { |
| // In the exactFirst=true case the search may |
| // produce one extra path |
| break; |
| } |
| } |
| |
| return results; |
| } catch (IOException bogus) { |
| throw new RuntimeException(bogus); |
| } |
| } |
| |
| @Override |
| public long getCount() { |
| return count; |
| } |
| |
| /** Returns all prefix paths to initialize the search. */ |
| protected List<FSTUtil.Path<Pair<Long,BytesRef>>> getFullPrefixPaths(List<FSTUtil.Path<Pair<Long,BytesRef>>> prefixPaths, |
| Automaton lookupAutomaton, |
| FST<Pair<Long,BytesRef>> fst) |
| throws IOException { |
| return prefixPaths; |
| } |
| |
| final Automaton toAutomaton(final BytesRef surfaceForm, final TokenStreamToAutomaton ts2a) throws IOException { |
| // Analyze surface form: |
| Automaton automaton; |
| try (TokenStream ts = indexAnalyzer.tokenStream("", surfaceForm.utf8ToString())) { |
| |
| // Create corresponding automaton: labels are bytes |
| // from each analyzed token, with byte 0 used as |
| // separator between tokens: |
| automaton = ts2a.toAutomaton(ts); |
| } |
| |
| automaton = replaceSep(automaton); |
| automaton = convertAutomaton(automaton); |
| |
| // TODO: LUCENE-5660 re-enable this once we disallow massive suggestion strings |
| // assert SpecialOperations.isFinite(automaton); |
| |
| // Get all paths from the automaton (there can be |
| // more than one path, eg if the analyzer created a |
| // graph using SynFilter or WDF): |
| return automaton; |
| } |
| |
| final Automaton toLookupAutomaton(final CharSequence key) throws IOException { |
| // TODO: is there a Reader from a CharSequence? |
| // Turn tokenstream into automaton: |
| Automaton automaton = null; |
| try (TokenStream ts = queryAnalyzer.tokenStream("", key.toString())) { |
| automaton = getTokenStreamToAutomaton().toAutomaton(ts); |
| } |
| |
| automaton = replaceSep(automaton); |
| |
| // TODO: we can optimize this somewhat by determinizing |
| // while we convert |
| automaton = Operations.determinize(automaton, DEFAULT_DETERMINIZE_WORK_LIMIT); |
| return automaton; |
| } |
| |
| /** |
| * Returns the weight associated with an input string, |
| * or null if it does not exist. |
| */ |
| public Object get(CharSequence key) { |
| throw new UnsupportedOperationException(); |
| } |
| |
| /** cost -> weight */ |
| private static int decodeWeight(long encoded) { |
| return (int)(Integer.MAX_VALUE - encoded); |
| } |
| |
| /** weight -> cost */ |
| private static int encodeWeight(long value) { |
| if (value < 0 || value > Integer.MAX_VALUE) { |
| throw new UnsupportedOperationException("cannot encode value: " + value); |
| } |
| return Integer.MAX_VALUE - (int)value; |
| } |
| |
| static final Comparator<Pair<Long,BytesRef>> weightComparator = new Comparator<Pair<Long,BytesRef>> () { |
| @Override |
| public int compare(Pair<Long,BytesRef> left, Pair<Long,BytesRef> right) { |
| return left.output1.compareTo(right.output1); |
| } |
| }; |
| } |
