content: package org.apache.lucene.search.grouping;
  
  /**
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  import java.io.IOException;
  import java.util.*;
  
  import org.apache.lucene.index.IndexReader;
  import org.apache.lucene.search.*;
  
  /** FirstPassGroupingCollector is the first of two passes necessary
   *  to collect grouped hits.  This pass gathers the top N sorted
   *  groups. Concrete subclasses define what a group is and how it
   *  is internally collected.
   *
   *  <p>See {@link org.apache.lucene.search.grouping} for more
   *  details including a full code example.</p>
   *
   * @lucene.experimental
   */
  abstract public class AbstractFirstPassGroupingCollector<GROUP_VALUE_TYPE> extends Collector {
  
    private final Sort groupSort;
    private final FieldComparator[] comparators;
    private final int[] reversed;
    private final int topNGroups;
    private final HashMap<GROUP_VALUE_TYPE, CollectedSearchGroup<GROUP_VALUE_TYPE>> groupMap;
    private final int compIDXEnd;
  
    // Set once we reach topNGroups unique groups:
    private TreeSet<CollectedSearchGroup<GROUP_VALUE_TYPE>> orderedGroups;
    private int docBase;
    private int spareSlot;
  
    /**
     * Create the first pass collector.
     *
     *  @param groupSort The {@link Sort} used to sort the
     *    groups.  The top sorted document within each group
     *    according to groupSort, determines how that group
     *    sorts against other groups.  This must be non-null,
     *    ie, if you want to groupSort by relevance use
     *    Sort.RELEVANCE.
     *  @param topNGroups How many top groups to keep.
     *  @throws IOException If I/O related errors occur
     */
    public AbstractFirstPassGroupingCollector(Sort groupSort, int topNGroups) throws IOException {
      if (topNGroups < 1) {
        throw new IllegalArgumentException("topNGroups must be >= 1 (got " + topNGroups + ")");
      }
  
      // TODO: allow null groupSort to mean "by relevance",
      // and specialize it?
      this.groupSort = groupSort;
  
      this.topNGroups = topNGroups;
  
      final SortField[] sortFields = groupSort.getSort();
      comparators = new FieldComparator[sortFields.length];
      compIDXEnd = comparators.length - 1;
      reversed = new int[sortFields.length];
      for (int i = 0; i < sortFields.length; i++) {
        final SortField sortField = sortFields[i];
  
        // use topNGroups + 1 so we have a spare slot to use for comparing (tracked by this.spareSlot):
        comparators[i] = sortField.getComparator(topNGroups + 1, i);
        reversed[i] = sortField.getReverse() ? -1 : 1;
      }
  
      spareSlot = topNGroups;
      groupMap = new HashMap<GROUP_VALUE_TYPE, CollectedSearchGroup<GROUP_VALUE_TYPE>>(topNGroups);
    }
  
    /**
     * Returns top groups, starting from offset.  This may
     * return null, if no groups were collected, or if the
     * number of unique groups collected is <= offset.
     *
     * @param groupOffset The offset in the collected groups
     * @param fillFields Whether to fill to {@link SearchGroup#sortValues}
     * @return top groups, starting from offset
     */
    public Collection<SearchGroup<GROUP_VALUE_TYPE>> getTopGroups(int groupOffset, boolean fillFields) {
  
      //System.out.println("FP.getTopGroups groupOffset=" + groupOffset + " fillFields=" + fillFields + " groupMap.size()=" + groupMap.size());
  
      if (groupOffset < 0) {
        throw new IllegalArgumentException("groupOffset must be >= 0 (got " + groupOffset + ")");
      }
  
      if (groupMap.size() <= groupOffset) {
        return null;
      }
  
      if (orderedGroups == null) {
        buildSortedSet();
      }
  
      final Collection<SearchGroup<GROUP_VALUE_TYPE>> result = new ArrayList<SearchGroup<GROUP_VALUE_TYPE>>();
      int upto = 0;
      final int sortFieldCount = groupSort.getSort().length;
      for(CollectedSearchGroup<GROUP_VALUE_TYPE> group : orderedGroups) {
        if (upto++ < groupOffset) {
          continue;
        }
        //System.out.println("  group=" + (group.groupValue == null ? "null" : group.groupValue.utf8ToString()));
        SearchGroup<GROUP_VALUE_TYPE> searchGroup = new SearchGroup<GROUP_VALUE_TYPE>();
        searchGroup.groupValue = group.groupValue;
        if (fillFields) {
          searchGroup.sortValues = new Object[sortFieldCount];
          for(int sortFieldIDX=0;sortFieldIDX<sortFieldCount;sortFieldIDX++) {
            searchGroup.sortValues[sortFieldIDX] = comparators[sortFieldIDX].value(group.comparatorSlot);
          }
        }
        result.add(searchGroup);
      }
      //System.out.println("  return " + result.size() + " groups");
      return result;
    }
  
    @Override
    public void setScorer(Scorer scorer) throws IOException {
      for (FieldComparator comparator : comparators) {
        comparator.setScorer(scorer);
      }
    }
  
    @Override
    public void collect(int doc) throws IOException {
      //System.out.println("FP.collect doc=" + doc);
  
      // If orderedGroups != null we already have collected N groups and
      // can short circuit by comparing this document to the bottom group,
      // without having to find what group this document belongs to.
      
      // Even if this document belongs to a group in the top N, we'll know that
      // we don't have to update that group.
  
      // Downside: if the number of unique groups is very low, this is
      // wasted effort as we will most likely be updating an existing group.
      if (orderedGroups != null) {
        for (int compIDX = 0;; compIDX++) {
          final int c = reversed[compIDX] * comparators[compIDX].compareBottom(doc);
          if (c < 0) {
            // Definitely not competitive. So don't even bother to continue
            return;
          } else if (c > 0) {
            // Definitely competitive.
            break;
          } else if (compIDX == compIDXEnd) {
            // Here c=0. If we're at the last comparator, this doc is not
            // competitive, since docs are visited in doc Id order, which means
            // this doc cannot compete with any other document in the queue.
            return;
          }
        }
      }
  
      // TODO: should we add option to mean "ignore docs that
      // don't have the group field" (instead of stuffing them
      // under null group)?
      final GROUP_VALUE_TYPE groupValue = getDocGroupValue(doc);
  
      final CollectedSearchGroup<GROUP_VALUE_TYPE> group = groupMap.get(groupValue);
  
      if (group == null) {
  
        // First time we are seeing this group, or, we've seen
        // it before but it fell out of the top N and is now
        // coming back
  
        if (groupMap.size() < topNGroups) {
  
          // Still in startup transient: we have not
          // seen enough unique groups to start pruning them;
          // just keep collecting them
  
          // Add a new CollectedSearchGroup:
          CollectedSearchGroup<GROUP_VALUE_TYPE> sg = new CollectedSearchGroup<GROUP_VALUE_TYPE>();
          sg.groupValue = copyDocGroupValue(groupValue, null);
          sg.comparatorSlot = groupMap.size();
          sg.topDoc = docBase + doc;
          for (FieldComparator fc : comparators) {
            fc.copy(sg.comparatorSlot, doc);
          }
          groupMap.put(sg.groupValue, sg);
  
          if (groupMap.size() == topNGroups) {
            // End of startup transient: we now have max
            // number of groups; from here on we will drop
            // bottom group when we insert new one:
            buildSortedSet();
          }
  
          return;
        }
  
        // We already tested that the document is competitive, so replace
        // the bottom group with this new group.
  
        // java 6-only: final CollectedSearchGroup bottomGroup = orderedGroups.pollLast();
        final CollectedSearchGroup<GROUP_VALUE_TYPE> bottomGroup = orderedGroups.last();
        orderedGroups.remove(bottomGroup);
        assert orderedGroups.size() == topNGroups -1;
  
        groupMap.remove(bottomGroup.groupValue);
  
        // reuse the removed CollectedSearchGroup
        bottomGroup.groupValue = copyDocGroupValue(groupValue, bottomGroup.groupValue);
        bottomGroup.topDoc = docBase + doc;
  
        for (FieldComparator fc : comparators) {
          fc.copy(bottomGroup.comparatorSlot, doc);
        }
  
        groupMap.put(bottomGroup.groupValue, bottomGroup);
        orderedGroups.add(bottomGroup);
        assert orderedGroups.size() == topNGroups;
  
        final int lastComparatorSlot = orderedGroups.last().comparatorSlot;
        for (FieldComparator fc : comparators) {
          fc.setBottom(lastComparatorSlot);
        }
  
        return;
      }
  
      // Update existing group:
      for (int compIDX = 0;; compIDX++) {
        final FieldComparator fc = comparators[compIDX];
        fc.copy(spareSlot, doc);
  
        final int c = reversed[compIDX] * fc.compare(group.comparatorSlot, spareSlot);
        if (c < 0) {
          // Definitely not competitive.
          return;
        } else if (c > 0) {
          // Definitely competitive; set remaining comparators:
          for (int compIDX2=compIDX+1; compIDX2<comparators.length; compIDX2++) {
            comparators[compIDX2].copy(spareSlot, doc);
          }
          break;
        } else if (compIDX == compIDXEnd) {
          // Here c=0. If we're at the last comparator, this doc is not
          // competitive, since docs are visited in doc Id order, which means
          // this doc cannot compete with any other document in the queue.
          return;
        }
      }
  
      // Remove before updating the group since lookup is done via comparators
      // TODO: optimize this
  
      final CollectedSearchGroup<GROUP_VALUE_TYPE> prevLast;
      if (orderedGroups != null) {
        prevLast = orderedGroups.last();
        orderedGroups.remove(group);
        assert orderedGroups.size() == topNGroups-1;
      } else {
        prevLast = null;
      }
  
      group.topDoc = docBase + doc;
  
      // Swap slots
      final int tmp = spareSlot;
      spareSlot = group.comparatorSlot;
      group.comparatorSlot = tmp;
  
      // Re-add the changed group
      if (orderedGroups != null) {
        orderedGroups.add(group);
        assert orderedGroups.size() == topNGroups;
        final CollectedSearchGroup newLast = orderedGroups.last();
        // If we changed the value of the last group, or changed which group was last, then update bottom:
        if (group == newLast || prevLast != newLast) {
          for (FieldComparator fc : comparators) {
            fc.setBottom(newLast.comparatorSlot);
          }
        }
      }
    }
  
    private void buildSortedSet() {
      final Comparator<CollectedSearchGroup> comparator = new Comparator<CollectedSearchGroup>() {
        public int compare(CollectedSearchGroup o1, CollectedSearchGroup o2) {
          for (int compIDX = 0;; compIDX++) {
            FieldComparator fc = comparators[compIDX];
            final int c = reversed[compIDX] * fc.compare(o1.comparatorSlot, o2.comparatorSlot);
            if (c != 0) {
              return c;
            } else if (compIDX == compIDXEnd) {
              return o1.topDoc - o2.topDoc;
            }
          }
        }
      };
  
      orderedGroups = new TreeSet<CollectedSearchGroup<GROUP_VALUE_TYPE>>(comparator);
      orderedGroups.addAll(groupMap.values());
      assert orderedGroups.size() > 0;
  
      for (FieldComparator fc : comparators) {
        fc.setBottom(orderedGroups.last().comparatorSlot);
      }
    }
  
    @Override
    public boolean acceptsDocsOutOfOrder() {
      return false;
    }
  
    @Override
    public void setNextReader(IndexReader reader, int docBase) throws IOException {
      this.docBase = docBase;
      for (int i=0; i<comparators.length; i++) {
        comparators[i].setNextReader(reader, docBase);
      }
    }
  
    /**
     * Returns the group value for the specified doc.
     *
     * @param doc The specified doc
     * @return the group value for the specified doc
     */
    protected abstract GROUP_VALUE_TYPE getDocGroupValue(int doc);
  
    /**
     * Returns a copy of the specified group value by creating a new instance and copying the value from the specified
     * groupValue in the new instance. Or optionally the reuse argument can be used to copy the group value in.
     *
     * @param groupValue The group value to copy
     * @param reuse Optionally a reuse instance to prevent a new instance creation
     * @return a copy of the specified group value
     */
    protected abstract GROUP_VALUE_TYPE copyDocGroupValue(GROUP_VALUE_TYPE groupValue, GROUP_VALUE_TYPE reuse);
  }
  
  class CollectedSearchGroup<T> extends SearchGroup<T> {
    int topDoc;
    int comparatorSlot;
  }
  