lucene/grouping/src/java/org/apache/lucene/search/grouping/TopGroups.java - lucene-solr - Git at Google

 /*
  * 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.
  */
 package org.apache.lucene.search.grouping;

 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopFieldDocs;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.search.TotalHits.Relation;

 /** Represents result returned by a grouping search.
  *
  * @lucene.experimental */
 public class TopGroups<T> {
   /** Number of documents matching the search */
   public final int totalHitCount;

   /** Number of documents grouped into the topN groups */
   public final int totalGroupedHitCount;

   /** The total number of unique groups. If <code>null</code> this value is not computed. */
   public final Integer totalGroupCount;

   /** Group results in groupSort order */
   public final GroupDocs<T>[] groups;

   /** How groups are sorted against each other */
   public final SortField[] groupSort;

   /** How docs are sorted within each group */
   public final SortField[] withinGroupSort;

   /** Highest score across all hits, or
    *  <code>Float.NaN</code> if scores were not computed. */
   public final float maxScore;

   public TopGroups(SortField[] groupSort, SortField[] withinGroupSort, int totalHitCount, int totalGroupedHitCount, GroupDocs<T>[] groups, float maxScore) {
     this.groupSort = groupSort;
     this.withinGroupSort = withinGroupSort;
     this.totalHitCount = totalHitCount;
     this.totalGroupedHitCount = totalGroupedHitCount;
     this.groups = groups;
     this.totalGroupCount = null;
     this.maxScore = maxScore;
   }

   public TopGroups(TopGroups<T> oldTopGroups, Integer totalGroupCount) {
     this.groupSort = oldTopGroups.groupSort;
     this.withinGroupSort = oldTopGroups.withinGroupSort;
     this.totalHitCount = oldTopGroups.totalHitCount;
     this.totalGroupedHitCount = oldTopGroups.totalGroupedHitCount;
     this.groups = oldTopGroups.groups;
     this.maxScore = oldTopGroups.maxScore;
     this.totalGroupCount = totalGroupCount;
   }

   /** How the GroupDocs score (if any) should be merged. */
   public enum ScoreMergeMode {
     /** Set score to Float.NaN */
     None,
     /* Sum score across all shards for this group. */
     Total,
     /* Avg score across all shards for this group. */
     Avg,
   }

   /**
    * If either value is NaN then return the other value, otherwise
    * return the greater of the two values by calling Math.max.
    * @param a - one value
    * @param b - another value
    * @return ignoring any NaN return the greater of a and b
    */
   private static float nonNANmax(float a, float b) {
     if (Float.isNaN(a)) return b;
     if (Float.isNaN(b)) return a;
     return Math.max(a, b);
   }

   /** Merges an array of TopGroups, for example obtained
    *  from the second-pass collector across multiple
    *  shards.  Each TopGroups must have been sorted by the
    *  same groupSort and docSort, and the top groups passed
    *  to all second-pass collectors must be the same.
    *
    * <b>NOTE</b>: We can't always compute an exact totalGroupCount.
    * Documents belonging to a group may occur on more than
    * one shard and thus the merged totalGroupCount can be
    * higher than the actual totalGroupCount. In this case the
    * totalGroupCount represents a upper bound. If the documents
    * of one group do only reside in one shard then the
    * totalGroupCount is exact.
    *
    * <b>NOTE</b>: the topDocs in each GroupDocs is actually
    * an instance of TopDocsAndShards
    */
   public static <T> TopGroups<T> merge(TopGroups<T>[] shardGroups, Sort groupSort, Sort docSort, int docOffset, int docTopN, ScoreMergeMode scoreMergeMode) {

     //System.out.println("TopGroups.merge");

     if (shardGroups.length == 0) {
       return null;
     }

     int totalHitCount = 0;
     int totalGroupedHitCount = 0;
     // Optionally merge the totalGroupCount.
     Integer totalGroupCount = null;

     final int numGroups = shardGroups[0].groups.length;
     for(TopGroups<T> shard : shardGroups) {
       if (numGroups != shard.groups.length) {
         throw new IllegalArgumentException("number of groups differs across shards; you must pass same top groups to all shards' second-pass collector");
       }
       totalHitCount += shard.totalHitCount;
       totalGroupedHitCount += shard.totalGroupedHitCount;
       if (shard.totalGroupCount != null) {
         if (totalGroupCount == null) {
           totalGroupCount = 0;
         }

         totalGroupCount += shard.totalGroupCount;
       }
     }

     @SuppressWarnings({"unchecked","rawtypes"})
     final GroupDocs<T>[] mergedGroupDocs = new GroupDocs[numGroups];

     final TopDocs[] shardTopDocs;
     if (docSort.equals(Sort.RELEVANCE)) {
       shardTopDocs = new TopDocs[shardGroups.length];
     } else {
       shardTopDocs = new TopFieldDocs[shardGroups.length];
     }
     float totalMaxScore = Float.NaN;

     for(int groupIDX=0;groupIDX<numGroups;groupIDX++) {
       final T groupValue = shardGroups[0].groups[groupIDX].groupValue;
       //System.out.println("  merge groupValue=" + groupValue + " sortValues=" + Arrays.toString(shardGroups[0].groups[groupIDX].groupSortValues));
       float maxScore = Float.NaN;
       int totalHits = 0;
       double scoreSum = 0.0;
       for(int shardIDX=0;shardIDX<shardGroups.length;shardIDX++) {
         //System.out.println("    shard=" + shardIDX);
         final TopGroups<T> shard = shardGroups[shardIDX];
         final GroupDocs<?> shardGroupDocs = shard.groups[groupIDX];
         if (groupValue == null) {
           if (shardGroupDocs.groupValue != null) {
             throw new IllegalArgumentException("group values differ across shards; you must pass same top groups to all shards' second-pass collector");
           }
         } else if (!groupValue.equals(shardGroupDocs.groupValue)) {
           throw new IllegalArgumentException("group values differ across shards; you must pass same top groups to all shards' second-pass collector");
         }

         /*
         for(ScoreDoc sd : shardGroupDocs.scoreDocs) {
           System.out.println("      doc=" + sd.doc);
         }
         */

         if (docSort.equals(Sort.RELEVANCE)) {
           shardTopDocs[shardIDX] = new TopDocs(shardGroupDocs.totalHits,
                                                shardGroupDocs.scoreDocs);
         } else {
           shardTopDocs[shardIDX] = new TopFieldDocs(shardGroupDocs.totalHits,
               shardGroupDocs.scoreDocs,
               docSort.getSort());
         }
         maxScore =  nonNANmax(maxScore, shardGroupDocs.maxScore);
         assert shardGroupDocs.totalHits.relation == Relation.EQUAL_TO;
         totalHits += shardGroupDocs.totalHits.value;
         scoreSum += shardGroupDocs.score;
       }

       final TopDocs mergedTopDocs;
       if (docSort.equals(Sort.RELEVANCE)) {
         mergedTopDocs = TopDocs.merge(docOffset + docTopN, shardTopDocs);
       } else {
         mergedTopDocs = TopDocs.merge(docSort, docOffset + docTopN, (TopFieldDocs[]) shardTopDocs);
       }

       // Slice;
       final ScoreDoc[] mergedScoreDocs;
       if (docOffset == 0) {
         mergedScoreDocs = mergedTopDocs.scoreDocs;
       } else if (docOffset >= mergedTopDocs.scoreDocs.length) {
         mergedScoreDocs = new ScoreDoc[0];
       } else {
         mergedScoreDocs = new ScoreDoc[mergedTopDocs.scoreDocs.length - docOffset];
         System.arraycopy(mergedTopDocs.scoreDocs,
                          docOffset,
                          mergedScoreDocs,
                          0,
                          mergedTopDocs.scoreDocs.length - docOffset);
       }

       final float groupScore;
       switch(scoreMergeMode) {
       case None:
         groupScore = Float.NaN;
         break;
       case Avg:
         if (totalHits > 0) {
           groupScore = (float) (scoreSum / totalHits);
         } else {
           groupScore = Float.NaN;
         }
         break;
       case Total:
         groupScore = (float) scoreSum;
         break;
       default:
         throw new IllegalArgumentException("can't handle ScoreMergeMode " + scoreMergeMode);
       }

       //System.out.println("SHARDS=" + Arrays.toString(mergedTopDocs.shardIndex));
       mergedGroupDocs[groupIDX] = new GroupDocs<>(groupScore,
                                                    maxScore,
                                                    new TotalHits(totalHits, TotalHits.Relation.EQUAL_TO),
                                                    mergedScoreDocs,
                                                    groupValue,
                                                    shardGroups[0].groups[groupIDX].groupSortValues);
       totalMaxScore = nonNANmax(totalMaxScore, maxScore);
     }

     if (totalGroupCount != null) {
       TopGroups<T> result = new TopGroups<>(groupSort.getSort(),
                               docSort.getSort(),
                               totalHitCount,
                               totalGroupedHitCount,
                               mergedGroupDocs,
                               totalMaxScore);
       return new TopGroups<>(result, totalGroupCount);
     } else {
       return new TopGroups<>(groupSort.getSort(),
                               docSort.getSort(),
                               totalHitCount,
                               totalGroupedHitCount,
                               mergedGroupDocs,
                               totalMaxScore);
     }
   }
 }
	/*
	* 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.
	*/
	package org.apache.lucene.search.grouping;

	import org.apache.lucene.search.ScoreDoc;
	import org.apache.lucene.search.Sort;
	import org.apache.lucene.search.SortField;
	import org.apache.lucene.search.TopDocs;
	import org.apache.lucene.search.TopFieldDocs;
	import org.apache.lucene.search.TotalHits;
	import org.apache.lucene.search.TotalHits.Relation;

	/** Represents result returned by a grouping search.
	*
	* @lucene.experimental */
	public class TopGroups<T> {
	/** Number of documents matching the search */
	public final int totalHitCount;

	/** Number of documents grouped into the topN groups */
	public final int totalGroupedHitCount;

	/** The total number of unique groups. If <code>null</code> this value is not computed. */
	public final Integer totalGroupCount;

	/** Group results in groupSort order */
	public final GroupDocs<T>[] groups;

	/** How groups are sorted against each other */
	public final SortField[] groupSort;

	/** How docs are sorted within each group */
	public final SortField[] withinGroupSort;

	/** Highest score across all hits, or
	* <code>Float.NaN</code> if scores were not computed. */
	public final float maxScore;

	public TopGroups(SortField[] groupSort, SortField[] withinGroupSort, int totalHitCount, int totalGroupedHitCount, GroupDocs<T>[] groups, float maxScore) {
	this.groupSort = groupSort;
	this.withinGroupSort = withinGroupSort;
	this.totalHitCount = totalHitCount;
	this.totalGroupedHitCount = totalGroupedHitCount;
	this.groups = groups;
	this.totalGroupCount = null;
	this.maxScore = maxScore;
	}

	public TopGroups(TopGroups<T> oldTopGroups, Integer totalGroupCount) {
	this.groupSort = oldTopGroups.groupSort;
	this.withinGroupSort = oldTopGroups.withinGroupSort;
	this.totalHitCount = oldTopGroups.totalHitCount;
	this.totalGroupedHitCount = oldTopGroups.totalGroupedHitCount;
	this.groups = oldTopGroups.groups;
	this.maxScore = oldTopGroups.maxScore;
	this.totalGroupCount = totalGroupCount;
	}

	/** How the GroupDocs score (if any) should be merged. */
	public enum ScoreMergeMode {
	/** Set score to Float.NaN */
	None,
	/* Sum score across all shards for this group. */
	Total,
	/* Avg score across all shards for this group. */
	Avg,
	}

	/**
	* If either value is NaN then return the other value, otherwise
	* return the greater of the two values by calling Math.max.
	* @param a - one value
	* @param b - another value
	* @return ignoring any NaN return the greater of a and b
	*/
	private static float nonNANmax(float a, float b) {
	if (Float.isNaN(a)) return b;
	if (Float.isNaN(b)) return a;
	return Math.max(a, b);
	}

	/** Merges an array of TopGroups, for example obtained
	* from the second-pass collector across multiple
	* shards. Each TopGroups must have been sorted by the
	* same groupSort and docSort, and the top groups passed
	* to all second-pass collectors must be the same.
	*
	* <b>NOTE</b>: We can't always compute an exact totalGroupCount.
	* Documents belonging to a group may occur on more than
	* one shard and thus the merged totalGroupCount can be
	* higher than the actual totalGroupCount. In this case the
	* totalGroupCount represents a upper bound. If the documents
	* of one group do only reside in one shard then the
	* totalGroupCount is exact.
	*
	* <b>NOTE</b>: the topDocs in each GroupDocs is actually
	* an instance of TopDocsAndShards
	*/
	public static <T> TopGroups<T> merge(TopGroups<T>[] shardGroups, Sort groupSort, Sort docSort, int docOffset, int docTopN, ScoreMergeMode scoreMergeMode) {

	//System.out.println("TopGroups.merge");

	if (shardGroups.length == 0) {
	return null;
	}

	int totalHitCount = 0;
	int totalGroupedHitCount = 0;
	// Optionally merge the totalGroupCount.
	Integer totalGroupCount = null;

	final int numGroups = shardGroups[0].groups.length;
	for(TopGroups<T> shard : shardGroups) {
	if (numGroups != shard.groups.length) {
	throw new IllegalArgumentException("number of groups differs across shards; you must pass same top groups to all shards' second-pass collector");
	}
	totalHitCount += shard.totalHitCount;
	totalGroupedHitCount += shard.totalGroupedHitCount;
	if (shard.totalGroupCount != null) {
	if (totalGroupCount == null) {
	totalGroupCount = 0;
	}

	totalGroupCount += shard.totalGroupCount;
	}
	}

	@SuppressWarnings({"unchecked","rawtypes"})
	final GroupDocs<T>[] mergedGroupDocs = new GroupDocs[numGroups];

	final TopDocs[] shardTopDocs;
	if (docSort.equals(Sort.RELEVANCE)) {
	shardTopDocs = new TopDocs[shardGroups.length];
	} else {
	shardTopDocs = new TopFieldDocs[shardGroups.length];
	}
	float totalMaxScore = Float.NaN;

	for(int groupIDX=0;groupIDX<numGroups;groupIDX++) {
	final T groupValue = shardGroups[0].groups[groupIDX].groupValue;
	//System.out.println(" merge groupValue=" + groupValue + " sortValues=" + Arrays.toString(shardGroups[0].groups[groupIDX].groupSortValues));
	float maxScore = Float.NaN;
	int totalHits = 0;
	double scoreSum = 0.0;
	for(int shardIDX=0;shardIDX<shardGroups.length;shardIDX++) {
	//System.out.println(" shard=" + shardIDX);
	final TopGroups<T> shard = shardGroups[shardIDX];
	final GroupDocs<?> shardGroupDocs = shard.groups[groupIDX];
	if (groupValue == null) {
	if (shardGroupDocs.groupValue != null) {
	throw new IllegalArgumentException("group values differ across shards; you must pass same top groups to all shards' second-pass collector");
	}
	} else if (!groupValue.equals(shardGroupDocs.groupValue)) {
	throw new IllegalArgumentException("group values differ across shards; you must pass same top groups to all shards' second-pass collector");
	}

	/*
	for(ScoreDoc sd : shardGroupDocs.scoreDocs) {
	System.out.println(" doc=" + sd.doc);
	}
	*/

	if (docSort.equals(Sort.RELEVANCE)) {
	shardTopDocs[shardIDX] = new TopDocs(shardGroupDocs.totalHits,
	shardGroupDocs.scoreDocs);
	} else {
	shardTopDocs[shardIDX] = new TopFieldDocs(shardGroupDocs.totalHits,
	shardGroupDocs.scoreDocs,
	docSort.getSort());
	}
	maxScore = nonNANmax(maxScore, shardGroupDocs.maxScore);
	assert shardGroupDocs.totalHits.relation == Relation.EQUAL_TO;
	totalHits += shardGroupDocs.totalHits.value;
	scoreSum += shardGroupDocs.score;
	}

	final TopDocs mergedTopDocs;
	if (docSort.equals(Sort.RELEVANCE)) {
	mergedTopDocs = TopDocs.merge(docOffset + docTopN, shardTopDocs);
	} else {
	mergedTopDocs = TopDocs.merge(docSort, docOffset + docTopN, (TopFieldDocs[]) shardTopDocs);
	}

	// Slice;
	final ScoreDoc[] mergedScoreDocs;
	if (docOffset == 0) {
	mergedScoreDocs = mergedTopDocs.scoreDocs;
	} else if (docOffset >= mergedTopDocs.scoreDocs.length) {
	mergedScoreDocs = new ScoreDoc[0];
	} else {
	mergedScoreDocs = new ScoreDoc[mergedTopDocs.scoreDocs.length - docOffset];
	System.arraycopy(mergedTopDocs.scoreDocs,
	docOffset,
	mergedScoreDocs,
	0,
	mergedTopDocs.scoreDocs.length - docOffset);
	}

	final float groupScore;
	switch(scoreMergeMode) {
	case None:
	groupScore = Float.NaN;
	break;
	case Avg:
	if (totalHits > 0) {
	groupScore = (float) (scoreSum / totalHits);
	} else {
	groupScore = Float.NaN;
	}
	break;
	case Total:
	groupScore = (float) scoreSum;
	break;
	default:
	throw new IllegalArgumentException("can't handle ScoreMergeMode " + scoreMergeMode);
	}

	//System.out.println("SHARDS=" + Arrays.toString(mergedTopDocs.shardIndex));
	mergedGroupDocs[groupIDX] = new GroupDocs<>(groupScore,
	maxScore,
	new TotalHits(totalHits, TotalHits.Relation.EQUAL_TO),
	mergedScoreDocs,
	groupValue,
	shardGroups[0].groups[groupIDX].groupSortValues);
	totalMaxScore = nonNANmax(totalMaxScore, maxScore);
	}

	if (totalGroupCount != null) {
	TopGroups<T> result = new TopGroups<>(groupSort.getSort(),
	docSort.getSort(),
	totalHitCount,
	totalGroupedHitCount,
	mergedGroupDocs,
	totalMaxScore);
	return new TopGroups<>(result, totalGroupCount);
	} else {
	return new TopGroups<>(groupSort.getSort(),
	docSort.getSort(),
	totalHitCount,
	totalGroupedHitCount,
	mergedGroupDocs,
	totalMaxScore);
	}
	}
	}