source: trunk/indexers/lucene-gs/src/org/greenstone/LuceneWrapper/GS2LuceneQuery.java@ 12775

/**
 *
 * @author  [email protected]
 * @author  [email protected]
 * @author  [email protected]
 * @author  [email protected]
 * @version 
 */

package org.nzdl.gsdl.LuceneWrap;


import java.io.*;
import java.util.*;
import java.util.regex.*;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermFreqVector;
import org.apache.lucene.queryParser.ParseException;
import org.apache.lucene.queryParser.QueryParser;
import org.apache.lucene.search.BooleanQuery.TooManyClauses;
import org.apache.lucene.search.Filter;
import org.apache.lucene.search.Hit;
import org.apache.lucene.search.Hits;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryFilter;
import org.apache.lucene.search.RangeFilter;
import org.apache.lucene.search.Searcher;
import org.apache.lucene.search.Sort;


public class GS2LuceneQuery
{
    // Use the standard set of English stop words by default
    static private String[] stop_words = StandardAnalyzer.STOP_WORDS;


    static public void main (String args[])
    {
    if (args.length == 0) {
        System.out.println("Usage: GS2LuceneQuery <index directory> [-fuzziness value] [-filter filter_string] [-sort sort_field] [-dco AND|OR] [-startresults number -endresults number]");
        return;
    }

        try {
        Searcher searcher = new IndexSearcher(args[0]);
        IndexReader reader = ((IndexSearcher) searcher).getIndexReader();

        // Create one query parser with the standard set of stop words, and one with none
        QueryParser query_parser = new QueryParser("TX", new StandardAnalyzer(stop_words));
        QueryParser query_parser_no_stop_words = new QueryParser("TX", new StandardAnalyzer(new String[] { }));

        Sort sorter = new Sort();
        Filter filter = null;
        String fuzziness = null;

        // Paging
        int start_results = 1;
        int end_results = -1;

            // New code to allow the default conjunction operator to be
            // definable
            String default_conjuction_operator = "OR";
            for (int i = 1; i < args.length; i++) {
        if (args[i].equals("-sort")) {
            i++;
            sorter = new Sort(args[i]);
        }
        if (args[i].equals("-filter")) {
            i++;
            
            // Parse up filter
            filter = parseFilterString(args[i]);
        }
        if (args[i].equals("-dco")) {
            i++;
            default_conjuction_operator = args[i];
        }
        if (args[i].equals("-fuzziness")) {
            i++;
            fuzziness = args[i];
        }
        if (args[i].equals("-startresults")) {
            i++;
            if (args[i].matches("\\d+")) {
            start_results = Integer.parseInt(args[i]);
            }
        }
        if (args[i].equals("-endresults")) {
            i++;
            if (args[i].matches("\\d+")) {
            end_results = Integer.parseInt(args[i]);
            }
        }
        }
        
        // Lucene does "OR" queries by default; do an "AND" query if specified
        if (default_conjuction_operator.equals("AND")) {
        query_parser.setDefaultOperator(query_parser.AND_OPERATOR);
        query_parser_no_stop_words.setDefaultOperator(query_parser.AND_OPERATOR);
        }

        BufferedReader in = new BufferedReader(new InputStreamReader(System.in, "UTF-8"));
        while (true) {
        // Read the query from STDIN
        String query_string = in.readLine();
        if (query_string == null || query_string.length() == -1) {
            break;
        }
        System.out.println("<ResultSet>");
        System.out.println("  <QueryString>" + xmlSafe(query_string) + "</QueryString>");
        if (filter != null) {
            System.out.println("  <FilterString>" + filter.toString() + "</FilterString>");
        }

        try {
            Query query_including_stop_words = query_parser_no_stop_words.parse(query_string);
            query_including_stop_words = query_including_stop_words.rewrite(reader);

            Query query = parseQuery(reader, query_parser, query_string, fuzziness);
            query = query.rewrite(reader);

            // Perform the query
            Hits hits;
            if (filter != null) {
            hits = searcher.search(query, filter, sorter);
            }
            else {
            hits = searcher.search(query, sorter);
            }

            // Return the list of expanded query terms and their frequencies
            HashMap term_counts = new HashMap();
            HashMap term_fields = new HashMap();
            HashSet terms = new HashSet();
            query.extractTerms(terms);
            Iterator iter = terms.iterator();
            while (iter.hasNext()) {
            Term term = (Term) iter.next();
            // If you wanted to limit this to just TX terms add
            // something like this:
            //if (term.field().equals("TX"))
            term_counts.put(term.text(), new Integer(0));
            term_fields.put(term.text(), term.field());
            }

            // Do we need to use a hit iterator to get sorted results?
            System.out.println("  <MatchingDocsInfo num=\"" + hits.length() + "\"/>");
            System.out.println("  <StartResults num=\"" + start_results + "\" />");
            System.out.println("  <EndsResults num=\"" + end_results + "\" />");

            int counter = 1;
            Iterator hit_iter = hits.iterator();
            while (hit_iter.hasNext()) {
            Hit hit = (Hit) hit_iter.next();
            Document doc = hit.getDocument();
            String node_id = doc.get("nodeID");
            
            // May not be paging results
            if (start_results == 1 && end_results == -1) {
                System.out.println("  <Match id=\"" + node_id + "\" />");
            }
            // Otherwise skip up until page offset
            else if (start_results <= counter && counter <= end_results) {
                System.out.println("  <Match id=\"" + node_id + "\" />");
            }
            // And skip all the rest
            
            // From the document, extract the Term Vector for the
            // TX field
            TermFreqVector term_freq_vector = reader.getTermFreqVector(hit.getId(), "TX");
            if (term_freq_vector != null && term_freq_vector.size() > 0) {
                int[] term_frequencies = term_freq_vector.getTermFrequencies();
                // Now for each query term, determine the
                // frequency - which may of course be 0.
                Set term_counts_set = term_counts.keySet();
                Iterator terms_iter = term_counts_set.iterator();
                while (terms_iter.hasNext()) {
                
                String term = (String) terms_iter.next();
                Integer count_integer = (Integer) term_counts.get(term);
                int count = count_integer.intValue();
                int index = term_freq_vector.indexOf(term);
                // If the term has a count, then add to
                // the total count for this term
                if (index != -1) {
                    count += term_frequencies[index];
                }
                // Store the result
                term_counts.put(term, new Integer(count));
                count_integer = null;
                term = null;
                }
                terms_iter = null;
                term_counts_set = null;
            }
            else {
                ///ystem.err.println("Error! Missing term vector for document " + hit.getId());
            }
            ++counter;
            }
            
            // Retrieve all the useful terms
            Set term_counts_set = term_counts.keySet();
            System.out.println("  <QueryTermsInfo num=\"" + term_counts_set.size() + "\"/>");
            // Iterate over them
            Iterator terms_iter = term_counts_set.iterator();
            while (terms_iter.hasNext()) {
            String term = (String) terms_iter.next();
            Integer count = (Integer) term_counts.get(term);
            String field = (String) term_fields.get(term);

            // Ignore any terms with zero frequency, because they don't exist in the matching
            // documents. It seems that this should never happen, but it's a consequence of
            // how the terms are identified. The terms are found by rewriting the query (above).
            // At this point, the query hasn't been run, so each query term is expanded without
            // knowing whether the expanded term will actually appear in one of the resulting
            // documents. For example, "otago" may be expanded to "otaio" and "otaqo", but if
            // the search is for "otago AND auckland", no matching documents may include "otaio".
            // Hopefully that made some sense...
            if (count.intValue() > 0) {
                System.out.println("  <Term value=\"" + term + "\" field=\"" + field + "\" freq=\"" + count.intValue() + "\" />");
            }
            count = null;
            term = null;
            }

            // Cleanup
            terms_iter = null;
            term_counts_set = null;

            // Return the list of stop words removed from the query
            HashSet terms_including_stop_words = new HashSet();
            query_including_stop_words.extractTerms(terms_including_stop_words);
            Iterator terms_including_stop_words_iter = terms_including_stop_words.iterator();
            while (terms_including_stop_words_iter.hasNext()) {
            Term term = (Term) terms_including_stop_words_iter.next();
            if (!terms.contains(term)) {
                System.out.println("  <StopWord value=\"" + term.text() + "\"/>");
            }
            }
        }
        catch (ParseException parse_exception) {
            System.out.println("  <Error type=\"PARSE_EXCEPTION\"/>");
        }
        catch (TooManyClauses too_many_clauses_exception) {
            System.out.println("  <Error type=\"TOO_MANY_CLAUSES\"/>");
        }

        System.out.println("</ResultSet>");
        }
        
        searcher.close();
    }
    catch (IOException exception) {
        exception.printStackTrace();
    }
    }

    private static String xmlSafe(String text) {
        return text.replaceAll("\\&", "\\&amp;");
    }

    private static Query parseQuery(IndexReader reader, QueryParser query_parser, String query_string, String fuzziness)
    throws java.io.IOException, org.apache.lucene.queryParser.ParseException
    {
    // Split query string into the search terms and the filter terms
    // * The first +(...) term contains the search terms so count
    //   up '(' and stop when we finish matching ')'
    int offset = 0;
    int paren_count = 0;
    boolean seen_paren = false;
    while (offset < query_string.length() && (!seen_paren || paren_count > 0)) {
        if (query_string.charAt(offset) == '(') {
        paren_count++;
        seen_paren = true;
        }
        if (query_string.charAt(offset) == ')') {
        paren_count--;
        }
        offset++;
    }
    String query_prefix = query_string.substring(0, offset);
    String query_suffix = query_string.substring(offset);

    ///ystem.err.println("Prefix: " + query_prefix);
    ///ystem.err.println("Suffix: " + query_suffix);

    Query query = query_parser.parse(query_prefix);
    query = query.rewrite(reader);

    // If this is a fuzzy search, then we need to add the fuzzy
    // flag to each of the query terms
    if (fuzziness != null && query.toString().length() > 0) {
        // Revert the query to a string
        System.err.println("Rewritten query: " + query.toString());
        // Search through the string for TX:<term> query terms
        // and append the ~ operator. Not that this search will
        // not change phrase searches (TX:"<term> <term>") as
        // fuzzy searching is not possible for these entries.
        // Yahoo! Time for a state machine!
        StringBuffer mutable_query_string = new StringBuffer(query.toString());
        int o = 0; // Offset
        // 0 = BASE, 1 = SEEN_T, 2 = SEEN_TX, 3 = SEEN_TX:
        int s = 0; // State
        while (o < mutable_query_string.length()) {
        char c = mutable_query_string.charAt(o);
        if (s == 0 && c == 'T') {
            ///ystem.err.println("Found T!");
            s = 1;
        }
        else if (s == 1) {
            if (c == 'X') {
            ///ystem.err.println("Found X!");
            s = 2;
            }
            else {
            s = 0; // Reset
            }
        }
        else if (s == 2) {
            if (c == ':') {
            ///ystem.err.println("Found TX:!");
            s = 3;
            }
            else {
            s = 0; // Reset
            }
        }
        else if (s == 3) {
            // Don't process phrases
            if (c == '"') {
            ///ystem.err.println("Stupid phrase...");
            s = 0; // Reset
            }
            // Found the end of the term... add the
            // fuzzy search indicator
            // Nor outside the scope of parentheses
            else if (Character.isWhitespace(c) || c == ')') {
            ///ystem.err.println("Yahoo! Found fuzzy term.");
            mutable_query_string.insert(o, '~' + fuzziness);
            o++;
            s = 0; // Reset
            }
        }
        o++;
        }
        // If we were in the state of looking for the end of a
        // term - then we just found it!
        if (s == 3) {
        mutable_query_string.append('~' + fuzziness);
        }
        // Reparse the query
        ///ystem.err.println("Fuzzy query: " + mutable_query_string.toString() + query_suffix);
        query = query_parser.parse(mutable_query_string.toString() + query_suffix);
    }
    else {
        query = query_parser.parse(query_prefix + query_suffix);
    }
    
    return query;
    }


  /**
   *  @todo Michael to comment
   */
  private static Filter parseFilterString(String filter_string)
  {
    Filter result = null;
    Pattern pattern = Pattern.compile("\\s*\\+(\\w+)\\:([\\{\\[])(\\d+)\\s+TO\\s+(\\d+)([\\}\\]])\\s*");
    Matcher matcher = pattern.matcher(filter_string);
    if (matcher.matches()) {
        String field_name = matcher.group(1);
        boolean include_lower = matcher.group(2).equals("[");
        String lower_term = matcher.group(3);
        String upper_term = matcher.group(4);
        boolean include_upper = matcher.group(5).equals("]");
        result = new RangeFilter(field_name, lower_term, upper_term, include_lower, include_upper);
    }
    else {
        System.err.println("Error: Could not understand filter string \"" + filter_string + "\"");
    }
    return result;
  }
  /** parseFilterString() **/
}