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 org.apache.lucene.index.IndexReader;
  import org.apache.lucene.search.*;
  
  import java.io.IOException;
  import java.util.*;
  
  /**
   * A base implementation of {@link AbstractAllGroupHeadsCollector} for retrieving the most relevant groups when grouping
   * on a string based group field. More specifically this all concrete implementations of this base implementation
   * use {@link org.apache.lucene.search.FieldCache.StringIndex}.
   *
   * @lucene.experimental
   */
  public abstract class TermAllGroupHeadsCollector<GH extends AbstractAllGroupHeadsCollector.GroupHead> extends AbstractAllGroupHeadsCollector<GH> {
  
    private static final int DEFAULT_INITIAL_SIZE = 128;
  
    final String groupField;
    FieldCache.StringIndex groupIndex;
    IndexReader indexReader;
    int docBase;
  
    protected TermAllGroupHeadsCollector(String groupField, int numberOfSorts) {
      super(numberOfSorts);
      this.groupField = groupField;
    }
  
    /**
     * Creates an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments.
     * This factory method decides with implementation is best suited.
     *
     * Delegates to {@link #create(String, org.apache.lucene.search.Sort, int)} with an initialSize of 128.
     *
     * @param groupField      The field to group by
     * @param sortWithinGroup The sort within each group
     * @return an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments
     * @throws IOException If I/O related errors occur
     */
    public static AbstractAllGroupHeadsCollector create(String groupField, Sort sortWithinGroup) throws IOException {
      return create(groupField, sortWithinGroup, DEFAULT_INITIAL_SIZE);
    }
  
    /**
     * Creates an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments.
     * This factory method decides with implementation is best suited.
     *
     * @param groupField      The field to group by
     * @param sortWithinGroup The sort within each group
     * @param initialSize The initial allocation size of the internal int set and group list which should roughly match
     *                    the total number of expected unique groups. Be aware that the heap usage is
     *                    4 bytes * initialSize.
     * @return an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments
     * @throws IOException If I/O related errors occur
     */
    public static AbstractAllGroupHeadsCollector create(String groupField, Sort sortWithinGroup, int initialSize) throws IOException {
      boolean sortAllScore = true;
      boolean sortAllFieldValue = true;
  
      for (SortField sortField : sortWithinGroup.getSort()) {
        if (sortField.getType() == SortField.SCORE) {
          sortAllFieldValue = false;
        } else if (needGeneralImpl(sortField)) {
          return new GeneralAllGroupHeadsCollector(groupField, sortWithinGroup);
        } else {
          sortAllScore = false;
        }
      }
  
      if (sortAllScore) {
        return new ScoreAllGroupHeadsCollector(groupField, sortWithinGroup, initialSize);
      } else if (sortAllFieldValue) {
        return new OrdAllGroupHeadsCollector(groupField, sortWithinGroup, initialSize);
      } else {
        return new OrdScoreAllGroupHeadsCollector(groupField, sortWithinGroup, initialSize);
      }
    }
  
    // Returns when a sort field needs the general impl.
    private static boolean needGeneralImpl(SortField sortField) {
      int sortType = sortField.getType();
      // Note (MvG): We can also make an optimized impl when sorting is SortField.DOC
      return sortType != SortField.STRING_VAL && sortType != SortField.STRING && sortType != SortField.SCORE;
    }
  
    // A general impl that works for any group sort.
    static class GeneralAllGroupHeadsCollector extends TermAllGroupHeadsCollector<GeneralAllGroupHeadsCollector.GroupHead> {
  
      private final Sort sortWithinGroup;
      private final Map<String, GroupHead> groups;
  
      private Scorer scorer;
  
      GeneralAllGroupHeadsCollector(String groupField, Sort sortWithinGroup) throws IOException {
        super(groupField, sortWithinGroup.getSort().length);
        this.sortWithinGroup = sortWithinGroup;
        groups = new HashMap<String, GroupHead>();
  
        final SortField[] sortFields = sortWithinGroup.getSort();
        for (int i = 0; i < sortFields.length; i++) {
          reversed[i] = sortFields[i].getReverse() ? -1 : 1;
        }
      }
  
      protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
        final int ord = groupIndex.order[doc];
        final String groupValue = ord == 0 ? null : groupIndex.lookup[ord];
        GroupHead groupHead = groups.get(groupValue);
        if (groupHead == null) {
          groupHead = new GroupHead(groupValue, sortWithinGroup, doc);
          groups.put(groupValue == null ? null : groupValue, groupHead);
          temporalResult.stop = true;
        } else {
          temporalResult.stop = false;
        }
        temporalResult.groupHead = groupHead;
      }
  
      protected Collection<GroupHead> getCollectedGroupHeads() {
        return groups.values();
      }
  
      public void setNextReader(IndexReader reader, int docBase) throws IOException {
        this.indexReader = reader;
        this.docBase = docBase;
        groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
  
        for (GroupHead groupHead : groups.values()) {
          for (int i = 0; i < groupHead.comparators.length; i++) {
            groupHead.comparators[i].setNextReader(reader, docBase);
          }
        }
      }
  
      public void setScorer(Scorer scorer) throws IOException {
        this.scorer = scorer;
        for (GroupHead groupHead : groups.values()) {
          for (FieldComparator comparator : groupHead.comparators) {
            comparator.setScorer(scorer);
          }
        }
      }
  
      class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
  
        final FieldComparator[] comparators;
  
        private GroupHead(String groupValue, Sort sort, int doc) throws IOException {
          super(groupValue, doc + docBase);
          final SortField[] sortFields = sort.getSort();
          comparators = new FieldComparator[sortFields.length];
          for (int i = 0; i < sortFields.length; i++) {
            comparators[i] = sortFields[i].getComparator(1, i);
            comparators[i].setNextReader(indexReader, docBase);
            comparators[i].setScorer(scorer);
            comparators[i].copy(0, doc);
            comparators[i].setBottom(0);
          }
        }
  
        public int compare(int compIDX, int doc) throws IOException {
          return comparators[compIDX].compareBottom(doc);
        }
  
        public void updateDocHead(int doc) throws IOException {
          for (FieldComparator comparator : comparators) {
            comparator.copy(0, doc);
            comparator.setBottom(0);
          }
          this.doc = doc + docBase;
        }
      }
    }
  
  
    // AbstractAllGroupHeadsCollector optimized for ord fields and scores.
    static class OrdScoreAllGroupHeadsCollector extends TermAllGroupHeadsCollector<OrdScoreAllGroupHeadsCollector.GroupHead> {
  
      private final SentinelIntSet ordSet;
      private final List<GroupHead> collectedGroups;
      private final SortField[] fields;
  
      private FieldCache.StringIndex[] sortsIndex;
      private Scorer scorer;
      private GroupHead[] segmentGroupHeads;
  
      OrdScoreAllGroupHeadsCollector(String groupField, Sort sortWithinGroup, int initialSize) {
        super(groupField, sortWithinGroup.getSort().length);
        ordSet = new SentinelIntSet(initialSize, -1);
        collectedGroups = new ArrayList<GroupHead>(initialSize);
  
        final SortField[] sortFields = sortWithinGroup.getSort();
        fields = new SortField[sortFields.length];
        sortsIndex = new FieldCache.StringIndex[sortFields.length];
        for (int i = 0; i < sortFields.length; i++) {
          reversed[i] = sortFields[i].getReverse() ? -1 : 1;
          fields[i] = sortFields[i];
        }
      }
  
      protected Collection<GroupHead> getCollectedGroupHeads() {
        return collectedGroups;
      }
  
      public void setScorer(Scorer scorer) throws IOException {
        this.scorer = scorer;
      }
  
      protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
        int key = groupIndex.order[doc];
        GroupHead groupHead;
        if (!ordSet.exists(key)) {
          ordSet.put(key);
          String term = key == 0 ? null : groupIndex.lookup[key];
          groupHead = new GroupHead(doc, term);
          collectedGroups.add(groupHead);
          segmentGroupHeads[key] = groupHead;
          temporalResult.stop = true;
        } else {
          temporalResult.stop = false;
          groupHead = segmentGroupHeads[key];
        }
        temporalResult.groupHead = groupHead;
      }
  
      public void setNextReader(IndexReader reader, int docBase) throws IOException {
        this.indexReader = reader;
        this.docBase = docBase;
        groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
        for (int i = 0; i < fields.length; i++) {
          if (fields[i].getType() == SortField.SCORE) {
            continue;
          }
  
          sortsIndex[i] = FieldCache.DEFAULT.getStringIndex(reader, fields[i].getField());
        }
  
        // Clear ordSet and fill it with previous encountered groups that can occur in the current segment.
        ordSet.clear();
        segmentGroupHeads = new GroupHead[groupIndex.lookup.length];
        for (GroupHead collectedGroup : collectedGroups) {
          int ord = groupIndex.binarySearchLookup(collectedGroup.groupValue);
          if (ord >= 0) {
            ordSet.put(ord);
            segmentGroupHeads[ord] = collectedGroup;
  
            for (int i = 0; i < sortsIndex.length; i++) {
              if (fields[i].getType() == SortField.SCORE) {
                continue;
              }
  
              collectedGroup.sortOrds[i] = sortsIndex[i].binarySearchLookup(collectedGroup.sortValues[i]);
            }
          }
        }
      }
  
      class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
  
        String[] sortValues;
        int[] sortOrds;
        float[] scores;
  
        private GroupHead(int doc, String groupValue) throws IOException {
          super(groupValue, doc + docBase);
          sortValues = new String[sortsIndex.length];
          sortOrds = new int[sortsIndex.length];
          scores = new float[sortsIndex.length];
          for (int i = 0; i < sortsIndex.length; i++) {
            if (fields[i].getType() == SortField.SCORE) {
              scores[i] = scorer.score();
            } else {
              sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
              sortOrds[i] = sortsIndex[i].order[doc];
            }
          }
  
        }
  
        public int compare(int compIDX, int doc) throws IOException {
          if (fields[compIDX].getType() == SortField.SCORE) {
            float score = scorer.score();
            if (scores[compIDX] < score) {
              return 1;
            } else if (scores[compIDX] > score) {
              return -1;
            }
            return 0;
          } else {
            if (sortOrds[compIDX] < 0) {
              // The current segment doesn't contain the sort value we encountered before. Therefore the ord is negative.
              final String val1 = sortValues[compIDX];
              final String val2 = sortsIndex[compIDX].lookup[sortsIndex[compIDX].order[doc]];
              if (val1 == null) {
                if (val2 == null) {
                  return 0;
                }
                return -1;
              } else if (val2 == null) {
                return 1;
              }
              return val1.compareTo(val2);
            } else {
              return sortOrds[compIDX] - sortsIndex[compIDX].order[doc];
            }
          }
        }
  
        public void updateDocHead(int doc) throws IOException {
          for (int i = 0; i < sortsIndex.length; i++) {
            if (fields[i].getType() == SortField.SCORE) {
              scores[i] = scorer.score();
            } else {
              sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
              sortOrds[i] = sortsIndex[i].order[doc];
            }
          }
          this.doc = doc + docBase;
        }
  
      }
  
    }
  
  
    // AbstractAllGroupHeadsCollector optimized for ord fields.
    static class OrdAllGroupHeadsCollector extends TermAllGroupHeadsCollector<OrdAllGroupHeadsCollector.GroupHead> {
  
      private final SentinelIntSet ordSet;
      private final List<GroupHead> collectedGroups;
      private final SortField[] fields;
  
      private FieldCache.StringIndex[] sortsIndex;
      private GroupHead[] segmentGroupHeads;
  
      OrdAllGroupHeadsCollector(String groupField, Sort sortWithinGroup, int initialSize) {
        super(groupField, sortWithinGroup.getSort().length);
        ordSet = new SentinelIntSet(initialSize, -1);
        collectedGroups = new ArrayList<GroupHead>(initialSize);
  
        final SortField[] sortFields = sortWithinGroup.getSort();
        fields = new SortField[sortFields.length];
        sortsIndex = new FieldCache.StringIndex[sortFields.length];
        for (int i = 0; i < sortFields.length; i++) {
          reversed[i] = sortFields[i].getReverse() ? -1 : 1;
          fields[i] = sortFields[i];
        }
      }
  
      protected Collection<GroupHead> getCollectedGroupHeads() {
        return collectedGroups;
      }
  
      public void setScorer(Scorer scorer) throws IOException {
      }
  
      protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
        int key = groupIndex.order[doc];
        GroupHead groupHead;
        if (!ordSet.exists(key)) {
          ordSet.put(key);
          String term = key == 0 ? null : groupIndex.lookup[key];
          groupHead = new GroupHead(doc, term);
          collectedGroups.add(groupHead);
          segmentGroupHeads[key] = groupHead;
          temporalResult.stop = true;
        } else {
          temporalResult.stop = false;
          groupHead = segmentGroupHeads[key];
        }
        temporalResult.groupHead = groupHead;
      }
  
      public void setNextReader(IndexReader reader, int docBase) throws IOException {
        this.indexReader = reader;
        this.docBase = docBase;
        groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
        for (int i = 0; i < fields.length; i++) {
          sortsIndex[i] = FieldCache.DEFAULT.getStringIndex(reader, fields[i].getField());
        }
  
        // Clear ordSet and fill it with previous encountered groups that can occur in the current segment.
        ordSet.clear();
        segmentGroupHeads = new GroupHead[groupIndex.lookup.length];
        for (GroupHead collectedGroup : collectedGroups) {
          int groupOrd = groupIndex.binarySearchLookup(collectedGroup.groupValue);
          if (groupOrd >= 0) {
            ordSet.put(groupOrd);
            segmentGroupHeads[groupOrd] = collectedGroup;
  
            for (int i = 0; i < sortsIndex.length; i++) {
              collectedGroup.sortOrds[i] = sortsIndex[i].binarySearchLookup(collectedGroup.sortValues[i]);
            }
          }
        }
      }
  
      class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
  
        String[] sortValues;
        int[] sortOrds;
  
        private GroupHead(int doc, String groupValue) throws IOException {
          super(groupValue, doc + docBase);
          sortValues = new String[sortsIndex.length];
          sortOrds = new int[sortsIndex.length];
          for (int i = 0; i < sortsIndex.length; i++) {
            sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
            sortOrds[i] = sortsIndex[i].order[doc];
          }
        }
  
        public int compare(int compIDX, int doc) throws IOException {
          if (sortOrds[compIDX] < 0) {
            // The current segment doesn't contain the sort value we encountered before. Therefore the ord is negative.
            final String val1 = sortValues[compIDX];
            final String val2 = sortsIndex[compIDX].lookup[sortsIndex[compIDX].order[doc]];
            if (val1 == null) {
              if (val2 == null) {
                return 0;
              }
              return -1;
            } else if (val2 == null) {
              return 1;
            }
            return val1.compareTo(val2);
          } else {
            return sortOrds[compIDX] - sortsIndex[compIDX].order[doc];
          }
        }
  
        public void updateDocHead(int doc) throws IOException {
          for (int i = 0; i < sortsIndex.length; i++) {
            sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
            sortOrds[i] = sortsIndex[i].order[doc];
          }
          this.doc = doc + docBase;
        }
  
      }
  
    }
  
  
    // AbstractAllGroupHeadsCollector optimized for scores.
    static class ScoreAllGroupHeadsCollector extends TermAllGroupHeadsCollector<ScoreAllGroupHeadsCollector.GroupHead> {
  
      private final SentinelIntSet ordSet;
      private final List<GroupHead> collectedGroups;
      private final SortField[] fields;
  
      private Scorer scorer;
      private GroupHead[] segmentGroupHeads;
  
      ScoreAllGroupHeadsCollector(String groupField, Sort sortWithinGroup, int initialSize) {
        super(groupField, sortWithinGroup.getSort().length);
        ordSet = new SentinelIntSet(initialSize, -1);
        collectedGroups = new ArrayList<GroupHead>(initialSize);
  
        final SortField[] sortFields = sortWithinGroup.getSort();
        fields = new SortField[sortFields.length];
        for (int i = 0; i < sortFields.length; i++) {
          reversed[i] = sortFields[i].getReverse() ? -1 : 1;
          fields[i] = sortFields[i];
        }
      }
  
      protected Collection<GroupHead> getCollectedGroupHeads() {
        return collectedGroups;
      }
  
      public void setScorer(Scorer scorer) throws IOException {
        this.scorer = scorer;
      }
  
      protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
        int key = groupIndex.order[doc];
        GroupHead groupHead;
        if (!ordSet.exists(key)) {
          ordSet.put(key);
          String term = key == 0 ? null : groupIndex.lookup[key];
          groupHead = new GroupHead(doc, term);
          collectedGroups.add(groupHead);
          segmentGroupHeads[key] = groupHead;
          temporalResult.stop = true;
        } else {
          temporalResult.stop = false;
          groupHead = segmentGroupHeads[key];
        }
        temporalResult.groupHead = groupHead;
      }
  
      public void setNextReader(IndexReader reader, int docBase) throws IOException {
        this.indexReader = reader;
        this.docBase = docBase;
        groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
  
        // Clear ordSet and fill it with previous encountered groups that can occur in the current segment.
        ordSet.clear();
        segmentGroupHeads = new GroupHead[groupIndex.lookup.length];
        for (GroupHead collectedGroup : collectedGroups) {
          int ord = groupIndex.binarySearchLookup(collectedGroup.groupValue);
          if (ord >= 0) {
            ordSet.put(ord);
            segmentGroupHeads[ord] = collectedGroup;
          }
        }
      }
  
      class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
  
        float[] scores;
  
        private GroupHead(int doc, String groupValue) throws IOException {
          super(groupValue, doc + docBase);
          scores = new float[fields.length];
          float score = scorer.score();
          for (int i = 0; i < scores.length; i++) {
            scores[i] = score;
          }
        }
  
        public int compare(int compIDX, int doc) throws IOException {
          float score = scorer.score();
          if (scores[compIDX] < score) {
            return 1;
          } else if (scores[compIDX] > score) {
            return -1;
          }
          return 0;
        }
  
        public void updateDocHead(int doc) throws IOException {
          float score = scorer.score();
          for (int i = 0; i < scores.length; i++) {
            scores[i] = score;
          }
          this.doc = doc + docBase;
        }
  
      }
  
    }
  
  }
  