lucene-java-3.5.0/lucene/contrib/grouping/src/java/org/apache/lucene/search/grouping/AbstractAllGroupHeadsCollector.java

   1 package org.apache.lucene.search.grouping;
   2
   3 /*
   4  * Licensed to the Apache Software Foundation (ASF) under one or more
   5  * contributor license agreements.  See the NOTICE file distributed with
   6  * this work for additional information regarding copyright ownership.
   7  * The ASF licenses this file to You under the Apache License, Version 2.0
   8  * (the "License"); you may not use this file except in compliance with
   9  * the License.  You may obtain a copy of the License at
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  13  * Unless required by applicable law or agreed to in writing, software
  14  * distributed under the License is distributed on an "AS IS" BASIS,
  15  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16  * See the License for the specific language governing permissions and
  17  * limitations under the License.
  18  */
  19
  20 import org.apache.lucene.index.IndexReader;
  21 import org.apache.lucene.search.Collector;
  22 import org.apache.lucene.util.FixedBitSet;
  23
  24 import java.io.IOException;
  25 import java.util.Collection;
  26
  27 /**
  28  * This collector specializes in collecting the most relevant document (group head) for each group that match the query.
  29  *
  30  * @lucene.experimental
  31  */
  32 public abstract class AbstractAllGroupHeadsCollector<GH extends AbstractAllGroupHeadsCollector.GroupHead> extends Collector {
  33
  34   protected final int[] reversed;
  35   protected final int compIDXEnd;
  36   protected final TemporalResult temporalResult;
  37
  38   protected AbstractAllGroupHeadsCollector(int numberOfSorts) {
  39     this.reversed = new int[numberOfSorts];
  40     this.compIDXEnd = numberOfSorts - 1;
  41     temporalResult = new TemporalResult();
  42   }
  43
  44   /**
  45    * @param maxDoc The maxDoc of the top level {@link IndexReader}.
  46    * @return an {@link FixedBitSet} containing all group heads.
  47    */
  48   public FixedBitSet retrieveGroupHeads(int maxDoc) {
  49     FixedBitSet bitSet = new FixedBitSet(maxDoc);
  50
  51     Collection<GH> groupHeads = getCollectedGroupHeads();
  52     for (GroupHead groupHead : groupHeads) {
  53       bitSet.set(groupHead.doc);
  54     }
  55
  56     return bitSet;
  57   }
  58
  59   /**
  60    * @return an int array containing all group heads. The size of the array is equal to number of collected unique groups.
  61    */
  62   public int[] retrieveGroupHeads() {
  63     Collection<GH> groupHeads = getCollectedGroupHeads();
  64     int[] docHeads = new int[groupHeads.size()];
  65
  66     int i = 0;
  67     for (GroupHead groupHead : groupHeads) {
  68       docHeads[i++] = groupHead.doc;
  69     }
  70
  71     return docHeads;
  72   }
  73
  74   /**
  75    * @return the number of group heads found for a query.
  76    */
  77   public int groupHeadsSize() {
  78     return getCollectedGroupHeads().size();
  79   }
  80
  81   /**
  82    * Returns the group head and puts it into {@link #temporalResult}.
  83    * If the group head wasn't encountered before then it will be added to the collected group heads.
  84    * <p/>
  85    * The {@link TemporalResult#stop} property will be <code>true</code> if the group head wasn't encountered before
  86    * otherwise <code>false</code>.
  87    *
  88    * @param doc The document to retrieve the group head for.
  89    * @throws IOException If I/O related errors occur
  90    */
  91   protected abstract void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException;
  92
  93   /**
  94    * Returns the collected group heads.
  95    * Subsequent calls should return the same group heads.
  96    *
  97    * @return the collected group heads
  98    */
  99   protected abstract Collection<GH> getCollectedGroupHeads();
 100
 101   public void collect(int doc) throws IOException {
 102     retrieveGroupHeadAndAddIfNotExist(doc);
 103     if (temporalResult.stop) {
 104       return;
 105     }
 106     GH groupHead = temporalResult.groupHead;
 107
 108     // Ok now we need to check if the current doc is more relevant then current doc for this group
 109     for (int compIDX = 0; ; compIDX++) {
 110       final int c = reversed[compIDX] * groupHead.compare(compIDX, doc);
 111       if (c < 0) {
 112         // Definitely not competitive. So don't even bother to continue
 113         return;
 114       } else if (c > 0) {
 115         // Definitely competitive.
 116         break;
 117       } else if (compIDX == compIDXEnd) {
 118         // Here c=0. If we're at the last comparator, this doc is not
 119         // competitive, since docs are visited in doc Id order, which means
 120         // this doc cannot compete with any other document in the queue.
 121         return;
 122       }
 123     }
 124     groupHead.updateDocHead(doc);
 125   }
 126
 127   public boolean acceptsDocsOutOfOrder() {
 128     return true;
 129   }
 130
 131   /**
 132    * Contains the result of group head retrieval.
 133    * To prevent new object creations of this class for every collect.
 134    */
 135   protected class TemporalResult {
 136
 137     public GH groupHead;
 138     public boolean stop;
 139
 140   }
 141
 142   /**
 143    * Represents a group head. A group head is the most relevant document for a particular group.
 144    * The relevancy is based is usually based on the sort.
 145    *
 146    * The group head contains a group value with its associated most relevant document id.
 147    */
 148   public static abstract class GroupHead<GROUP_VALUE_TYPE> {
 149
 150     public final GROUP_VALUE_TYPE groupValue;
 151     public int doc;
 152
 153     protected GroupHead(GROUP_VALUE_TYPE groupValue, int doc) {
 154       this.groupValue = groupValue;
 155       this.doc = doc;
 156     }
 157
 158     /**
 159      * Compares the specified document for a specified comparator against the current most relevant document.
 160      *
 161      * @param compIDX The comparator index of the specified comparator.
 162      * @param doc The specified document.
 163      * @return -1 if the specified document wasn't competitive against the current most relevant document, 1 if the
 164      *         specified document was competitive against the current most relevant document. Otherwise 0.
 165      * @throws IOException If I/O related errors occur
 166      */
 167     protected abstract int compare(int compIDX, int doc) throws IOException;
 168
 169     /**
 170      * Updates the current most relevant document with the specified document.
 171      *
 172      * @param doc The specified document
 173      * @throws IOException If I/O related errors occur
 174      */
 175     protected abstract void updateDocHead(int doc) throws IOException;
 176
 177   }
 178
 179 }