pylucene 3.5.0-3

[pylucene.git] / lucene-java-3.5.0 / lucene / contrib / spellchecker / src / java / org / apache / lucene / search / suggest / tst / TSTLookup.java

diff --git a/lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTLookup.java b/lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTLookup.java
new file mode 100644 (file)
index 0000000..497e05f
--- /dev/null
+++ b/lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTLookup.java
@@ -0,0 +1,174 @@
+package org.apache.lucene.search.suggest.tst;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.lucene.search.suggest.Lookup;
+import org.apache.lucene.search.suggest.SortedTermFreqIteratorWrapper;
+import org.apache.lucene.search.spell.SortedIterator;
+import org.apache.lucene.search.spell.TermFreqIterator;
+
+public class TSTLookup extends Lookup {
+  TernaryTreeNode root = new TernaryTreeNode();
+  TSTAutocomplete autocomplete = new TSTAutocomplete();
+
+  @Override
+  public void build(TermFreqIterator tfit) throws IOException {
+    root = new TernaryTreeNode();
+    // buffer first
+    if (!(tfit instanceof SortedIterator)) {
+      // make sure it's sorted
+      tfit = new SortedTermFreqIteratorWrapper(tfit);
+    }
+
+    ArrayList<String> tokens = new ArrayList<String>();
+    ArrayList<Float> vals = new ArrayList<Float>();
+    while (tfit.hasNext()) {
+      tokens.add(tfit.next());
+      vals.add(new Float(tfit.freq()));
+    }
+    autocomplete.balancedTree(tokens.toArray(), vals.toArray(), 0, tokens.size() - 1, root);
+  }
+
+  @Override
+  public boolean add(String key, Object value) {
+    autocomplete.insert(root, key, value, 0);
+    // XXX we don't know if a new node was created
+    return true;
+  }
+
+  @Override
+  public Object get(String key) {
+    List<TernaryTreeNode> list = autocomplete.prefixCompletion(root, key, 0);
+    if (list == null || list.isEmpty()) {
+      return null;
+    }
+    for (TernaryTreeNode n : list) {
+      if (n.token.equals(key)) {
+        return n.val;
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public List<LookupResult> lookup(String key, boolean onlyMorePopular, int num) {
+    List<TernaryTreeNode> list = autocomplete.prefixCompletion(root, key, 0);
+    List<LookupResult> res = new ArrayList<LookupResult>();
+    if (list == null || list.size() == 0) {
+      return res;
+    }
+    int maxCnt = Math.min(num, list.size());
+    if (onlyMorePopular) {
+      LookupPriorityQueue queue = new LookupPriorityQueue(num);
+      for (TernaryTreeNode ttn : list) {
+        queue.insertWithOverflow(new LookupResult(ttn.token, (Float)ttn.val));
+      }
+      for (LookupResult lr : queue.getResults()) {
+        res.add(lr);
+      }
+    } else {
+      for (int i = 0; i < maxCnt; i++) {
+        TernaryTreeNode ttn = list.get(i);
+        res.add(new LookupResult(ttn.token, (Float)ttn.val));
+      }
+    }
+    return res;
+  }
+  
+  public static final String FILENAME = "tst.dat";
+  
+  private static final byte LO_KID = 0x01;
+  private static final byte EQ_KID = 0x02;
+  private static final byte HI_KID = 0x04;
+  private static final byte HAS_TOKEN = 0x08;
+  private static final byte HAS_VALUE = 0x10;
+
+  @Override
+  public synchronized boolean load(File storeDir) throws IOException {
+    File data = new File(storeDir, FILENAME);
+    if (!data.exists() || !data.canRead()) {
+      return false;
+    }
+    DataInputStream in = new DataInputStream(new FileInputStream(data));
+    root = new TernaryTreeNode();
+    try {
+      readRecursively(in, root);
+    } finally {
+      in.close();
+    }
+    return true;
+  }
+  
+  // pre-order traversal
+  private void readRecursively(DataInputStream in, TernaryTreeNode node) throws IOException {
+    node.splitchar = in.readChar();
+    byte mask = in.readByte();
+    if ((mask & HAS_TOKEN) != 0) {
+      node.token = in.readUTF();
+    }
+    if ((mask & HAS_VALUE) != 0) {
+      node.val = new Float(in.readFloat());
+    }
+    if ((mask & LO_KID) != 0) {
+      node.loKid = new TernaryTreeNode();
+      readRecursively(in, node.loKid);
+    }
+    if ((mask & EQ_KID) != 0) {
+      node.eqKid = new TernaryTreeNode();
+      readRecursively(in, node.eqKid);
+    }
+    if ((mask & HI_KID) != 0) {
+      node.hiKid = new TernaryTreeNode();
+      readRecursively(in, node.hiKid);
+    }
+  }
+
+  @Override
+  public synchronized boolean store(File storeDir) throws IOException {
+    if (!storeDir.exists() || !storeDir.isDirectory() || !storeDir.canWrite()) {
+      return false;
+    }
+    File data = new File(storeDir, FILENAME);
+    DataOutputStream out = new DataOutputStream(new FileOutputStream(data));
+    try {
+      writeRecursively(out, root);
+      out.flush();
+    } finally {
+      out.close();
+    }
+    return true;
+  }
+  
+  // pre-order traversal
+  private void writeRecursively(DataOutputStream out, TernaryTreeNode node) throws IOException {
+    // write out the current node
+    out.writeChar(node.splitchar);
+    // prepare a mask of kids
+    byte mask = 0;
+    if (node.eqKid != null) mask |= EQ_KID;
+    if (node.loKid != null) mask |= LO_KID;
+    if (node.hiKid != null) mask |= HI_KID;
+    if (node.token != null) mask |= HAS_TOKEN;
+    if (node.val != null) mask |= HAS_VALUE;
+    out.writeByte(mask);
+    if (node.token != null) out.writeUTF(node.token);
+    if (node.val != null) out.writeFloat((Float)node.val);
+    // recurse and write kids
+    if (node.loKid != null) {
+      writeRecursively(out, node.loKid);
+    }
+    if (node.eqKid != null) {
+      writeRecursively(out, node.eqKid);
+    }
+    if (node.hiKid != null) {
+      writeRecursively(out, node.hiKid);
+    }
+  }
+}