FederatedQueryRunner.java

  1. /*
  2.  *  FederatedQueryRunner.java
  3.  *
  4.  *  Copyright (c) 2007-2011, The University of Sheffield.
  5.  *
  6.  *  This file is part of GATE Mímir (see http://gate.ac.uk/family/mimir.html),
  7.  *  and is free software, licenced under the GNU Lesser General Public License,
  8.  *  Version 3, June 2007 (also included with this distribution as file
  9.  *  LICENCE-LGPL3.html).
  10.  *
  11.  *  Ian Roberts, 15 Dec 2011
  12.  * 
  13.  *  $Id: FederatedQueryRunner.java 16414 2012-12-10 10:54:47Z valyt $
  14.  */
  15. package gate.mimir.search;

  17. import gate.mimir.index.IndexException;
  18. import gate.mimir.search.query.Binding;

  20. import it.unimi.dsi.fastutil.ints.IntArrayList;
  21. import it.unimi.dsi.fastutil.ints.IntBigArrayBigList;
  22. import it.unimi.dsi.fastutil.ints.IntBigList;
  23. import it.unimi.dsi.fastutil.ints.IntList;
  24. import it.unimi.dsi.fastutil.longs.LongBigArrayBigList;
  25. import it.unimi.dsi.fastutil.longs.LongBigList;

  27. import java.io.IOException;
  28. import java.io.Serializable;
  29. import java.util.List;
  30. import java.util.Map;
  31. import java.util.Set;

  33. import org.slf4j.Logger;
  34. import org.slf4j.LoggerFactory;

  36. /**
  37.  * A {@link QueryRunner} that presents a set of sub-indexes (represented by
  38.  * their own QueryRunners) as a single index.
  39.  */
  40. public class FederatedQueryRunner implements QueryRunner {
  41.   
  42.   private static final Logger log = LoggerFactory.getLogger(FederatedQueryRunner.class);
  43.   
  44.   /**
  45.    * The total number of result documents (or -1 if not yet known).
  46.    */
  47.   private long documentsCount = -1;
  48.   
  49.   /**
  50.    * The query runners for the sub-indexes.
  51.    */
  52.   protected QueryRunner[] subRunners;
  53.   
  54.   /**
  55.    * The next rank that needs to be merged from each sub runner.
  56.    */
  57.   protected long[] nextSubRunnerRank;
  58.   
  59.   /**
  60.    * For each result document rank, this list supplies the index for the
  61.    * sub-runner that supplied the document.
  62.    */
  63.   protected IntBigList rank2runnerIndex;
  64.   
  65.   /**
  66.    * Which of the sub-runners has provided the previous document. This is an
  67.    * instance field so that we can rotate the sub-runners (when the scores are
  68.    * equal)
  69.    */
  70.   private int bestSubRunnerIndex = -1;
  71.   
  72.   /**
  73.    * For each result document rank, this list supplies the rank of the document
  74.    * in sub-runner that supplied it.
  75.    */
  76.   protected LongBigList rank2subRank;
  77.   
  78.   public FederatedQueryRunner(QueryRunner[] subrunners) {
  79.     this.subRunners = subrunners;
  80.     this.nextSubRunnerRank = null;
  81.     this.rank2runnerIndex = new IntBigArrayBigList();
  82.     this.rank2subRank = new LongBigArrayBigList();
  83.   }

  85.   /* (non-Javadoc)
  86.    * @see gate.mimir.search.QueryRunner#getDocumentsCount()
  87.    */
  88.   @Override
  89.   public long getDocumentsCount() {
  90.     if(documentsCount < 0) {
  91.       long newDocumentsCount = 0;
  92.       for(QueryRunner subRunner : subRunners) {
  93.         long subDocumentsCount = subRunner.getDocumentsCount();
  94.         if(subDocumentsCount < 0) {
  95.           return -1;
  96.         } else {
  97.           newDocumentsCount += subDocumentsCount;
  98.         }
  99.       }
  100.       synchronized(this) {
  101.         // initialize the nextSubRunnerRank array
  102.         nextSubRunnerRank = new long[subRunners.length];
  103.         for(int i = 0; i < nextSubRunnerRank.length; i++) {
  104.           if(subRunners[i].getDocumentsCount() == 0) {
  105.             nextSubRunnerRank[i] = -1;
  106.           }
  107.         }
  108.         documentsCount = newDocumentsCount;
  109.       }
  110.     }
  111.     return documentsCount;
  112.   }

  114.   
  115.   
  116.   /* (non-Javadoc)
  117.    * @see gate.mimir.search.QueryRunner#getDocumentsCountSync()
  118.    */
  119.   @Override
  120.   public long getDocumentsCountSync() {
  121.     for(QueryRunner subRunner : subRunners) {
  122.       subRunner.getDocumentsCountSync();
  123.     }
  124.     return getDocumentsCount();
  125.   }

  127.   /* (non-Javadoc)
  128.    * @see gate.mimir.search.QueryRunner#getCurrentDocumentsCount()
  129.    */
  130.   @Override
  131.   public long getDocumentsCurrentCount() {
  132.     if(documentsCount >= 0) {
  133.       return documentsCount;
  134.     } else {
  135.       int newDocumentsCount = 0;
  136.       for(QueryRunner subRunner : subRunners) {
  137.         newDocumentsCount += subRunner.getDocumentsCurrentCount();
  138.       }
  139.       return newDocumentsCount;
  140.     }
  141.   }
  142.   
  143.   /**
  144.    * Ensure that the given rank is resolved to the appropriate sub-runner rank.
  145.    * @throws IndexOutOfBoundsException if rank is beyond the last document.
  146.    */
  147.   private final synchronized void checkRank(long rank) throws IndexOutOfBoundsException, IOException {
  148.     if(rank < 0) throw new IndexOutOfBoundsException(
  149.       "Document rank " + rank + " is negative.");
  150.     long maxRank = getDocumentsCount(); 
  151.     if(rank >= maxRank) throw new IndexOutOfBoundsException(
  152.         "Document rank too large (" + rank + " >= " + maxRank + ").");    
  153.     // quick check to see if we need to do anything else
  154.     if(rank < rank2runnerIndex.size64()) {
  155.       return;
  156.     }
  157.     for(long nextRank = rank2runnerIndex.size64(); nextRank <= rank; nextRank++) {
  158.       boolean allOut = true;
  159.       // start with the runner next the previously chosen one
  160.       bestSubRunnerIndex = (bestSubRunnerIndex + 1) % subRunners.length;
  161.       double maxScore = Double.NEGATIVE_INFINITY;
  162.       if(nextSubRunnerRank[bestSubRunnerIndex] >= 0) {
  163.         maxScore = subRunners[bestSubRunnerIndex].getDocumentScore(
  164.             nextSubRunnerRank[bestSubRunnerIndex]);
  165.         allOut = false;
  166.       }
  167.       // now check all remaining runners, in round-robin fashion
  168.       final int from = bestSubRunnerIndex + 1;
  169.       final int to = bestSubRunnerIndex + subRunners.length;
  170.       for(int bigI = from; bigI < to; bigI++) {
  171.         int i = bigI % subRunners.length;
  172.         if(nextSubRunnerRank[i] >= 0) {
  173.           allOut = false;
  174.           if(subRunners[i].getDocumentScore(nextSubRunnerRank[i]) > maxScore) {
  175.             bestSubRunnerIndex = i;
  176.             maxScore = subRunners[i].getDocumentScore(nextSubRunnerRank[i]);
  177.           }
  178.         }
  179.       }
  180.       if(allOut) {
  181.         // we ran out of docs
  182.         throw new IndexOutOfBoundsException("Requested rank was " + rank +
  183.           " but ran out of documents at " + nextRank + "!");
  184.       }
  185.       // consume the next doc from bestSubRunnerIndex
  186.       rank2runnerIndex.add(bestSubRunnerIndex);
  187.       rank2subRank.add(nextSubRunnerRank[bestSubRunnerIndex]);
  188.       if(nextSubRunnerRank[bestSubRunnerIndex] <
  189.           subRunners[bestSubRunnerIndex].getDocumentsCount() -1) {
  190.         nextSubRunnerRank[bestSubRunnerIndex]++;
  191.       } else {
  192.         // this runner has run out of documents
  193.         nextSubRunnerRank[bestSubRunnerIndex] = -1;
  194.       }
  195.     }
  196.   }
  197.   
  198.   /* (non-Javadoc)
  199.    * @see gate.mimir.search.QueryRunner#getDocumentID(int)
  200.    */
  201.   @Override
  202.   public long getDocumentID(long rank) throws IndexOutOfBoundsException,
  203.     IOException {
  204.     checkRank(rank);
  205.     long subId = subRunners[rank2runnerIndex.getInt(rank)].getDocumentID(
  206.       rank2subRank.getLong(rank));
  207.     return subId * subRunners.length + rank2runnerIndex.getInt(rank);
  208.   }

  210.   /* (non-Javadoc)
  211.    * @see gate.mimir.search.QueryRunner#getDocumentScore(int)
  212.    */
  213.   @Override
  214.   public double getDocumentScore(long rank) throws IndexOutOfBoundsException,
  215.     IOException {
  216.     checkRank(rank);
  217.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentScore(
  218.         rank2subRank.getLong(rank));
  219.   }

  221.   /* (non-Javadoc)
  222.    * @see gate.mimir.search.QueryRunner#getDocumentHits(int)
  223.    */
  224.   @Override
  225.   public List<Binding> getDocumentHits(long rank)
  226.     throws IndexOutOfBoundsException, IOException {
  227.     checkRank(rank);
  228.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentHits(
  229.         rank2subRank.getLong(rank));
  230.   }

  232.   /* (non-Javadoc)
  233.    * @see gate.mimir.search.QueryRunner#getDocumentText(int, int, int)
  234.    */
  235.   @Override
  236.   public String[][] getDocumentText(long rank, int termPosition, int length)
  237.     throws IndexException, IndexOutOfBoundsException, IOException {
  238.     checkRank(rank);
  239.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentText(
  240.         rank2subRank.getLong(rank), termPosition, length);
  241.   }

  243.   /* (non-Javadoc)
  244.    * @see gate.mimir.search.QueryRunner#getDocumentURI(int)
  245.    */
  246.   @Override
  247.   public String getDocumentURI(long rank) throws IndexException,
  248.     IndexOutOfBoundsException, IOException {
  249.     checkRank(rank);
  250.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentURI(
  251.         rank2subRank.getLong(rank));
  252.   }

  254.   /* (non-Javadoc)
  255.    * @see gate.mimir.search.QueryRunner#getDocumentTitle(int)
  256.    */
  257.   @Override
  258.   public String getDocumentTitle(long rank) throws IndexException,
  259.     IndexOutOfBoundsException, IOException {
  260.     checkRank(rank);
  261.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentTitle(
  262.         rank2subRank.getLong(rank));
  263.   }

  265.   /* (non-Javadoc)
  266.    * @see gate.mimir.search.QueryRunner#getDocumentMetadataField(int, java.lang.String)
  267.    */
  268.   @Override
  269.   public Serializable getDocumentMetadataField(long rank, String fieldName)
  270.     throws IndexException, IndexOutOfBoundsException, IOException {
  271.     checkRank(rank);
  272.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentMetadataField(
  273.         rank2subRank.getLong(rank), fieldName);
  274.   }

  276.   /* (non-Javadoc)
  277.    * @see gate.mimir.search.QueryRunner#getDocumentMetadataFields(int, java.util.Set)
  278.    */
  279.   @Override
  280.   public Map<String, Serializable> getDocumentMetadataFields(long rank,
  281.                                                              Set<String> fieldNames)
  282.     throws IndexException, IndexOutOfBoundsException, IOException {
  283.     checkRank(rank);
  284.     return subRunners[rank2runnerIndex.getInt(rank)].getDocumentMetadataFields(
  285.           rank2subRank.getLong(rank), fieldNames);
  286.   }

  288.   /* (non-Javadoc)
  289.    * @see gate.mimir.search.QueryRunner#renderDocument(int, java.lang.Appendable)
  290.    */
  291.   @Override
  292.   public void renderDocument(long rank, Appendable out) throws IOException,
  293.     IndexException {
  294.     checkRank(rank);
  295.     subRunners[rank2runnerIndex.getInt(rank)].renderDocument(
  296.         rank2subRank.getLong(rank), out);
  297.   }

  299.   /* (non-Javadoc)
  300.    * @see gate.mimir.search.QueryRunner#close()
  301.    */
  302.   @Override
  303.   public void close() throws IOException {
  304.     for(QueryRunner r : subRunners) {
  305.       try{
  306.         r.close();
  307.       } catch (Throwable t) {
  308.         log.error("Error while closing sub-runner", t);
  309.       }
  310.     }
  311.   }
  312. }
Created with JaCoCo 0.8.5.201910111838