n shortest paths to retrieve top-ranked
+ * suggestions.
+ *
+ * NOTE: Although the {@link TermFreqIterator} API specifies
+ * floating point weights, input weights should be whole numbers.
+ * Input weights will be cast to a java integer, and any
+ * negative, infinite, or NaN values will be rejected.
+ *
+ * @see Util#shortestPaths(FST, FST.Arc, int)
+ */
+public class WFSTCompletionLookup extends Lookup {
+
+ /**
+ * File name for the automaton.
+ *
+ * @see #store(File)
+ * @see #load(File)
+ */
+ private static final String FILENAME = "wfst.bin";
+
+ /**
+ * FST, weights are encoded as costs: (Integer.MAX_VALUE-weight)
+ */
+ // NOTE: like FSTSuggester, this is really a WFSA, if you want to
+ // customize the code to add some output you should use PairOutputs.
+ private FST fst = null;
+
+ /**
+ * True if exact match suggestions should always be returned first.
+ */
+ private final boolean exactFirst;
+
+ /**
+ * Calls {@link #WFSTCompletionLookup(boolean) WFSTCompletionLookup(true)}
+ */
+ public WFSTCompletionLookup() {
+ this(true);
+ }
+
+ /**
+ * Creates a new suggester.
+ *
+ * @param exactFirst true if suggestions that match the
+ * prefix exactly should always be returned first, regardless
+ * of score. This has no performance impact, but could result
+ * in low-quality suggestions.
+ */
+ public WFSTCompletionLookup(boolean exactFirst) {
+ this.exactFirst = exactFirst;
+ }
+
+ @Override
+ public void build(TermFreqIterator iterator) throws IOException {
+ String prefix = getClass().getSimpleName();
+ File directory = Sort.defaultTempDir();
+ File tempInput = File.createTempFile(prefix, ".input", directory);
+ File tempSorted = File.createTempFile(prefix, ".sorted", directory);
+
+ Sort.ByteSequencesWriter writer = new Sort.ByteSequencesWriter(tempInput);
+ Sort.ByteSequencesReader reader = null;
+ BytesRef scratch = new BytesRef();
+
+ boolean success = false;
+ try {
+ byte [] buffer = new byte [0];
+ ByteArrayDataOutput output = new ByteArrayDataOutput(buffer);
+ while (iterator.hasNext()) {
+ String key = iterator.next();
+ UnicodeUtil.UTF16toUTF8(key, 0, key.length(), scratch);
+
+ if (scratch.length + 5 >= buffer.length) {
+ buffer = ArrayUtil.grow(buffer, scratch.length + 5);
+ }
+
+ output.reset(buffer);
+ output.writeBytes(scratch.bytes, scratch.offset, scratch.length);
+ output.writeByte((byte)0); // separator: not used, just for sort order
+ output.writeInt((int)encodeWeight(iterator.freq()));
+ writer.write(buffer, 0, output.getPosition());
+ }
+ writer.close();
+ new Sort().sort(tempInput, tempSorted);
+ reader = new Sort.ByteSequencesReader(tempSorted);
+
+ PositiveIntOutputs outputs = PositiveIntOutputs.getSingleton(true);
+ Builder builder = new Builder(FST.INPUT_TYPE.BYTE1, outputs);
+
+ BytesRef previous = null;
+ BytesRef suggestion = new BytesRef();
+ IntsRef scratchInts = new IntsRef();
+ ByteArrayDataInput input = new ByteArrayDataInput();
+ while (reader.read(scratch)) {
+ suggestion.bytes = scratch.bytes;
+ suggestion.offset = scratch.offset;
+ suggestion.length = scratch.length - 5; // int + separator
+
+ input.reset(scratch.bytes);
+ input.skipBytes(suggestion.length + 1); // suggestion + separator
+ long cost = input.readInt();
+
+ if (previous == null) {
+ previous = new BytesRef();
+ } else if (suggestion.equals(previous)) {
+ continue; // for duplicates, the best suggestion is actually added
+ }
+ Util.toIntsRef(suggestion, scratchInts);
+ builder.add(scratchInts, cost);
+ previous.copyBytes(suggestion);
+ }
+ fst = builder.finish();
+ success = true;
+ } finally {
+ if (success) {
+ IOUtils.close(reader, writer);
+ } else {
+ IOUtils.closeWhileHandlingException(reader, writer);
+ }
+
+ tempInput.delete();
+ tempSorted.delete();
+ }
+ }
+
+ @Override
+ public boolean store(File storeDir) throws IOException {
+ fst.save(new File(storeDir, FILENAME));
+ return true;
+ }
+
+ @Override
+ public boolean load(File storeDir) throws IOException {
+ this.fst = FST.read(new File(storeDir, FILENAME), PositiveIntOutputs.getSingleton(true));
+ return true;
+ }
+
+ @Override
+ public List lookup(String key, boolean onlyMorePopular, int num) {
+ assert num > 0;
+ BytesRef scratch = new BytesRef(key);
+ int prefixLength = scratch.length;
+ Arc arc = new Arc();
+
+ // match the prefix portion exactly
+ Long prefixOutput = null;
+ try {
+ prefixOutput = lookupPrefix(scratch, arc);
+ } catch (IOException bogus) { throw new RuntimeException(bogus); }
+
+ if (prefixOutput == null) {
+ return Collections.emptyList();
+ }
+
+ List results = new ArrayList(num);
+ if (exactFirst && arc.isFinal()) {
+ results.add(new LookupResult(scratch.utf8ToString(), decodeWeight(prefixOutput + arc.nextFinalOutput)));
+ if (--num == 0) {
+ return results; // that was quick
+ }
+ }
+
+ // complete top-N
+ MinResult completions[] = null;
+ try {
+ completions = Util.shortestPaths(fst, arc, num);
+ } catch (IOException bogus) { throw new RuntimeException(bogus); }
+
+ BytesRef suffix = new BytesRef(8);
+ for (MinResult completion : completions) {
+ scratch.length = prefixLength;
+ // append suffix
+ Util.toBytesRef(completion.input, suffix);
+ scratch.append(suffix);
+
+ results.add(new LookupResult(scratch.utf8ToString(), decodeWeight(prefixOutput + completion.output)));
+ }
+ return results;
+ }
+
+ private Long lookupPrefix(BytesRef scratch, Arc arc) throws /*Bogus*/IOException {
+ assert 0 == fst.outputs.getNoOutput().longValue();
+ long output = 0;
+ BytesReader bytesReader = fst.getBytesReader(0);
+
+ fst.getFirstArc(arc);
+
+ byte[] bytes = scratch.bytes;
+ int pos = scratch.offset;
+ int end = pos + scratch.length;
+ while (pos < end) {
+ if (fst.findTargetArc(bytes[pos++] & 0xff, arc, arc, bytesReader) == null) {
+ return null;
+ } else {
+ output += arc.output.longValue();
+ }
+ }
+
+ return output;
+ }
+
+ @Override
+ public boolean add(String key, Object value) {
+ return false; // Not supported.
+ }
+
+ /**
+ * Returns the weight associated with an input string,
+ * or null if it does not exist.
+ */
+ @Override
+ public Float get(String key) {
+ Arc arc = new Arc();
+ Long result = null;
+ try {
+ result = lookupPrefix(new BytesRef(key), arc);
+ } catch (IOException bogus) { throw new RuntimeException(bogus); }
+ if (result == null || !arc.isFinal()) {
+ return null;
+ } else {
+ return decodeWeight(result.intValue() + arc.nextFinalOutput.intValue());
+ }
+ }
+
+ /** cost -> weight */
+ private static float decodeWeight(long encoded) {
+ return Integer.MAX_VALUE - encoded;
+ }
+
+ /** weight -> cost */
+ private static long encodeWeight(float value) {
+ if (Float.isNaN(value) || Float.isInfinite(value) || value < 0 || value > Integer.MAX_VALUE) {
+ throw new UnsupportedOperationException("cannot encode value: " + value);
+ }
+ return Integer.MAX_VALUE - (int)value;
+ }
+}
Property changes on: modules/suggest/src/java/org/apache/lucene/search/suggest/fst/WFSTCompletionLookup.java
___________________________________________________________________
Added: svn:eol-style
## -0,0 +1 ##
+native
Index: lucene/core/src/test/org/apache/lucene/util/fst/TestFSTs.java
===================================================================
--- lucene/core/src/test/org/apache/lucene/util/fst/TestFSTs.java (revision 1290919)
+++ lucene/core/src/test/org/apache/lucene/util/fst/TestFSTs.java (working copy)
@@ -58,6 +58,7 @@
import org.apache.lucene.util.UnicodeUtil;
import org.apache.lucene.util._TestUtil;
import org.apache.lucene.util.fst.FST.Arc;
+import org.apache.lucene.util.fst.FST.BytesReader;
@UseNoMemoryExpensiveCodec
public class TestFSTs extends LuceneTestCase {
@@ -1975,6 +1976,119 @@
assertEquals(42, arc.output.longValue());
}
+ public void testShortestPaths() throws Exception {
+ final PositiveIntOutputs outputs = PositiveIntOutputs.getSingleton(true);
+ final Builder builder = new Builder(FST.INPUT_TYPE.BYTE1, outputs);
+
+ final IntsRef scratch = new IntsRef();
+ builder.add(Util.toIntsRef(new BytesRef("aab"), scratch), 22L);
+ builder.add(Util.toIntsRef(new BytesRef("aac"), scratch), 7L);
+ builder.add(Util.toIntsRef(new BytesRef("ax"), scratch), 17L);
+ final FST fst = builder.finish();
+ Writer w = new OutputStreamWriter(new FileOutputStream("out.dot"));
+ Util.toDot(fst, w, false, false);
+ w.close();
+
+ Util.MinResult[] r = Util.shortestPaths(fst,
+ fst.getFirstArc(new FST.Arc()),
+ 3);
+ assertEquals(3, r.length);
+
+ assertEquals(Util.toIntsRef(new BytesRef("aac"), scratch), r[0].input);
+ assertEquals(7, r[0].output);
+
+ assertEquals(Util.toIntsRef(new BytesRef("ax"), scratch), r[1].input);
+ assertEquals(17, r[1].output);
+
+ assertEquals(Util.toIntsRef(new BytesRef("aab"), scratch), r[2].input);
+ assertEquals(22, r[2].output);
+ }
+
+ public void testShortestPathsRandom() throws Exception {
+ int numWords = atLeast(1000);
+
+ final TreeMap slowCompletor = new TreeMap();
+ final TreeSet allPrefixes = new TreeSet();
+
+ final PositiveIntOutputs outputs = PositiveIntOutputs.getSingleton(true);
+ final Builder builder = new Builder(FST.INPUT_TYPE.BYTE1, outputs);
+ final IntsRef scratch = new IntsRef();
+
+ for (int i = 0; i < numWords; i++) {
+ String s;
+ while (true) {
+ s = _TestUtil.randomSimpleString(random);
+ if (!slowCompletor.containsKey(s)) {
+ break;
+ }
+ }
+
+ for (int j = 1; j < s.length(); j++) {
+ allPrefixes.add(s.substring(0, j));
+ }
+ int weight = _TestUtil.nextInt(random, 1, 100); // weights 1..100
+ slowCompletor.put(s, (long)weight);
+ }
+
+ for (Map.Entry e : slowCompletor.entrySet()) {
+ //System.out.println("add: " + e);
+ builder.add(Util.toIntsRef(new BytesRef(e.getKey()), scratch), e.getValue());
+ }
+
+ final FST fst = builder.finish();
+ //System.out.println("SAVE out.dot");
+ //Writer w = new OutputStreamWriter(new FileOutputStream("out.dot"));
+ //Util.toDot(fst, w, false, false);
+ //w.close();
+
+ BytesReader reader = fst.getBytesReader(0);
+
+ //System.out.println("testing: " + allPrefixes.size() + " prefixes");
+ for (String prefix : allPrefixes) {
+ // 1. run prefix against fst, then complete by value
+ //System.out.println("TEST: " + prefix);
+
+ long prefixOutput = 0;
+ FST.Arc arc = fst.getFirstArc(new FST.Arc());
+ for(int idx=0;idx matches = new ArrayList();
+
+ // TODO: could be faster... but its slowCompletor for a reason
+ for (Map.Entry e : slowCompletor.entrySet()) {
+ if (e.getKey().startsWith(prefix)) {
+ //System.out.println(" consider " + e.getKey());
+ matches.add(new Util.MinResult(Util.toIntsRef(new BytesRef(e.getKey().substring(prefix.length())), new IntsRef()),
+ e.getValue() - prefixOutput));
+ }
+ }
+
+ assertTrue(matches.size() > 0);
+ Collections.sort(matches);
+ if (matches.size() > topN) {
+ matches.subList(topN, matches.size()).clear();
+ }
+
+ assertEquals(matches.size(), r.length);
+
+ for(int hit=0;hit builder = new Builder(FST.INPUT_TYPE.BYTE1, outputs);
Index: lucene/core/src/java/org/apache/lucene/util/fst/Util.java
===================================================================
--- lucene/core/src/java/org/apache/lucene/util/fst/Util.java (revision 1290919)
+++ lucene/core/src/java/org/apache/lucene/util/fst/Util.java (working copy)
@@ -234,7 +234,302 @@
}
}
}
-
+
+ private static class FSTPath implements Comparable {
+ public FST.Arc arc;
+ public long cost;
+ public final IntsRef input = new IntsRef();
+
+ public FSTPath(long cost, FST.Arc arc) {
+ this.arc = new FST.Arc().copyFrom(arc);
+ this.cost = cost;
+ }
+
+ @Override
+ public String toString() {
+ return "input=" + input + " cost=" + cost;
+ }
+
+ @Override
+ public int compareTo(FSTPath other) {
+ if (cost < other.cost) {
+ return -1;
+ } else if (cost > other.cost) {
+ return 1;
+ } else {
+ return input.compareTo(other.input);
+ }
+ }
+ }
+
+ private static class TopNSearcher {
+
+ private final FST fst;
+ private final FST.Arc fromNode;
+ private final int topN;
+
+ // Set once the queue has filled:
+ FSTPath bottom = null;
+
+ TreeSet queue = null;
+
+ public TopNSearcher(FST fst, FST.Arc fromNode, int topN) {
+ this.fst = fst;
+ this.topN = topN;
+ this.fromNode = fromNode;
+ }
+
+ // If back plus this arc is competitive then add to queue:
+ private void addIfCompetitive(FSTPath path) {
+
+ assert queue != null;
+
+ long cost = path.cost + path.arc.output;
+ //System.out.println(" addIfCompetitive bottom=" + bottom + " queue.size()=" + queue.size());
+
+ if (bottom != null) {
+
+ if (cost > bottom.cost) {
+ // Doesn't compete
+ return;
+ } else if (cost == bottom.cost) {
+ // Tie break by alpha sort on the input:
+ path.input.grow(path.input.length+1);
+ path.input.ints[path.input.length++] = path.arc.label;
+ final int cmp = bottom.input.compareTo(path.input);
+ path.input.length--;
+ assert cmp != 0;
+ if (cmp < 0) {
+ // Doesn't compete
+ return;
+ }
+ }
+ // Competes
+ } else {
+ // Queue isn't full yet, so any path we hit competes:
+ }
+
+ final FSTPath newPath = new FSTPath(cost, path.arc);
+
+ newPath.input.grow(path.input.length+1);
+ System.arraycopy(path.input.ints, 0, newPath.input.ints, 0, path.input.length);
+ newPath.input.ints[path.input.length] = path.arc.label;
+ newPath.input.length = path.input.length+1;
+
+ //System.out.println(" add path=" + newPath);
+ queue.add(newPath);
+ if (bottom != null) {
+ final FSTPath removed = queue.pollLast();
+ assert removed == bottom;
+ bottom = queue.last();
+ //System.out.println(" now re-set bottom: " + bottom + " queue=" + queue);
+ } else if (queue.size() == topN) {
+ // Queue just filled up:
+ bottom = queue.last();
+ //System.out.println(" now set bottom: " + bottom);
+ }
+ }
+
+ public MinResult[] search() throws IOException {
+ //System.out.println(" search topN=" + topN);
+ final FST.Arc scratchArc = new FST.Arc();
+
+ final List results = new ArrayList();
+
+ final Long NO_OUTPUT = fst.outputs.getNoOutput();
+
+ // TODO: we could enable FST to sorting arcs by weight
+ // as it freezes... can easily do this on first pass
+ // (w/o requiring rewrite)
+
+ // TODO: maybe we should make an FST.INPUT_TYPE.BYTE0.5!?
+ // (nibbles)
+
+ // For each top N path:
+ while (results.size() < topN) {
+ //System.out.println("\nfind next path");
+
+ FSTPath path;
+
+ if (queue == null) {
+
+ if (results.size() != 0) {
+ // Ran out of paths
+ break;
+ }
+
+ // First pass (top path): start from original fromNode
+ if (topN > 1) {
+ queue = new TreeSet();
+ }
+
+ long minArcCost = Long.MAX_VALUE;
+ FST.Arc minArc = null;
+
+ path = new FSTPath(0, fromNode);
+ fst.readFirstTargetArc(fromNode, path.arc);
+
+ // Bootstrap: find the min starting arc
+ while (true) {
+ long arcScore = path.arc.output;
+ //System.out.println(" arc=" + (char) path.arc.label + " cost=" + arcScore);
+ if (arcScore < minArcCost) {
+ minArcCost = arcScore;
+ minArc = scratchArc.copyFrom(path.arc);
+ //System.out.println(" **");
+ }
+ if (queue != null) {
+ addIfCompetitive(path);
+ }
+ if (path.arc.isLast()) {
+ break;
+ }
+ fst.readNextArc(path.arc);
+ }
+
+ assert minArc != null;
+
+ if (queue != null) {
+ // Remove top path since we are now going to
+ // pursue it:
+ path = queue.pollFirst();
+ //System.out.println(" remove init path=" + path);
+ assert path.arc.label == minArc.label;
+ if (bottom != null && queue.size() == topN-1) {
+ bottom = queue.last();
+ //System.out.println(" set init bottom: " + bottom);
+ }
+ } else {
+ path.arc.copyFrom(minArc);
+ path.input.grow(1);
+ path.input.ints[0] = minArc.label;
+ path.input.length = 1;
+ path.cost = minArc.output;
+ }
+
+ } else {
+ path = queue.pollFirst();
+ if (path == null) {
+ // There were less than topN paths available:
+ break;
+ }
+ }
+
+ if (path.arc.label == FST.END_LABEL) {
+ //System.out.println(" empty string! cost=" + path.cost);
+ // Empty string!
+ path.input.length--;
+ results.add(new MinResult(path.input, path.cost));
+ continue;
+ }
+
+ if (results.size() == topN-1) {
+ // Last path -- don't bother w/ queue anymore:
+ queue = null;
+ }
+
+ //System.out.println(" path: " + path);
+
+ // We take path and find its "0 output completion",
+ // ie, just keep traversing the first arc with
+ // NO_OUTPUT that we can find, since this must lead
+ // to the minimum path that completes from
+ // path.arc.
+
+ // For each input letter:
+ while (true) {
+
+ //System.out.println("\n cycle path: " + path);
+
+ fst.readFirstTargetArc(path.arc, path.arc);
+
+ // For each arc leaving this node:
+ boolean foundZero = false;
+ while(true) {
+ //System.out.println(" arc=" + (char) path.arc.label + " cost=" + path.arc.output);
+ if (path.arc.output == NO_OUTPUT) {
+ if (queue == null) {
+ foundZero = true;
+ break;
+ } else if (!foundZero) {
+ scratchArc.copyFrom(path.arc);
+ foundZero = true;
+ } else {
+ addIfCompetitive(path);
+ }
+ } else if (queue != null) {
+ addIfCompetitive(path);
+ }
+ if (path.arc.isLast()) {
+ break;
+ }
+ fst.readNextArc(path.arc);
+ }
+
+ assert foundZero;
+
+ if (queue != null) {
+ // TODO: maybe we can save this copyFrom if we
+ // are more clever above... eg on finding the
+ // first NO_OUTPUT arc we'd switch to using
+ // scratchArc
+ path.arc.copyFrom(scratchArc);
+ }
+
+ if (path.arc.label == FST.END_LABEL) {
+ // Add final output:
+ //System.out.println(" done!: " + path);
+ results.add(new MinResult(path.input, path.cost + path.arc.output));
+ break;
+ } else {
+ path.input.grow(1+path.input.length);
+ path.input.ints[path.input.length] = path.arc.label;
+ path.input.length++;
+ path.cost += path.arc.output;
+ }
+ }
+ }
+
+ return results.toArray(new MinResult[results.size()]);
+ }
+ }
+
+ // TODO: parameterize the FST type and allow passing in a
+ // comparator; eg maybe your output is a PairOutput and
+ // one of the outputs in the pair is monotonic so you
+ // compare by that
+
+ public final static class MinResult implements Comparable {
+ public final IntsRef input;
+ public final long output;
+ public MinResult(IntsRef input, long output) {
+ this.input = input;
+ this.output = output;
+ }
+
+ @Override
+ public int compareTo(MinResult other) {
+ if (output < other.output) {
+ return -1;
+ } else if (output > other.output) {
+ return 1;
+ } else {
+ return input.compareTo(other.input);
+ }
+ }
+ }
+
+ /** Starting from node, find the top N min cost (Long
+ * output) completions to a final node.
+ *
+ * NOTE: you must share the outputs when you build the
+ * FST (pass doShare=true to {@link
+ * PositiveIntOutputs#getSingleton}). */
+
+ public static MinResult[] shortestPaths(FST fst, FST.Arc fromNode, int topN) throws IOException {
+ return new TopNSearcher(fst, fromNode, topN).search();
+ }
+
/**
* Dumps an {@link FST} to a GraphViz's dot language description
* for visualization. Example of use:
Index: lucene/CHANGES.txt
===================================================================
--- lucene/CHANGES.txt (revision 1290919)
+++ lucene/CHANGES.txt (working copy)
@@ -858,6 +858,9 @@
* LUCENE-3725: Added optional packing to FST building; this uses extra
RAM during building but results in a smaller FST. (Mike McCandless)
+
+* LUCENE-3714: Add top N shortest cost paths search for FST.
+ (Robert Muir, Dawid Weiss, Mike McCandless)
Bug fixes
Index: lucene/contrib/CHANGES.txt
===================================================================
--- lucene/contrib/CHANGES.txt (revision 1290919)
+++ lucene/contrib/CHANGES.txt (working copy)
@@ -150,6 +150,9 @@
* LUCENE-3602: Added query time joining under the join module. (Martijn van Groningen, Michael McCandless)
+ * LUCENE-3714: Add WFSTCompletionLookup suggester that supports more fine-grained
+ ranking for suggestions. (Mike McCandless, Dawid Weiss, Robert Muir)
+
API Changes
* LUCENE-3596: DirectoryTaxonomyWriter.openIndexWriter() now takes an