+++ /dev/null
-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);
- }
- }
-}
fnp / pylucene.git / blobdiff