Lucene example source code file (JaspellLookup.java)

This example Lucene source code file (JaspellLookup.java) is included in the DevDaily.com "Java Source Code Warehouse" project. The intent of this project is to help you "Learn Java by Example" ^TM.

Java - Lucene tags/keywords

eq_kid, file, file, filename, hi_kid, io, ioexception, ioexception, jaspellternarysearchtrie, lo_kid, override, override, string, tstnode, tstnode, util

The Lucene JaspellLookup.java source code

package org.apache.lucene.search.suggest.jaspell;

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.spell.SortedIterator;
import org.apache.lucene.search.spell.TermFreqIterator;
import org.apache.lucene.search.suggest.Lookup;
import org.apache.lucene.search.suggest.UnsortedTermFreqIteratorWrapper;
import org.apache.lucene.search.suggest.jaspell.JaspellTernarySearchTrie.TSTNode;

public class JaspellLookup extends Lookup {
  JaspellTernarySearchTrie trie = new JaspellTernarySearchTrie();
  private boolean usePrefix = true;
  private int editDistance = 2;

  @Override
  public void build(TermFreqIterator tfit) throws IOException {
    if (tfit instanceof SortedIterator) {
      // make sure it's unsorted
      tfit = new UnsortedTermFreqIteratorWrapper(tfit);
    }
    trie = new JaspellTernarySearchTrie();
    trie.setMatchAlmostDiff(editDistance);
    while (tfit.hasNext()) {
      String key = tfit.next();
      float freq = tfit.freq();
      if (key.length() == 0) {
        continue;
      }
      trie.put(key, new Float(freq));
    }
  }

  @Override
  public boolean add(String key, Object value) {
    trie.put(key, value);
    // XXX
    return false;
  }

  @Override
  public Object get(String key) {
    return trie.get(key);
  }

  @Override
  public List<LookupResult> lookup(String key, boolean onlyMorePopular, int num) {
    List<LookupResult> res = new ArrayList();
    List<String> list;
    int count = onlyMorePopular ? num * 2 : num;
    if (usePrefix) {
      list = trie.matchPrefix(key, count);
    } else {
      list = trie.matchAlmost(key, count);
    }
    if (list == null || list.size() == 0) {
      return res;
      
    }
    int maxCnt = Math.min(num, list.size());
    if (onlyMorePopular) {
      LookupPriorityQueue queue = new LookupPriorityQueue(num);
      for (String s : list) {
        float freq = (Float)trie.get(s);
        queue.insertWithOverflow(new LookupResult(s, freq));
      }
      for (LookupResult lr : queue.getResults()) {
        res.add(lr);
      }
    } else {
      for (int i = 0; i < maxCnt; i++) {
        String s = list.get(i);
        float freq = (Float)trie.get(s);
        res.add(new LookupResult(s, freq));
      }      
    }
    return res;
  }

  public static final String FILENAME = "jaspell.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_VALUE = 0x08;
 
  
  @Override
  public 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));
    TSTNode root = trie.new TSTNode('\0', null);
    try {
      readRecursively(in, root);
      trie.setRoot(root);
    } finally {
      in.close();
    }
    return true;
  }
  
  private void readRecursively(DataInputStream in, TSTNode node) throws IOException {
    node.splitchar = in.readChar();
    byte mask = in.readByte();
    if ((mask & HAS_VALUE) != 0) {
      node.data = new Float(in.readFloat());
    }
    if ((mask & LO_KID) != 0) {
      TSTNode kid = trie.new TSTNode('\0', node);
      node.relatives[TSTNode.LOKID] = kid;
      readRecursively(in, kid);
    }
    if ((mask & EQ_KID) != 0) {
      TSTNode kid = trie.new TSTNode('\0', node);
      node.relatives[TSTNode.EQKID] = kid;
      readRecursively(in, kid);
    }
    if ((mask & HI_KID) != 0) {
      TSTNode kid = trie.new TSTNode('\0', node);
      node.relatives[TSTNode.HIKID] = kid;
      readRecursively(in, kid);
    }
  }

  @Override
  public boolean store(File storeDir) throws IOException {
    if (!storeDir.exists() || !storeDir.isDirectory() || !storeDir.canWrite()) {
      return false;
    }
    TSTNode root = trie.getRoot();
    if (root == null) { // empty tree
      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;
  }
  
  private void writeRecursively(DataOutputStream out, TSTNode node) throws IOException {
    if (node == null) {
      return;
    }
    out.writeChar(node.splitchar);
    byte mask = 0;
    if (node.relatives[TSTNode.LOKID] != null) mask |= LO_KID;
    if (node.relatives[TSTNode.EQKID] != null) mask |= EQ_KID;
    if (node.relatives[TSTNode.HIKID] != null) mask |= HI_KID;
    if (node.data != null) mask |= HAS_VALUE;
    out.writeByte(mask);
    if (node.data != null) {
      out.writeFloat((Float)node.data);
    }
    writeRecursively(out, node.relatives[TSTNode.LOKID]);
    writeRecursively(out, node.relatives[TSTNode.EQKID]);
    writeRecursively(out, node.relatives[TSTNode.HIKID]);
  }
}