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 java.io.IOException;
21 import java.util.Collection;
22 import java.util.HashMap;
25 import org.apache.lucene.index.IndexReader;
26 import org.apache.lucene.search.*;
29 * SecondPassGroupingCollector is the second of two passes
30 * necessary to collect grouped docs. This pass gathers the
31 * top N documents per top group computed from the
32 * first pass. Concrete subclasses define what a group is and how it
33 * is internally collected.
35 * <p>See {@link org.apache.lucene.search.grouping} for more
36 * details including a full code example.</p>
38 * @lucene.experimental
40 public abstract class AbstractSecondPassGroupingCollector<GROUP_VALUE_TYPE> extends Collector {
42 protected final Map<GROUP_VALUE_TYPE, SearchGroupDocs<GROUP_VALUE_TYPE>> groupMap;
43 private final int maxDocsPerGroup;
44 protected SearchGroupDocs<GROUP_VALUE_TYPE>[] groupDocs;
45 private final Collection<SearchGroup<GROUP_VALUE_TYPE>> groups;
46 private final Sort withinGroupSort;
47 private final Sort groupSort;
49 private int totalHitCount;
50 private int totalGroupedHitCount;
52 public AbstractSecondPassGroupingCollector(Collection<SearchGroup<GROUP_VALUE_TYPE>> groups, Sort groupSort, Sort withinGroupSort,
53 int maxDocsPerGroup, boolean getScores, boolean getMaxScores, boolean fillSortFields)
56 //System.out.println("SP init");
57 if (groups.size() == 0) {
58 throw new IllegalArgumentException("no groups to collect (groups.size() is 0)");
61 this.groupSort = groupSort;
62 this.withinGroupSort = withinGroupSort;
64 this.maxDocsPerGroup = maxDocsPerGroup;
65 groupMap = new HashMap<GROUP_VALUE_TYPE, SearchGroupDocs<GROUP_VALUE_TYPE>>(groups.size());
67 for (SearchGroup<GROUP_VALUE_TYPE> group : groups) {
68 //System.out.println(" prep group=" + (group.groupValue == null ? "null" : group.groupValue.utf8ToString()));
69 final TopDocsCollector collector;
70 if (withinGroupSort == null) {
72 collector = TopScoreDocCollector.create(maxDocsPerGroup, true);
75 collector = TopFieldCollector.create(withinGroupSort, maxDocsPerGroup, fillSortFields, getScores, getMaxScores, true);
77 groupMap.put(group.groupValue,
78 new SearchGroupDocs<GROUP_VALUE_TYPE>(group.groupValue,
84 public void setScorer(Scorer scorer) throws IOException {
85 for (SearchGroupDocs<GROUP_VALUE_TYPE> group : groupMap.values()) {
86 group.collector.setScorer(scorer);
91 public void collect(int doc) throws IOException {
93 SearchGroupDocs<GROUP_VALUE_TYPE> group = retrieveGroup(doc);
95 totalGroupedHitCount++;
96 group.collector.collect(doc);
101 * Returns the group the specified doc belongs to or <code>null</code> if no group could be retrieved.
103 * @param doc The specified doc
104 * @return the group the specified doc belongs to or <code>null</code> if no group could be retrieved
105 * @throws IOException If an I/O related error occurred
107 protected abstract SearchGroupDocs<GROUP_VALUE_TYPE> retrieveGroup(int doc) throws IOException;
110 public void setNextReader(IndexReader reader, int docBase) throws IOException {
111 //System.out.println("SP.setNextReader");
112 for (SearchGroupDocs<GROUP_VALUE_TYPE> group : groupMap.values()) {
113 group.collector.setNextReader(reader, docBase);
118 public boolean acceptsDocsOutOfOrder() {
122 public TopGroups<GROUP_VALUE_TYPE> getTopGroups(int withinGroupOffset) {
123 @SuppressWarnings("unchecked")
124 final GroupDocs<GROUP_VALUE_TYPE>[] groupDocsResult = (GroupDocs<GROUP_VALUE_TYPE>[]) new GroupDocs[groups.size()];
127 for(SearchGroup group : groups) {
128 final SearchGroupDocs<GROUP_VALUE_TYPE> groupDocs = groupMap.get(group.groupValue);
129 final TopDocs topDocs = groupDocs.collector.topDocs(withinGroupOffset, maxDocsPerGroup);
130 groupDocsResult[groupIDX++] = new GroupDocs<GROUP_VALUE_TYPE>(topDocs.getMaxScore(),
133 groupDocs.groupValue,
137 return new TopGroups<GROUP_VALUE_TYPE>(groupSort.getSort(),
138 withinGroupSort == null ? null : withinGroupSort.getSort(),
139 totalHitCount, totalGroupedHitCount, groupDocsResult);
143 // TODO: merge with SearchGroup or not?
144 // ad: don't need to build a new hashmap
145 // disad: blows up the size of SearchGroup if we need many of them, and couples implementations
146 public class SearchGroupDocs<GROUP_VALUE_TYPE> {
148 public final GROUP_VALUE_TYPE groupValue;
149 public final TopDocsCollector collector;
151 public SearchGroupDocs(GROUP_VALUE_TYPE groupValue, TopDocsCollector collector) {
152 this.groupValue = groupValue;
153 this.collector = collector;