--- /dev/null
+package org.apache.lucene.search.suggest.fst;
+
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.IntsRef;
+import org.apache.lucene.util.fst.Builder;
+import org.apache.lucene.util.fst.FST;
+import org.apache.lucene.util.fst.FST.Arc;
+import org.apache.lucene.util.fst.NoOutputs;
+import org.apache.lucene.util.fst.Outputs;
+
+import org.apache.lucene.search.suggest.Lookup;
+import org.apache.lucene.search.suggest.tst.TSTLookup;
+import org.apache.lucene.search.spell.TermFreqIterator;
+import org.apache.lucene.store.InputStreamDataInput;
+import org.apache.lucene.store.OutputStreamDataOutput;
+
+/**
+ * Finite state automata based implementation of {@link Lookup} query
+ * suggestion/ autocomplete interface.
+ *
+ * <h2>Implementation details</h2>
+ *
+ * <p>The construction step in {@link #build(TermFreqIterator)} works as follows:
+ * <ul>
+ * <li>A set of input terms (String) and weights (float) is given.</li>
+ * <li>The range of weights is determined and then all weights are discretized into a fixed set
+ * of values ({@link #buckets}).
+ * Note that this means that minor changes in weights may be lost during automaton construction.
+ * In general, this is not a big problem because the "priorities" of completions can be split
+ * into a fixed set of classes (even as rough as: very frequent, frequent, baseline, marginal).
+ * If you need exact, fine-grained weights, use {@link TSTLookup} instead.<li>
+ * <li>All terms in the input are preprended with a synthetic pseudo-character being the weight
+ * of that term. For example a term <code>abc</code> with a discretized weight equal '1' would
+ * become <code>1abc</code>.</li>
+ * <li>The terms are sorted by their raw value of utf16 character values (including the synthetic
+ * term in front).</li>
+ * <li>A finite state automaton ({@link FST}) is constructed from the input. The root node has
+ * arcs labeled with all possible weights. We cache all these arcs, highest-weight first.</li>
+ * </ul>
+ *
+ * <p>At runtime, in {@link #lookup(String, boolean, int)}, the automaton is utilized as follows:
+ * <ul>
+ * <li>For each possible term weight encoded in the automaton (cached arcs from the root above),
+ * starting with the highest one, we descend along the path of the input key. If the key is not
+ * a prefix of a sequence in the automaton (path ends prematurely), we exit immediately.
+ * No completions.
+ * <li>Otherwise, we have found an internal automaton node that ends the key. <b>The entire
+ * subautomaton (all paths) starting from this node form the key's completions.</b> We start
+ * the traversal of this subautomaton. Every time we reach a final state (arc), we add a single
+ * suggestion to the list of results (the weight of this suggestion is constant and equal to the
+ * root path we started from). The tricky part is that because automaton edges are sorted and
+ * we scan depth-first, we can terminate the entire procedure as soon as we collect enough
+ * suggestions the user requested.
+ * <li>In case the number of suggestions collected in the step above is still insufficient,
+ * we proceed to the next (smaller) weight leaving the root node and repeat the same
+ * algorithm again.
+ * </li>
+ * </ul>
+ *
+ * <h2>Runtime behavior and performance characteristic</h2>
+ *
+ * <p>The algorithm described above is optimized for finding suggestions to short prefixes
+ * in a top-weights-first order. This is probably the most common use case: it allows
+ * presenting suggestions early and sorts them by the global frequency (and then alphabetically).
+ *
+ * <p>If there is an exact match in the automaton, it is returned first on the results
+ * list (even with by-weight sorting).
+ *
+ * <p>Note that the maximum lookup time for <b>any prefix</b>
+ * is the time of descending to the subtree, plus traversal of the subtree up to the number
+ * of requested suggestions (because they are already presorted by weight on the root level
+ * and alphabetically at any node level).
+ *
+ * <p>To order alphabetically only (no ordering by priorities), use identical term weights
+ * for all terms. Alphabetical suggestions are returned even if non-constant weights are
+ * used, but the algorithm for doing this is suboptimal.
+ *
+ * <p>"alphabetically" in any of the documentation above indicates utf16 codepoint order,
+ * nothing else.
+ */
+public class FSTLookup extends Lookup {
+
+ public FSTLookup() {
+ this(10, true);
+ }
+
+ public FSTLookup(int buckets, boolean exactMatchFirst) {
+ this.buckets = buckets;
+ this.exactMatchFirst = exactMatchFirst;
+ }
+
+ /** A structure for a single entry (for sorting/ preprocessing). */
+ private static class Entry {
+ char [] term;
+ float weight;
+
+ public Entry(char [] term, float freq) {
+ this.term = term;
+ this.weight = freq;
+ }
+ }
+
+ /** Serialized automaton file name (storage). */
+ public static final String FILENAME = "fst.dat";
+
+ /** An empty result. */
+ private static final List<LookupResult> EMPTY_RESULT = Collections.emptyList();
+
+ /**
+ * The number of separate buckets for weights (discretization). The more buckets,
+ * the more fine-grained term weights (priorities) can be assigned. The speed of lookup
+ * will not decrease for prefixes which have highly-weighted completions (because these
+ * are filled-in first), but will decrease significantly for low-weighted terms (but
+ * these should be infrequent, so it is all right).
+ *
+ * <p>The number of buckets must be within [1, 255] range.
+ */
+ private final int buckets;
+
+ /**
+ * If <code>true</code>, exact suggestions are returned first, even if they are prefixes
+ * of other strings in the automaton (possibly with larger weights).
+ */
+ private final boolean exactMatchFirst;
+
+ /**
+ * Finite state automaton encoding all the lookup terms. See class
+ * notes for details.
+ */
+ private FST<Object> automaton;
+
+ /**
+ * An array of arcs leaving the root automaton state and encoding weights of all
+ * completions in their sub-trees.
+ */
+ private Arc<Object> [] rootArcs;
+
+ /* */
+ @Override
+ public void build(TermFreqIterator tfit) throws IOException {
+ // Buffer the input because we will need it twice: for calculating
+ // weights distribution and for the actual automata building.
+ List<Entry> entries = new ArrayList<Entry>();
+ while (tfit.hasNext()) {
+ String term = tfit.next();
+ char [] termChars = new char [term.length() + 1]; // add padding for weight.
+ for (int i = 0; i < term.length(); i++)
+ termChars[i + 1] = term.charAt(i);
+ entries.add(new Entry(termChars, tfit.freq()));
+ }
+
+ // Distribute weights into at most N buckets. This is a form of discretization to
+ // limit the number of possible weights so that they can be efficiently encoded in the
+ // automaton.
+ //
+ // It is assumed the distribution of weights is _linear_ so proportional division
+ // of [min, max] range will be enough here. Other approaches could be to sort
+ // weights and divide into proportional ranges.
+ if (entries.size() > 0) {
+ redistributeWeightsProportionalMinMax(entries, buckets);
+ encodeWeightPrefix(entries);
+ }
+
+ // Build the automaton (includes input sorting) and cache root arcs in order from the highest,
+ // to the lowest weight.
+ this.automaton = buildAutomaton(entries);
+ cacheRootArcs();
+ }
+
+ /**
+ * Cache the root node's output arcs starting with completions with the highest weights.
+ */
+ @SuppressWarnings("unchecked")
+ private void cacheRootArcs() throws IOException {
+ if (automaton != null) {
+ List<Arc<Object>> rootArcs = new ArrayList<Arc<Object>>();
+ Arc<Object> arc = automaton.getFirstArc(new Arc<Object>());
+ automaton.readFirstTargetArc(arc, arc);
+ while (true) {
+ rootArcs.add(new Arc<Object>().copyFrom(arc));
+ if (arc.isLast())
+ break;
+ automaton.readNextArc(arc);
+ }
+
+ Collections.reverse(rootArcs); // we want highest weights first.
+ this.rootArcs = rootArcs.toArray(new Arc[rootArcs.size()]);
+ }
+ }
+
+ /**
+ * Not implemented.
+ */
+ @Override
+ public boolean add(String key, Object value) {
+ // This implementation does not support ad-hoc additions (all input
+ // must be sorted for the builder).
+ return false;
+ }
+
+ /**
+ * Get the (approximated) weight of a single key (if there is a perfect match
+ * for it in the automaton).
+ *
+ * @return Returns the approximated weight of the input key or <code>null</code>
+ * if not found.
+ */
+ @Override
+ public Float get(String key) {
+ return getExactMatchStartingFromRootArc(0, key);
+ }
+
+ /**
+ * Returns the first exact match by traversing root arcs, starting from
+ * the arc <code>i</code>.
+ *
+ * @param i The first root arc index in {@link #rootArcs} to consider when
+ * matching.
+ */
+ private Float getExactMatchStartingFromRootArc(int i, String key) {
+ // Get the UTF-8 bytes representation of the input key.
+ try {
+ final FST.Arc<Object> scratch = new FST.Arc<Object>();
+ for (; i < rootArcs.length; i++) {
+ final FST.Arc<Object> rootArc = rootArcs[i];
+ final FST.Arc<Object> arc = scratch.copyFrom(rootArc);
+
+ // Descend into the automaton using the key as prefix.
+ if (descendWithPrefix(arc, key)) {
+ automaton.readFirstTargetArc(arc, arc);
+ if (arc.label == FST.END_LABEL) {
+ // Prefix-encoded weight.
+ return rootArc.label / (float) buckets;
+ }
+ }
+ }
+ } catch (IOException e) {
+ // Should never happen, but anyway.
+ throw new RuntimeException(e);
+ }
+
+ return null;
+ }
+
+ /**
+ * Lookup autocomplete suggestions to <code>key</code>.
+ *
+ * @param key The prefix to which suggestions should be sought.
+ * @param onlyMorePopular Return most popular suggestions first. This is the default
+ * behavior for this implementation. Setting it to <code>false</code> has no effect (use
+ * constant term weights to sort alphabetically only).
+ * @param num At most this number of suggestions will be returned.
+ * @return Returns the suggestions, sorted by their approximated weight first (decreasing)
+ * and then alphabetically (utf16 codepoint order).
+ */
+ @Override
+ public List<LookupResult> lookup(String key, boolean onlyMorePopular, int num) {
+ if (key.length() == 0 || automaton == null) {
+ // Keep the result an ArrayList to keep calls monomorphic.
+ return EMPTY_RESULT;
+ }
+
+ try {
+ if (!onlyMorePopular && rootArcs.length > 1) {
+ // We could emit a warning here (?). An optimal strategy for alphabetically sorted
+ // suggestions would be to add them with a constant weight -- this saves unnecessary
+ // traversals and sorting.
+ return lookupSortedAlphabetically(key, num);
+ } else {
+ return lookupSortedByWeight(key, num, false);
+ }
+ } catch (IOException e) {
+ // Should never happen, but anyway.
+ throw new RuntimeException(e);
+ }
+ }
+
+ /**
+ * Lookup suggestions sorted alphabetically <b>if weights are not constant</b>. This
+ * is a workaround: in general, use constant weights for alphabetically sorted result.
+ */
+ private List<LookupResult> lookupSortedAlphabetically(String key, int num) throws IOException {
+ // Greedily get num results from each weight branch.
+ List<LookupResult> res = lookupSortedByWeight(key, num, true);
+
+ // Sort and trim.
+ Collections.sort(res, new Comparator<LookupResult>() {
+ // not till java6 @Override
+ public int compare(LookupResult o1, LookupResult o2) {
+ return o1.key.compareTo(o2.key);
+ }
+ });
+ if (res.size() > num) {
+ res = res.subList(0, num);
+ }
+ return res;
+ }
+
+ /**
+ * Lookup suggestions sorted by weight (descending order).
+ *
+ * @param collectAll If <code>true</code>, the routine terminates immediately when <code>num</code>
+ * suggestions have been collected. If <code>false</code>, it will collect suggestions from
+ * all weight arcs (needed for {@link #lookupSortedAlphabetically}.
+ */
+ private ArrayList<LookupResult> lookupSortedByWeight(String key, int num, boolean collectAll) throws IOException {
+ // Don't overallocate the results buffers. This also serves the purpose of allowing
+ // the user of this class to request all matches using Integer.MAX_VALUE as the number
+ // of results.
+ final ArrayList<LookupResult> res = new ArrayList<LookupResult>(Math.min(10, num));
+ final StringBuilder output = new StringBuilder(key);
+ final int matchLength = key.length() - 1;
+
+ for (int i = 0; i < rootArcs.length; i++) {
+ final FST.Arc<Object> rootArc = rootArcs[i];
+ final FST.Arc<Object> arc = new FST.Arc<Object>().copyFrom(rootArc);
+
+ // Descend into the automaton using the key as prefix.
+ if (descendWithPrefix(arc, key)) {
+ // Prefix-encoded weight.
+ final float weight = rootArc.label / (float) buckets;
+
+ // A subgraph starting from the current node has the completions
+ // of the key prefix. The arc we're at is the last key's byte,
+ // so we will collect it too.
+ output.setLength(matchLength);
+ if (collect(res, num, weight, output, arc) && !collectAll) {
+ // We have enough suggestions to return immediately. Keep on looking for an
+ // exact match, if requested.
+ if (exactMatchFirst) {
+ if (!checkExistingAndReorder(res, key)) {
+ Float exactMatchWeight = getExactMatchStartingFromRootArc(i, key);
+ if (exactMatchWeight != null) {
+ // Insert as the first result and truncate at num.
+ while (res.size() >= num) {
+ res.remove(res.size() - 1);
+ }
+ res.add(0, new LookupResult(key, exactMatchWeight));
+ }
+ }
+ }
+ break;
+ }
+ }
+ }
+ return res;
+ }
+
+ /**
+ * Checks if the list of {@link LookupResult}s already has a <code>key</code>. If so,
+ * reorders that {@link LookupResult} to the first position.
+ *
+ * @return Returns <code>true<code> if and only if <code>list</code> contained <code>key</code>.
+ */
+ private boolean checkExistingAndReorder(ArrayList<LookupResult> list, String key) {
+ // We assume list does not have duplicates (because of how the FST is created).
+ for (int i = list.size(); --i >= 0;) {
+ if (key.equals(list.get(i).key)) {
+ // Key found. Unless already at i==0, remove it and push up front so that the ordering
+ // remains identical with the exception of the exact match.
+ list.add(0, list.remove(i));
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Descend along the path starting at <code>arc</code> and going through
+ * bytes in <code>utf8</code> argument.
+ *
+ * @param arc The starting arc. This argument is modified in-place.
+ * @param term The term to descend with.
+ * @return If <code>true</code>, <code>arc</code> will be set to the arc matching
+ * last byte of <code>utf8</code>. <code>false</code> is returned if no such
+ * prefix <code>utf8</code> exists.
+ */
+ private boolean descendWithPrefix(Arc<Object> arc, String term) throws IOException {
+ final int max = term.length();
+
+ for (int i = 0; i < max; i++) {
+ if (automaton.findTargetArc(term.charAt(i) & 0xffff, arc, arc) == null) {
+ // No matching prefixes, return an empty result.
+ return false;
+ }
+ }
+
+ return true;
+ }
+
+ /**
+ * Recursive collect lookup results from the automaton subgraph starting at <code>arc</code>.
+ *
+ * @param num Maximum number of results needed (early termination).
+ * @param weight Weight of all results found during this collection.
+ */
+ private boolean collect(List<LookupResult> res, int num, float weight, StringBuilder output, Arc<Object> arc) throws IOException {
+ output.append((char) arc.label);
+
+ automaton.readFirstTargetArc(arc, arc);
+ while (true) {
+ if (arc.label == FST.END_LABEL) {
+ res.add(new LookupResult(output.toString(), weight));
+ if (res.size() >= num)
+ return true;
+ } else {
+ int save = output.length();
+ if (collect(res, num, weight, output, new Arc<Object>().copyFrom(arc))) {
+ return true;
+ }
+ output.setLength(save);
+ }
+
+ if (arc.isLast()) {
+ break;
+ }
+ automaton.readNextArc(arc);
+ }
+ return false;
+ }
+
+ /**
+ * Builds the final automaton from a list of entries.
+ */
+ private FST<Object> buildAutomaton(List<Entry> entries) throws IOException {
+ if (entries.size() == 0)
+ return null;
+
+ // Sort by utf16 (raw char value)
+ final Comparator<Entry> comp = new Comparator<Entry>() {
+ public int compare(Entry o1, Entry o2) {
+ char [] ch1 = o1.term;
+ char [] ch2 = o2.term;
+ int len1 = ch1.length;
+ int len2 = ch2.length;
+
+ int max = Math.min(len1, len2);
+ for (int i = 0; i < max; i++) {
+ int v = ch1[i] - ch2[i];
+ if (v != 0) return v;
+ }
+ return len1 - len2;
+ }
+ };
+ Collections.sort(entries, comp);
+
+ // Avoid duplicated identical entries, if possible. This is required because
+ // it breaks automaton construction otherwise.
+ int len = entries.size();
+ int j = 0;
+ for (int i = 1; i < len; i++) {
+ if (comp.compare(entries.get(j), entries.get(i)) != 0) {
+ entries.set(++j, entries.get(i));
+ }
+ }
+ entries = entries.subList(0, j + 1);
+
+ // Build the automaton.
+ final Outputs<Object> outputs = NoOutputs.getSingleton();
+ final Object empty = outputs.getNoOutput();
+ final Builder<Object> builder =
+ new Builder<Object>(FST.INPUT_TYPE.BYTE4, outputs);
+ final IntsRef scratchIntsRef = new IntsRef(10);
+ for (Entry e : entries) {
+ final int termLength = scratchIntsRef.length = e.term.length;
+
+ scratchIntsRef.grow(termLength);
+ final int [] ints = scratchIntsRef.ints;
+ final char [] chars = e.term;
+ for (int i = termLength; --i >= 0;) {
+ ints[i] = chars[i];
+ }
+ builder.add(scratchIntsRef, empty);
+ }
+ return builder.finish();
+ }
+
+ /**
+ * Prepends the entry's weight to each entry, encoded as a single byte, so that the
+ * root automaton node fans out to all possible priorities, starting with the arc that has
+ * the highest weights.
+ */
+ private void encodeWeightPrefix(List<Entry> entries) {
+ for (Entry e : entries) {
+ int weight = (int) e.weight;
+ assert (weight >= 0 && weight <= buckets) :
+ "Weight out of range: " + weight + " [" + buckets + "]";
+
+ // There should be a single empty char reserved in front for the weight.
+ e.term[0] = (char) weight;
+ }
+ }
+
+ /**
+ * Split [min, max] range into buckets, reassigning weights. Entries' weights are
+ * remapped to [0, buckets] range (so, buckets + 1 buckets, actually).
+ */
+ private void redistributeWeightsProportionalMinMax(List<Entry> entries, int buckets) {
+ float min = entries.get(0).weight;
+ float max = min;
+ for (Entry e : entries) {
+ min = Math.min(e.weight, min);
+ max = Math.max(e.weight, max);
+ }
+
+ final float range = max - min;
+ for (Entry e : entries) {
+ e.weight = (int) (buckets * ((e.weight - min) / range)); // int cast equiv. to floor()
+ }
+ }
+
+ /**
+ * Deserialization from disk.
+ */
+ @Override
+ public synchronized boolean load(File storeDir) throws IOException {
+ File data = new File(storeDir, FILENAME);
+ if (!data.exists() || !data.canRead()) {
+ return false;
+ }
+
+ InputStream is = new BufferedInputStream(new FileInputStream(data));
+ try {
+ this.automaton = new FST<Object>(new InputStreamDataInput(is), NoOutputs.getSingleton());
+ cacheRootArcs();
+ } finally {
+ IOUtils.close(is);
+ }
+ return true;
+ }
+
+ /**
+ * Serialization to disk.
+ */
+ @Override
+ public synchronized boolean store(File storeDir) throws IOException {
+ if (!storeDir.exists() || !storeDir.isDirectory() || !storeDir.canWrite()) {
+ return false;
+ }
+
+ if (this.automaton == null)
+ return false;
+
+ File data = new File(storeDir, FILENAME);
+ OutputStream os = new BufferedOutputStream(new FileOutputStream(data));
+ try {
+ this.automaton.save(new OutputStreamDataOutput(os));
+ } finally {
+ IOUtils.close(os);
+ }
+
+ return true;
+ }
+}
fnp / pylucene.git / blobdiff