StructuredQueryRanker.java

package ivory.sqe.retrieval;

import ivory.core.ConfigurationException;
import ivory.core.RetrievalEnvironment;
import ivory.smrf.model.GlobalEvidence;
import ivory.smrf.retrieval.Accumulator;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

import org.apache.hadoop.fs.FileSystem;

import com.google.gson.JsonObject;

import edu.umd.cloud9.collection.DocnoMapping;

public class StructuredQueryRanker {
  private RetrievalEnvironment env;
  private final PriorityQueue<Accumulator> sortedAccumulators = new PriorityQueue<Accumulator>();
  private final int numResults;
  private HashMap<String, Accumulator[]> results;

  public StructuredQueryRanker(String indexPath, FileSystem fs, int numResults) throws IOException,
  ConfigurationException {
    this.env = new RetrievalEnvironment(indexPath, fs);
    env.initialize(true);

    this.numResults = numResults;
    results = new HashMap<String, Accumulator[]>();
  }

  public Accumulator[] rank(String qid, JsonObject query, int queryLength) {
    GlobalEvidence globalEvidence = new GlobalEvidence(env.getDocumentCount(), env.getCollectionSize(), queryLength);

    PostingsReaderWrapper structureReader;
    structureReader = new PostingsReaderWrapper(query, env, globalEvidence);

    sortedAccumulators.clear();
    Accumulator a = new Accumulator(0, 0.0f);

    // NodeWeight that must be achieved to enter result set.
    double scoreThreshold = Double.NEGATIVE_INFINITY;

    int docno = Integer.MAX_VALUE;
    int nextDocno = structureReader.getNextCandidate(docno);
    if(nextDocno < docno){
      docno = nextDocno;
    }
    while (docno < Integer.MAX_VALUE) {
      float score = 0.0f;

      // Document-at-a-time scoring.
      NodeWeight sc = structureReader.computeScore(docno);
      score = sc.getScore();

      // Keep track of numResults best accumulators.
      if (score > scoreThreshold) {
        a.docno = docno;
        a.score = score;
        sortedAccumulators.add(a);

        if (sortedAccumulators.size() == numResults + 1) {
          a = sortedAccumulators.poll();
          scoreThreshold = sortedAccumulators.peek().score;
        } else {
          a = new Accumulator(0, 0.0f);
        }
      }     

      // Advance to next document
      docno = Integer.MAX_VALUE;
      nextDocno = structureReader.getNextCandidate(docno);
      if(nextDocno < docno){
        docno = nextDocno;
      }
    }

    // Grab the accumulators off the stack, in (reverse) order.
    Accumulator[] accs = new Accumulator[Math.min(numResults, sortedAccumulators.size())];
    for (int i = 0; i < accs.length; i++) {
      Accumulator acc = sortedAccumulators.poll();
      accs[accs.length - 1 - i] = acc;
    }

    this.results.put(qid, accs);

    return accs;
  }

  public DocnoMapping getDocnoMapping() throws IOException {
    return env.getDocnoMapping();
  }

  public Accumulator[] getResults(String queryID) {
    return results.get(queryID);
  }

  public Map<String, Accumulator[]> getResults() {
    return results;
  }
}