+++ /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, true);
- }
- } 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, false);
-
- // 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 greedy 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 greedy) throws IOException {
- 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) && greedy) {
- // We have enough suggestion to return immediately. Keep on looking for an
- // exact match, if requested.
- if (exactMatchFirst) {
- Float exactMatchWeight = getExactMatchStartingFromRootArc(i, key);
- if (exactMatchWeight != null) {
- res.add(0, new LookupResult(key, exactMatchWeight));
- while (res.size() > num) {
- res.remove(res.size() - 1);
- }
- }
- }
- break;
- }
- }
- }
- return res;
- }
-
- /**
- * 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