lucene-java-3.4.0/lucene/contrib/grouping/src/java/org/apache/lucene/search/grouping/AbstractSecondPassGroupingCollector.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 java.io.IOException;
  21 import java.util.Collection;
  22 import java.util.HashMap;
  23 import java.util.Map;
  24
  25 import org.apache.lucene.index.IndexReader;
  26 import org.apache.lucene.search.*;
  27
  28 /**
  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.
  34  *
  35  * <p>See {@link org.apache.lucene.search.grouping} for more
  36  * details including a full code example.</p>
  37  *
  38  * @lucene.experimental
  39  */
  40 public abstract class AbstractSecondPassGroupingCollector<GROUP_VALUE_TYPE> extends Collector {
  41
  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;
  48
  49   private int totalHitCount;
  50   private int totalGroupedHitCount;
  51
  52   public AbstractSecondPassGroupingCollector(Collection<SearchGroup<GROUP_VALUE_TYPE>> groups, Sort groupSort, Sort withinGroupSort,
  53                                              int maxDocsPerGroup, boolean getScores, boolean getMaxScores, boolean fillSortFields)
  54     throws IOException {
  55
  56     //System.out.println("SP init");
  57     if (groups.size() == 0) {
  58       throw new IllegalArgumentException("no groups to collect (groups.size() is 0)");
  59     }
  60
  61     this.groupSort = groupSort;
  62     this.withinGroupSort = withinGroupSort;
  63     this.groups = groups;
  64     this.maxDocsPerGroup = maxDocsPerGroup;
  65     groupMap = new HashMap<GROUP_VALUE_TYPE, SearchGroupDocs<GROUP_VALUE_TYPE>>(groups.size());
  66
  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) {
  71         // Sort by score
  72         collector = TopScoreDocCollector.create(maxDocsPerGroup, true);
  73       } else {
  74         // Sort by fields
  75         collector = TopFieldCollector.create(withinGroupSort, maxDocsPerGroup, fillSortFields, getScores, getMaxScores, true);
  76       }
  77       groupMap.put(group.groupValue,
  78           new SearchGroupDocs<GROUP_VALUE_TYPE>(group.groupValue,
  79               collector));
  80     }
  81   }
  82
  83   @Override
  84   public void setScorer(Scorer scorer) throws IOException {
  85     for (SearchGroupDocs<GROUP_VALUE_TYPE> group : groupMap.values()) {
  86       group.collector.setScorer(scorer);
  87     }
  88   }
  89
  90   @Override
  91   public void collect(int doc) throws IOException {
  92     totalHitCount++;
  93     SearchGroupDocs<GROUP_VALUE_TYPE> group = retrieveGroup(doc);
  94     if (group != null) {
  95       totalGroupedHitCount++;
  96       group.collector.collect(doc);
  97     }
  98   }
  99
 100   /**
 101    * Returns the group the specified doc belongs to or <code>null</code> if no group could be retrieved.
 102    *
 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
 106    */
 107   protected abstract SearchGroupDocs<GROUP_VALUE_TYPE> retrieveGroup(int doc) throws IOException;
 108
 109   @Override
 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);
 114     }
 115   }
 116
 117   @Override
 118   public boolean acceptsDocsOutOfOrder() {
 119     return false;
 120   }
 121
 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()];
 125
 126     int groupIDX = 0;
 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(),
 131                                                                     topDocs.totalHits,
 132                                                                     topDocs.scoreDocs,
 133                                                                     groupDocs.groupValue,
 134                                                                     group.sortValues);
 135     }
 136
 137     return new TopGroups<GROUP_VALUE_TYPE>(groupSort.getSort(),
 138                                            withinGroupSort == null ? null : withinGroupSort.getSort(),
 139                                            totalHitCount, totalGroupedHitCount, groupDocsResult);
 140   }
 141
 142
 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> {
 147
 148     public final GROUP_VALUE_TYPE groupValue;
 149     public final TopDocsCollector collector;
 150
 151     public SearchGroupDocs(GROUP_VALUE_TYPE groupValue, TopDocsCollector collector) {
 152       this.groupValue = groupValue;
 153       this.collector = collector;
 154     }
 155   }
 156 }