content: package org.apache.lucene.search;
  
  /**
   * 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 org.apache.lucene.util.PriorityQueue;
  
  /** Represents hits returned by {@link
   * Searcher#search(Query,Filter,int)} and {@link
   * Searcher#search(Query,int)}. */
  public class TopDocs implements java.io.Serializable {
  
    /** The total number of hits for the query. */
    public int totalHits;
  
    /** The top hits for the query. */
    public ScoreDoc[] scoreDocs;
  
    /** Stores the maximum score value encountered, needed for normalizing. */
    private float maxScore;
    private double sum,max,min;
    /**
     * Returns the maximum score value encountered. Note that in case
     * scores are not tracked, this returns {@link Float#NaN}.
     */
    public float getMaxScore() {
      return maxScore;
    }
    
    /** Sets the maximum score value encountered. */
    public void setMaxScore(float maxScore) {
      this.maxScore=maxScore;
    }
  
    public double getSum() { return sum; }
    public void setSum(double sum) { this.sum=sum; }
    public double getMax() { return max; }
    public void setMax(double max) { this.max=max; }
    public double getMin() { return min; }
    public void setMin(double min) { this.min=min; }
  
    /** Constructs a TopDocs with a default maxScore=Float.NaN. */
    TopDocs(int totalHits, ScoreDoc[] scoreDocs) {
      this(totalHits, scoreDocs, Float.NaN, 0.0f, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY);
    }
  
    public TopDocs(int totalHits, ScoreDoc[] scoreDocs, float maxScore, double sum, double max, double min) {
      this.totalHits = totalHits;
      this.scoreDocs = scoreDocs;
      this.maxScore = maxScore;
      this.sum = sum;
      this.max = max;
      this.min = min;
    }
  
    // Refers to one hit:
    private static class ShardRef {
      // Which shard (index into shardHits[]):
      final int shardIndex;
  
      // Which hit within the shard:
      int hitIndex;
  
      public ShardRef(int shardIndex) {
        this.shardIndex = shardIndex;
      }
  
      @Override
      public String toString() {
        return "ShardRef(shardIndex=" + shardIndex + " hitIndex=" + hitIndex + ")";
      }
    };
  
    // Specialized MergeSortQueue that just merges by
    // relevance score, descending:
    private static class ScoreMergeSortQueue extends PriorityQueue<ShardRef> {
      final ScoreDoc[][] shardHits;
  
      public ScoreMergeSortQueue(TopDocs[] shardHits) {
        initialize(shardHits.length);
        this.shardHits = new ScoreDoc[shardHits.length][];
        for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
          this.shardHits[shardIDX] = shardHits[shardIDX].scoreDocs;
        }
      }
  
      // Returns true if first is < second
      public boolean lessThan(ShardRef first, ShardRef second) {
        assert first != second;
        final float firstScore = shardHits[first.shardIndex][first.hitIndex].score;
        final float secondScore = shardHits[second.shardIndex][second.hitIndex].score;
  
        if (firstScore < secondScore) {
          return false;
        } else if (firstScore > secondScore) {
          return true;
        } else {
          // Tie break: earlier shard wins
          if (first.shardIndex < second.shardIndex) {
            return true;
          } else if (first.shardIndex > second.shardIndex) {
            return false;
          } else {
            // Tie break in same shard: resolve however the
            // shard had resolved it:
            assert first.hitIndex != second.hitIndex;
            return first.hitIndex < second.hitIndex;
          }
        }
      }
    }
  
    private static class MergeSortQueue extends PriorityQueue<ShardRef> {
      // These are really FieldDoc instances:
      final ScoreDoc[][] shardHits;
      final FieldComparator[] comparators;
      final int[] reverseMul;
  
      public MergeSortQueue(Sort sort, TopDocs[] shardHits) throws IOException {
        initialize(shardHits.length);
        this.shardHits = new ScoreDoc[shardHits.length][];
        for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
          final ScoreDoc[] shard = shardHits[shardIDX].scoreDocs;
          //System.out.println("  init shardIdx=" + shardIDX + " hits=" + shard);
          if (shard != null) {
            this.shardHits[shardIDX] = shard;
            // Fail gracefully if API is misused:
            for(int hitIDX=0;hitIDX<shard.length;hitIDX++) {
              final ScoreDoc sd = shard[hitIDX];
              if (!(sd instanceof FieldDoc)) {
                throw new IllegalArgumentException("shard " + shardIDX + " was not sorted by the provided Sort (expected FieldDoc but got ScoreDoc)");
              }
              final FieldDoc fd = (FieldDoc) sd;
              if (fd.fields == null) {
                throw new IllegalArgumentException("shard " + shardIDX + " did not set sort field values (FieldDoc.fields is null); you must pass fillFields=true to IndexSearcher.search on each shard");
              }
            }
          }
        }
  
        final SortField[] sortFields = sort.getSort();
        comparators = new FieldComparator[sortFields.length];
        reverseMul = new int[sortFields.length];
        for(int compIDX=0;compIDX<sortFields.length;compIDX++) {
          final SortField sortField = sortFields[compIDX];
          comparators[compIDX] = sortField.getComparator(1, compIDX);
          reverseMul[compIDX] = sortField.getReverse() ? -1 : 1;
        }
      }
  
      // Returns true if first is < second
      @SuppressWarnings("unchecked")
      public boolean lessThan(ShardRef first, ShardRef second) {
        assert first != second;
        final FieldDoc firstFD = (FieldDoc) shardHits[first.shardIndex][first.hitIndex];
        final FieldDoc secondFD = (FieldDoc) shardHits[second.shardIndex][second.hitIndex];
        //System.out.println("  lessThan:\n     first=" + first + " doc=" + firstFD.doc + " score=" + firstFD.score + "\n    second=" + second + " doc=" + secondFD.doc + " score=" + secondFD.score);
  
        for(int compIDX=0;compIDX<comparators.length;compIDX++) {
          final FieldComparator comp = comparators[compIDX];
          //System.out.println("    cmp idx=" + compIDX + " cmp1=" + firstFD.fields[compIDX] + " cmp2=" + secondFD.fields[compIDX] + " reverse=" + reverseMul[compIDX]);
  
          final int cmp = reverseMul[compIDX] * comp.compareValues(firstFD.fields[compIDX], secondFD.fields[compIDX]);
          
          if (cmp != 0) {
            //System.out.println("    return " + (cmp < 0));
            return cmp < 0;
          }
        }
  
        // Tie break: earlier shard wins
        if (first.shardIndex < second.shardIndex) {
          //System.out.println("    return tb true");
          return true;
        } else if (first.shardIndex > second.shardIndex) {
          //System.out.println("    return tb false");
          return false;
        } else {
          // Tie break in same shard: resolve however the
          // shard had resolved it:
          //System.out.println("    return tb " + (first.hitIndex < second.hitIndex));
          assert first.hitIndex != second.hitIndex;
          return first.hitIndex < second.hitIndex;
        }
      }
    }
  
    /** Returns a new TopDocs, containing topN results across
     *  the provided TopDocs, sorting by the specified {@link
     *  Sort}.  Each of the TopDocs must have been sorted by
     *  the same Sort, and sort field values must have been
     *  filled (ie, <code>fillFields=true</code> must be
     *  passed to {@link
     *  TopFieldCollector#create}.
     *
     * <p>Pass sort=null to merge sort by score descending.
     *
     * @lucene.experimental */
    public static TopDocs merge(Sort sort, int topN, TopDocs[] shardHits) throws IOException {
  
      final PriorityQueue<ShardRef> queue;
      if (sort == null) {
        queue = new ScoreMergeSortQueue(shardHits);
      } else {
        queue = new MergeSortQueue(sort, shardHits);
      }
  
      int totalHitCount = 0;
      int availHitCount = 0;
      float maxScore = Float.MIN_VALUE;
      double sum = 0.0f;
      double max = Double.NEGATIVE_INFINITY;
      double min = Double.POSITIVE_INFINITY;
      for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
        final TopDocs shard = shardHits[shardIDX];
        if (shard.scoreDocs != null && shard.scoreDocs.length > 0) {
          totalHitCount += shard.totalHits;
          availHitCount += shard.scoreDocs.length;
          queue.add(new ShardRef(shardIDX));
          maxScore = Math.max(maxScore, shard.getMaxScore());
          sum += shard.getSum();
          max = Math.max(max, shard.getMax());
          min = Math.min(min, shard.getMin());
          //System.out.println("  maxScore now " + maxScore + " vs " + shard.getMaxScore());
        }
      }
  
      final ScoreDoc[] hits = new ScoreDoc[Math.min(topN, availHitCount)];
  
      int hitUpto = 0;
      while(hitUpto < hits.length) {
        assert queue.size() > 0;
        ShardRef ref = queue.pop();
        final ScoreDoc hit = shardHits[ref.shardIndex].scoreDocs[ref.hitIndex++];
        hit.shardIndex = ref.shardIndex;
        hits[hitUpto] = hit;
  
        //System.out.println("  hitUpto=" + hitUpto);
        //System.out.println("    doc=" + hits[hitUpto].doc + " score=" + hits[hitUpto].score);
  
        hitUpto++;
  
        if (ref.hitIndex < shardHits[ref.shardIndex].scoreDocs.length) {
          // Not done with this these TopDocs yet:
          queue.add(ref);
        }
      }
  
      if (sort == null) {
        return new TopDocs(totalHitCount, hits, maxScore, sum, max, min);
      } else {
        return new TopFieldDocs(totalHitCount, hits, sort.getSort(), maxScore, sum, max, min);
      }
    }
  }
  