SOLR-5773: CollapsingQParserPlugin should make elevated documents the group head

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1583500 13f79535-47bb-0310-9956-ffa450edef68

solr/CHANGES.txt

@@ -253,6 +253,9 @@ Other Changes
 * SOLR-5934: LBHttpSolrServer exception handling improvement and small test
   improvements. (Gregory Chanan via Mark Miller)
 
+* SOLR-5773: CollapsingQParserPlugin should make elevated documents the 
+  group head. (David Boychuck, Joel Bernstein)
+
 ==================  4.7.1  ==================
 
 Versions of Major Components

solr/core/src/java/org/apache/solr/search/CollapsingQParserPlugin.java

@@ -443,6 +443,7 @@ public class CollapsingQParserPlugin extends QParserPlugin {
     private int nullDoc;
     private FloatArrayList nullScores;
     private IntOpenHashSet boostDocs;
+    private int[] boostOrds;
 
     public CollapsingScoreCollector(int maxDoc,
                                     int segments,
@@ -455,11 +456,19 @@ public class CollapsingQParserPlugin extends QParserPlugin {
       this.boostDocs = boostDocs;
       if(this.boostDocs != null) {
         //Set the elevated docs now.
+        IntOpenHashSet boostG = new IntOpenHashSet();
         Iterator<IntCursor> it = this.boostDocs.iterator();
         while(it.hasNext()) {
           IntCursor cursor = it.next();
-          this.collapsedSet.set(cursor.value);
+          int i = cursor.value;
+          this.collapsedSet.set(i);
+          int ord = values.getOrd(i);
+          if(ord > -1) {
+            boostG.add(ord);
+          }
         }
+        boostOrds = boostG.toArray();
+        Arrays.sort(boostOrds);
       }
       this.values = values;
       int valueCount = values.getValueCount();
@@ -489,6 +498,7 @@ public class CollapsingQParserPlugin extends QParserPlugin {
     public void collect(int docId) throws IOException {
       int globalDoc = docId+this.docBase;
       int ord = values.getOrd(globalDoc);
+
       if(ord > -1) {
         float score = scorer.score();
         if(score > scores[ord]) {
@@ -520,6 +530,12 @@ public class CollapsingQParserPlugin extends QParserPlugin {
         this.collapsedSet.set(nullDoc);
       }
 
+      if(this.boostOrds != null) {
+        for(int i=0; i<this.boostOrds.length; i++) {
+          ords[boostOrds[i]] = -1;
+        }
+      }
+
       for(int i=0; i<ords.length; i++) {
         int doc = ords[i];
         if(doc > -1) {
@@ -539,6 +555,7 @@ public class CollapsingQParserPlugin extends QParserPlugin {
       while((docId = it.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
 
         int ord = values.getOrd(docId);
+
         if(ord > -1) {
           dummy.score = scores[ord];
         } else if(this.boostDocs != null && boostDocs.contains(docId)) {
@@ -600,14 +617,14 @@ public class CollapsingQParserPlugin extends QParserPlugin {
       this.needsScores = needsScores;
       this.boostDocs = boostDocs;
       if(funcQuery != null) {
-        this.fieldValueCollapse =  new ValueSourceCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs, funcQuery, searcher);
+        this.fieldValueCollapse =  new ValueSourceCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs, funcQuery, searcher, values);
       } else {
         if(fieldType instanceof TrieIntField) {
-          this.fieldValueCollapse = new IntValueCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs);
+          this.fieldValueCollapse = new IntValueCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs, values);
         } else if(fieldType instanceof TrieLongField) {
-          this.fieldValueCollapse =  new LongValueCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs);
+          this.fieldValueCollapse =  new LongValueCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs, values);
         } else if(fieldType instanceof TrieFloatField) {
-          this.fieldValueCollapse =  new FloatValueCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs);
+          this.fieldValueCollapse =  new FloatValueCollapse(maxDoc, field, nullPolicy, new int[valueCount], max, this.needsScores, boostDocs, values);
         } else {
           throw new IOException("min/max must be either TrieInt, TrieLong or TrieFloat.");
         }
@@ -696,6 +713,7 @@ public class CollapsingQParserPlugin extends QParserPlugin {
     protected float[] scores;
     protected FixedBitSet collapsedSet;
     protected IntOpenHashSet boostDocs;
+    protected int[] boostOrds;
     protected int nullDoc = -1;
     protected boolean needsScores;
     protected boolean max;
@@ -709,7 +727,8 @@ public class CollapsingQParserPlugin extends QParserPlugin {
                               int nullPolicy,
                               boolean max,
                               boolean needsScores,
-                              IntOpenHashSet boostDocs) {
+                              IntOpenHashSet boostDocs,
+                              SortedDocValues values) {
       this.field = field;
       this.nullPolicy = nullPolicy;
       this.max = max;
@@ -717,11 +736,19 @@ public class CollapsingQParserPlugin extends QParserPlugin {
       this.collapsedSet = new FixedBitSet(maxDoc);
       this.boostDocs = boostDocs;
       if(this.boostDocs != null) {
+        IntOpenHashSet boostG = new IntOpenHashSet();
         Iterator<IntCursor> it = boostDocs.iterator();
         while(it.hasNext()) {
           IntCursor cursor = it.next();
-          this.collapsedSet.set(cursor.value);
+          int i = cursor.value;
+          this.collapsedSet.set(i);
+          int ord = values.getOrd(i);
+          if(ord > -1) {
+            boostG.add(ord);
+          }
         }
+        this.boostOrds = boostG.toArray();
+        Arrays.sort(this.boostOrds);
       }
     }
 
@@ -730,6 +757,12 @@ public class CollapsingQParserPlugin extends QParserPlugin {
         this.collapsedSet.set(nullDoc);
       }
 
+      if(this.boostOrds != null) {
+        for(int i=0; i<this.boostOrds.length; i++) {
+          ords[boostOrds[i]] = -1;
+        }
+      }
+
       for(int i=0; i<ords.length; i++) {
         int doc = ords[i];
         if(doc > -1) {
@@ -770,8 +803,8 @@ public class CollapsingQParserPlugin extends QParserPlugin {
                             int[] ords,
                             boolean max,
                             boolean needsScores,
-                            IntOpenHashSet boostDocs) throws IOException {
+                            IntOpenHashSet boostDocs, SortedDocValues values) throws IOException {
-      super(maxDoc, field, nullPolicy, max, needsScores, boostDocs);
+      super(maxDoc, field, nullPolicy, max, needsScores, boostDocs, values);
       this.ords = ords;
       this.ordVals = new int[ords.length];
       Arrays.fill(ords, -1);
@@ -838,8 +871,8 @@ public class CollapsingQParserPlugin extends QParserPlugin {
                              int[] ords,
                              boolean max,
                              boolean needsScores,
-                             IntOpenHashSet boostDocs) throws IOException {
+                             IntOpenHashSet boostDocs, SortedDocValues values) throws IOException {
-      super(maxDoc, field, nullPolicy, max, needsScores, boostDocs);
+      super(maxDoc, field, nullPolicy, max, needsScores, boostDocs, values);
       this.ords = ords;
       this.ordVals = new long[ords.length];
       Arrays.fill(ords, -1);
@@ -907,8 +940,8 @@ public class CollapsingQParserPlugin extends QParserPlugin {
                               int[] ords,
                               boolean max,
                               boolean needsScores,
-                              IntOpenHashSet boostDocs) throws IOException {
+                              IntOpenHashSet boostDocs, SortedDocValues values) throws IOException {
-      super(maxDoc, field, nullPolicy, max, needsScores, boostDocs);
+      super(maxDoc, field, nullPolicy, max, needsScores, boostDocs, values);
       this.ords = ords;
       this.ordVals = new float[ords.length];
       Arrays.fill(ords, -1);
@@ -982,8 +1015,8 @@ public class CollapsingQParserPlugin extends QParserPlugin {
                                boolean max,
                                boolean needsScores,
                                IntOpenHashSet boostDocs,
-                               FunctionQuery funcQuery, IndexSearcher searcher) throws IOException {
+                               FunctionQuery funcQuery, IndexSearcher searcher, SortedDocValues values) throws IOException {
-      super(maxDoc, null, nullPolicy, max, needsScores, boostDocs);
+      super(maxDoc, null, nullPolicy, max, needsScores, boostDocs, values);
       this.valueSource = funcQuery.getValueSource();
       this.rcontext = ValueSource.newContext(searcher);
       this.ords = ords;

solr/core/src/test/org/apache/solr/search/TestCollapseQParserPlugin.java

@@ -23,9 +23,11 @@ import org.apache.solr.common.params.ModifiableSolrParams;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import com.carrotsearch.hppc.IntOpenHashSet;
 
 import java.io.IOException;
 import java.util.*;
+import java.util.Random;
 
 public class TestCollapseQParserPlugin extends SolrTestCaseJ4 {
 
@@ -146,6 +148,51 @@ public class TestCollapseQParserPlugin extends SolrTestCaseJ4 {
                          "//result/doc[3]/float[@name='id'][.='3.0']",
                          "//result/doc[4]/float[@name='id'][.='6.0']");
 
+    //Test SOLR-5773 with score collapse criteria
+    params = new ModifiableSolrParams();
+    params.add("q", "YYYY");
+    params.add("fq", "{!collapse field=group_s nullPolicy=collapse}");
+    params.add("defType", "edismax");
+    params.add("bf", "field(test_ti)");
+    params.add("qf", "term_s");
+    params.add("qt", "/elevate");
+    params.add("elevateIds", "1,5");
+    assertQ(req(params), "*[count(//doc)=3]",
+        "//result/doc[1]/float[@name='id'][.='1.0']",
+        "//result/doc[2]/float[@name='id'][.='5.0']",
+        "//result/doc[3]/float[@name='id'][.='3.0']");
+
+    //Test SOLR-5773 with max field collapse criteria
+    params = new ModifiableSolrParams();
+    params.add("q", "YYYY");
+    params.add("fq", "{!collapse field=group_s min=test_ti nullPolicy=collapse}");
+    params.add("defType", "edismax");
+    params.add("bf", "field(test_ti)");
+    params.add("qf", "term_s");
+    params.add("qt", "/elevate");
+    params.add("elevateIds", "1,5");
+    assertQ(req(params), "*[count(//doc)=3]",
+        "//result/doc[1]/float[@name='id'][.='1.0']",
+        "//result/doc[2]/float[@name='id'][.='5.0']",
+        "//result/doc[3]/float[@name='id'][.='4.0']");
+
+
+    //Test SOLR-5773 elevating documents with null group
+    params = new ModifiableSolrParams();
+    params.add("q", "YYYY");
+    params.add("fq", "{!collapse field=group_s}");
+    params.add("defType", "edismax");
+    params.add("bf", "field(test_ti)");
+    params.add("qf", "term_s");
+    params.add("qt", "/elevate");
+    params.add("elevateIds", "3,4");
+    assertQ(req(params), "*[count(//doc)=4]",
+        "//result/doc[1]/float[@name='id'][.='3.0']",
+        "//result/doc[2]/float[@name='id'][.='4.0']",
+        "//result/doc[3]/float[@name='id'][.='2.0']",
+        "//result/doc[4]/float[@name='id'][.='6.0']");
+
+
 
     //Test collapse by min int field and sort
     params = new ModifiableSolrParams();