lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/jaspell/JaspellLookup.java

   1 package org.apache.lucene.search.suggest.jaspell;
   2
   3 import java.io.DataInputStream;
   4 import java.io.DataOutputStream;
   5 import java.io.File;
   6 import java.io.FileInputStream;
   7 import java.io.FileOutputStream;
   8 import java.io.IOException;
   9 import java.util.ArrayList;
  10 import java.util.List;
  11
  12 import org.apache.lucene.search.spell.SortedIterator;
  13 import org.apache.lucene.search.spell.TermFreqIterator;
  14 import org.apache.lucene.search.suggest.Lookup;
  15 import org.apache.lucene.search.suggest.UnsortedTermFreqIteratorWrapper;
  16 import org.apache.lucene.search.suggest.jaspell.JaspellTernarySearchTrie.TSTNode;
  17
  18 public class JaspellLookup extends Lookup {
  19   JaspellTernarySearchTrie trie = new JaspellTernarySearchTrie();
  20   private boolean usePrefix = true;
  21   private int editDistance = 2;
  22
  23   @Override
  24   public void build(TermFreqIterator tfit) throws IOException {
  25     if (tfit instanceof SortedIterator) {
  26       // make sure it's unsorted
  27       tfit = new UnsortedTermFreqIteratorWrapper(tfit);
  28     }
  29     trie = new JaspellTernarySearchTrie();
  30     trie.setMatchAlmostDiff(editDistance);
  31     while (tfit.hasNext()) {
  32       String key = tfit.next();
  33       float freq = tfit.freq();
  34       if (key.length() == 0) {
  35         continue;
  36       }
  37       trie.put(key, new Float(freq));
  38     }
  39   }
  40
  41   @Override
  42   public boolean add(String key, Object value) {
  43     trie.put(key, value);
  44     // XXX
  45     return false;
  46   }
  47
  48   @Override
  49   public Object get(String key) {
  50     return trie.get(key);
  51   }
  52
  53   @Override
  54   public List<LookupResult> lookup(String key, boolean onlyMorePopular, int num) {
  55     List<LookupResult> res = new ArrayList<LookupResult>();
  56     List<String> list;
  57     int count = onlyMorePopular ? num * 2 : num;
  58     if (usePrefix) {
  59       list = trie.matchPrefix(key, count);
  60     } else {
  61       list = trie.matchAlmost(key, count);
  62     }
  63     if (list == null || list.size() == 0) {
  64       return res;
  65
  66     }
  67     int maxCnt = Math.min(num, list.size());
  68     if (onlyMorePopular) {
  69       LookupPriorityQueue queue = new LookupPriorityQueue(num);
  70       for (String s : list) {
  71         float freq = (Float)trie.get(s);
  72         queue.insertWithOverflow(new LookupResult(s, freq));
  73       }
  74       for (LookupResult lr : queue.getResults()) {
  75         res.add(lr);
  76       }
  77     } else {
  78       for (int i = 0; i < maxCnt; i++) {
  79         String s = list.get(i);
  80         float freq = (Float)trie.get(s);
  81         res.add(new LookupResult(s, freq));
  82       }
  83     }
  84     return res;
  85   }
  86
  87   public static final String FILENAME = "jaspell.dat";
  88   private static final byte LO_KID = 0x01;
  89   private static final byte EQ_KID = 0x02;
  90   private static final byte HI_KID = 0x04;
  91   private static final byte HAS_VALUE = 0x08;
  92
  93
  94   @Override
  95   public boolean load(File storeDir) throws IOException {
  96     File data = new File(storeDir, FILENAME);
  97     if (!data.exists() || !data.canRead()) {
  98       return false;
  99     }
 100     DataInputStream in = new DataInputStream(new FileInputStream(data));
 101     TSTNode root = trie.new TSTNode('\0', null);
 102     try {
 103       readRecursively(in, root);
 104       trie.setRoot(root);
 105     } finally {
 106       in.close();
 107     }
 108     return true;
 109   }
 110
 111   private void readRecursively(DataInputStream in, TSTNode node) throws IOException {
 112     node.splitchar = in.readChar();
 113     byte mask = in.readByte();
 114     if ((mask & HAS_VALUE) != 0) {
 115       node.data = new Float(in.readFloat());
 116     }
 117     if ((mask & LO_KID) != 0) {
 118       TSTNode kid = trie.new TSTNode('\0', node);
 119       node.relatives[TSTNode.LOKID] = kid;
 120       readRecursively(in, kid);
 121     }
 122     if ((mask & EQ_KID) != 0) {
 123       TSTNode kid = trie.new TSTNode('\0', node);
 124       node.relatives[TSTNode.EQKID] = kid;
 125       readRecursively(in, kid);
 126     }
 127     if ((mask & HI_KID) != 0) {
 128       TSTNode kid = trie.new TSTNode('\0', node);
 129       node.relatives[TSTNode.HIKID] = kid;
 130       readRecursively(in, kid);
 131     }
 132   }
 133
 134   @Override
 135   public boolean store(File storeDir) throws IOException {
 136     if (!storeDir.exists() || !storeDir.isDirectory() || !storeDir.canWrite()) {
 137       return false;
 138     }
 139     TSTNode root = trie.getRoot();
 140     if (root == null) { // empty tree
 141       return false;
 142     }
 143     File data = new File(storeDir, FILENAME);
 144     DataOutputStream out = new DataOutputStream(new FileOutputStream(data));
 145     try {
 146       writeRecursively(out, root);
 147       out.flush();
 148     } finally {
 149       out.close();
 150     }
 151     return true;
 152   }
 153
 154   private void writeRecursively(DataOutputStream out, TSTNode node) throws IOException {
 155     if (node == null) {
 156       return;
 157     }
 158     out.writeChar(node.splitchar);
 159     byte mask = 0;
 160     if (node.relatives[TSTNode.LOKID] != null) mask |= LO_KID;
 161     if (node.relatives[TSTNode.EQKID] != null) mask |= EQ_KID;
 162     if (node.relatives[TSTNode.HIKID] != null) mask |= HI_KID;
 163     if (node.data != null) mask |= HAS_VALUE;
 164     out.writeByte(mask);
 165     if (node.data != null) {
 166       out.writeFloat((Float)node.data);
 167     }
 168     writeRecursively(out, node.relatives[TSTNode.LOKID]);
 169     writeRecursively(out, node.relatives[TSTNode.EQKID]);
 170     writeRecursively(out, node.relatives[TSTNode.HIKID]);
 171   }
 172 }