Add support for skipping over non-competitive documents in FirstPassGroupingCollector

Author: alexmm-amzn

lucene/grouping/src/java/org/apache/lucene/search/grouping/FirstPassGroupingCollector.java

@@ -50,9 +50,10 @@ public class FirstPassGroupingCollector<T> extends SimpleCollector {
   private final LeafFieldComparator[] leafComparators;
   private final int[] reversed;
   private final int topNGroups;
-  private final boolean needsScores;
   private final HashMap<T, CollectedSearchGroup<T>> groupMap;
   private final int compIDXEnd;
+  private final ScoreMode scoreMode;
+  private final boolean canSetMinScore;
 
   // Set once we reach topNGroups unique groups:
   /**
@@ -62,6 +63,9 @@ public class FirstPassGroupingCollector<T> extends SimpleCollector {
 
   private int docBase;
   private int spareSlot;
+  private Scorable scorer;
+  private int bottomSlot;
+  private float minCompetitiveScore;
 
   /**
    * Create the first pass collector.
@@ -83,7 +87,6 @@ public FirstPassGroupingCollector(
     // and specialize it?
 
     this.topNGroups = topNGroups;
-    this.needsScores = groupSort.needsScores();
     final SortField[] sortFields = groupSort.getSort();
     comparators = new FieldComparator<?>[sortFields.length];
     leafComparators = new LeafFieldComparator[sortFields.length];
@@ -105,13 +108,21 @@ public FirstPassGroupingCollector(
       reversed[i] = sortField.getReverse() ? -1 : 1;
     }
 
+    if (SortField.FIELD_SCORE.equals(sortFields[0]) == true) {
+      scoreMode = ScoreMode.TOP_SCORES;
+      canSetMinScore = true;
+    } else {
+      scoreMode = groupSort.needsScores() ? ScoreMode.TOP_DOCS_WITH_SCORES : ScoreMode.TOP_DOCS;
+      canSetMinScore = false;
+    }
+
     spareSlot = topNGroups;
     groupMap = CollectionUtil.newHashMap(topNGroups);
   }
 
   @Override
   public ScoreMode scoreMode() {
-    return needsScores ? ScoreMode.COMPLETE : ScoreMode.COMPLETE_NO_SCORES;
+    return scoreMode;
   }
 
   /**
@@ -162,10 +173,12 @@ public Collection<SearchGroup<T>> getTopGroups(int groupOffset) throws IOExcepti
 
   @Override
   public void setScorer(Scorable scorer) throws IOException {
+    this.scorer = scorer;
     groupSelector.setScorer(scorer);
     for (LeafFieldComparator comparator : leafComparators) {
       comparator.setScorer(scorer);
     }
+    setMinCompetitiveScore(scorer);
   }
 
   private boolean isCompetitive(int doc) throws IOException {
@@ -273,9 +286,7 @@ private void collectNewGroup(final int doc) throws IOException {
       assert orderedGroups.size() == topNGroups;
 
       final int lastComparatorSlot = orderedGroups.last().comparatorSlot;
-      for (LeafFieldComparator fc : leafComparators) {
-        fc.setBottom(lastComparatorSlot);
-      }
+      setBottomSlot(lastComparatorSlot);
     }
   }
 
@@ -331,9 +342,7 @@ private void collectExistingGroup(final int doc, final CollectedSearchGroup<T> g
       // If we changed the value of the last group, or changed which group was last, then update
       // bottom:
       if (group == newLast || prevLast != newLast) {
-        for (LeafFieldComparator fc : leafComparators) {
-          fc.setBottom(newLast.comparatorSlot);
-        }
+        setBottomSlot(newLast.comparatorSlot);
       }
     }
   }
@@ -364,13 +373,12 @@ public int compare(CollectedSearchGroup<?> o1, CollectedSearchGroup<?> o2) {
     orderedGroups.addAll(groupMap.values());
     assert orderedGroups.size() > 0;
 
-    for (LeafFieldComparator fc : leafComparators) {
-      fc.setBottom(orderedGroups.last().comparatorSlot);
-    }
+    setBottomSlot(orderedGroups.last().comparatorSlot);
   }
 
   @Override
   protected void doSetNextReader(LeafReaderContext readerContext) throws IOException {
+    minCompetitiveScore = 0f;
     docBase = readerContext.docBase;
     for (int i = 0; i < comparators.length; i++) {
       leafComparators[i] = comparators[i].getLeafComparator(readerContext);
@@ -388,4 +396,25 @@ public GroupSelector<T> getGroupSelector() {
   private boolean isGroupMapFull() {
     return groupMap.size() >= topNGroups;
   }
+
+  private void setBottomSlot(final int bottomSlot) throws IOException {
+    for (LeafFieldComparator fc : leafComparators) {
+      fc.setBottom(bottomSlot);
+    }
+
+    this.bottomSlot = bottomSlot;
+    setMinCompetitiveScore(scorer);
+  }
+
+  private void setMinCompetitiveScore(final Scorable scorer) throws IOException {
+    if (canSetMinScore == false || isGroupMapFull() == false) {
+      return;
+    }
+
+    final float minScore = (float) comparators[0].value(bottomSlot);
+    if (minScore > minCompetitiveScore) {
+      scorer.setMinCompetitiveScore(minScore);
+      minCompetitiveScore = minScore;
+    }
+  }
 }

lucene/grouping/src/test/org/apache/lucene/search/grouping/TestGrouping.java

@@ -16,6 +16,8 @@
  */
 package org.apache.lucene.search.grouping;
 
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
+
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -1520,14 +1522,14 @@ private void assertEquals(
         "expected.groups.length != actual.groups.length",
         expected.groups.length,
         actual.groups.length);
-    assertEquals(
-        "expected.totalHitCount != actual.totalHitCount",
-        expected.totalHitCount,
-        actual.totalHitCount);
-    assertEquals(
-        "expected.totalGroupedHitCount != actual.totalGroupedHitCount",
-        expected.totalGroupedHitCount,
-        actual.totalGroupedHitCount);
+    assertThat(
+        "expected.totalHitCount >= actual.totalHitCount",
+        actual.totalHitCount,
+        lessThanOrEqualTo(expected.totalHitCount));
+    assertThat(
+        "expected.totalGroupedHitCount >= actual.totalGroupedHitCount",
+        actual.totalGroupedHitCount,
+        lessThanOrEqualTo(expected.totalGroupedHitCount));
     if (expected.totalGroupCount != null && verifyTotalGroupCount) {
       assertEquals(
           "expected.totalGroupCount != actual.totalGroupCount",
@@ -1556,7 +1558,10 @@ private void assertEquals(
 
       // TODO
       // assertEquals(expectedGroup.maxScore, actualGroup.maxScore);
-      assertEquals(expectedGroup.totalHits().value(), actualGroup.totalHits().value());
+      assertThat(
+          "expectedGroup.totalHits().value() >= actualGroup.totalHits().value()",
+          actualGroup.totalHits().value(),
+          lessThanOrEqualTo(expectedGroup.totalHits().value()));
 
       final ScoreDoc[] expectedFDs = expectedGroup.scoreDocs();
       final ScoreDoc[] actualFDs = actualGroup.scoreDocs();