1 package org.apache.lucene.search.grouping;
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
11 * http://www.apache.org/licenses/LICENSE-2.0
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.
20 import org.apache.lucene.index.IndexReader;
21 import org.apache.lucene.search.Collector;
22 import org.apache.lucene.util.FixedBitSet;
24 import java.io.IOException;
25 import java.util.Collection;
28 * This collector specializes in collecting the most relevant document (group head) for each group that match the query.
30 * @lucene.experimental
32 public abstract class AbstractAllGroupHeadsCollector<GH extends AbstractAllGroupHeadsCollector.GroupHead> extends Collector {
34 protected final int[] reversed;
35 protected final int compIDXEnd;
36 protected final TemporalResult temporalResult;
38 protected AbstractAllGroupHeadsCollector(int numberOfSorts) {
39 this.reversed = new int[numberOfSorts];
40 this.compIDXEnd = numberOfSorts - 1;
41 temporalResult = new TemporalResult();
45 * @param maxDoc The maxDoc of the top level {@link IndexReader}.
46 * @return an {@link FixedBitSet} containing all group heads.
48 public FixedBitSet retrieveGroupHeads(int maxDoc) {
49 FixedBitSet bitSet = new FixedBitSet(maxDoc);
51 Collection<GH> groupHeads = getCollectedGroupHeads();
52 for (GroupHead groupHead : groupHeads) {
53 bitSet.set(groupHead.doc);
60 * @return an int array containing all group heads. The size of the array is equal to number of collected unique groups.
62 public int[] retrieveGroupHeads() {
63 Collection<GH> groupHeads = getCollectedGroupHeads();
64 int[] docHeads = new int[groupHeads.size()];
67 for (GroupHead groupHead : groupHeads) {
68 docHeads[i++] = groupHead.doc;
75 * @return the number of group heads found for a query.
77 public int groupHeadsSize() {
78 return getCollectedGroupHeads().size();
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.
85 * The {@link TemporalResult#stop} property will be <code>true</code> if the group head wasn't encountered before
86 * otherwise <code>false</code>.
88 * @param doc The document to retrieve the group head for.
89 * @throws IOException If I/O related errors occur
91 protected abstract void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException;
94 * Returns the collected group heads.
95 * Subsequent calls should return the same group heads.
97 * @return the collected group heads
99 protected abstract Collection<GH> getCollectedGroupHeads();
101 public void collect(int doc) throws IOException {
102 retrieveGroupHeadAndAddIfNotExist(doc);
103 if (temporalResult.stop) {
106 GH groupHead = temporalResult.groupHead;
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);
112 // Definitely not competitive. So don't even bother to continue
115 // Definitely competitive.
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.
124 groupHead.updateDocHead(doc);
127 public boolean acceptsDocsOutOfOrder() {
132 * Contains the result of group head retrieval.
133 * To prevent new object creations of this class for every collect.
135 protected class TemporalResult {
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.
146 * The group head contains a group value with its associated most relevant document id.
148 public static abstract class GroupHead<GROUP_VALUE_TYPE> {
150 public final GROUP_VALUE_TYPE groupValue;
153 protected GroupHead(GROUP_VALUE_TYPE groupValue, int doc) {
154 this.groupValue = groupValue;
159 * Compares the specified document for a specified comparator against the current most relevant document.
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
167 protected abstract int compare(int compIDX, int doc) throws IOException;
170 * Updates the current most relevant document with the specified document.
172 * @param doc The specified document
173 * @throws IOException If I/O related errors occur
175 protected abstract void updateDocHead(int doc) throws IOException;