Support grouping + reranking for group.func

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

@@ -69,8 +69,7 @@
 import org.apache.solr.schema.FieldType;
 import org.apache.solr.schema.SchemaField;
 import org.apache.solr.schema.StrFieldSource;
-import org.apache.solr.search.grouping.collector.FilterCollector;
-import org.apache.solr.search.grouping.collector.RerankTermSecondPassGroupingCollector;
+import org.apache.solr.search.grouping.collector.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -1005,9 +1004,13 @@ protected Collector createSecondPassCollector() throws IOException {
       int groupdDocsToCollect = getMax(groupOffset, docsPerGroup, maxDoc);
       groupdDocsToCollect = Math.max(groupdDocsToCollect, 1);
       Sort withinGroupSort = this.withinGroupSort != null ? this.withinGroupSort : Sort.RELEVANCE;
-      secondPass = new FunctionSecondPassGroupingCollector(
-          topGroups, groupSort, withinGroupSort, groupdDocsToCollect, needScores, needScores, false, groupBy, context
-      );
+      if (query instanceof RankQuery){
+        secondPass = new RerankFunctionSecondPassGroupingCollector(topGroups, groupSort, withinGroupSort, (RankQuery)query, searcher, groupdDocsToCollect, needScores, needScores, false, groupBy, context);
+      } else {
+        secondPass = new FunctionSecondPassGroupingCollector(
+            topGroups, groupSort, withinGroupSort, groupdDocsToCollect, needScores, needScores, false, groupBy, context
+        );
+      }
 
       if (totalCount == TotalCount.grouped) {
         allGroupsCollector = new FunctionAllGroupsCollector(groupBy, context);
@@ -1029,6 +1032,21 @@ public AllGroupHeadsCollector<?> createAllGroupCollector() throws IOException {
     @Override
     protected void finish() throws IOException {
       result = secondPass != null ? secondPass.getTopGroups(0) : null;
+
+      if (result != null && query instanceof RankQuery && groupSort == Sort.RELEVANCE) {
+        // if we are sorting for relevance and query is a RankQuery, it may be that
+        // the order of the groups changed, we need to reorder
+        GroupDocs[] groups = result.groups;
+        Arrays.sort(groups, new Comparator<GroupDocs>() {
+          @Override
+          public int compare(GroupDocs o1, GroupDocs o2) {
+            if (o1.maxScore > o2.maxScore) return -1;
+            if (o1.maxScore < o2.maxScore) return 1;
+            return 0;
+          }
+        });
+      }
+
       if (main) {
         mainResult = createSimpleResponse();
         return;

solr/core/src/java/org/apache/solr/search/grouping/collector/RerankFunctionSecondPassGroupingCollector.java

@@ -0,0 +1,67 @@
+/*
+ * 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.solr.search.grouping.collector;
+
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Map;
+
+import org.apache.lucene.queries.function.ValueSource;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Sort;
+import org.apache.lucene.search.TopDocsCollector;
+import org.apache.lucene.search.grouping.SearchGroup;
+import org.apache.lucene.search.grouping.TopGroups;
+import org.apache.lucene.search.grouping.function.FunctionSecondPassGroupingCollector;
+import org.apache.lucene.util.mutable.MutableValue;
+import org.apache.solr.search.RankQuery;
+
+public class RerankFunctionSecondPassGroupingCollector extends FunctionSecondPassGroupingCollector {
+
+
+  private static final int DEFAULT_GROUPING_RERANKING = 10;
+
+  /**
+   * Constructs a {@link RerankFunctionSecondPassGroupingCollector} instance.
+   *
+   * @param searchGroups The {@link SearchGroup} instances collected during the first phase.
+   * @param groupSort The group sort
+   * @param withinGroupSort The sort inside a group
+   * @param query The rankquery used to rerank
+   * @param searcher The index searcher
+   * @param maxDocsPerGroup The maximum number of documents to collect inside a group
+   * @param getScores Whether to include the scores
+   * @param getMaxScores Whether to include the maximum score
+   * @param fillSortFields Whether to fill the sort values in {@link TopGroups#withinGroupSort}
+   * @param groupByVS The {@link ValueSource} to group by
+   * @param vsContext The value source context
+   * @throws IOException IOException When I/O related errors occur
+   */
+  public RerankFunctionSecondPassGroupingCollector(Collection<SearchGroup<MutableValue>> searchGroups, Sort groupSort, Sort withinGroupSort, RankQuery query, IndexSearcher searcher, int maxDocsPerGroup, boolean getScores, boolean getMaxScores, boolean fillSortFields, ValueSource groupByVS, Map<?, ?> vsContext) throws IOException {
+    super(searchGroups, groupSort, withinGroupSort, maxDocsPerGroup, getScores, getMaxScores, fillSortFields, groupByVS, vsContext);
+    for (SearchGroup<MutableValue> group : searchGroups) {
+      TopDocsCollector<?> collector;
+      if (query != null) {
+        collector = groupMap.get(group.groupValue).collector;
+        collector = query.getTopDocsCollector(collector, DEFAULT_GROUPING_RERANKING, groupSort, searcher);
+        groupMap.put(group.groupValue, new SearchGroupDocs<MutableValue>(group.groupValue, collector));
+      }
+    }
+  }
+}

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

@@ -21,6 +21,7 @@
 import org.apache.solr.common.params.ModifiableSolrParams;
 import org.apache.solr.common.util.NamedList;
 import org.apache.solr.core.SolrInfoMBean;
+import org.apache.solr.request.*;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -725,12 +726,35 @@ public void testRerankQueryAndGroupingRerankGroups() throws Exception {
         "//arr/lst[2]/result/doc/float[@name='id'][.='3.0']",
         "//arr/lst[3]/result/doc/float[@name='id'][.='2.0']"
     );
-    // test grouping by function
-    params.remove("group.field");
+    // test grouping by function:
+    // documents are
+    // 1. firstly scored by their test_tl score and then
+    // 2. then grouped by the value of log(test_tf)
+    // 3. finally reranked by test_ti
+    // rerank query will add the score so:
+
+    // 1. first pass (doc_id, score):    results = (1, 9)(2, 8) (3, 2) (4, 3) (5, 1)
+    // 2. grouping:                      group1(max score=9) = [(1, 9) (3, 2)],  group2(max score = 8) = [(2, 8) (4, 3) (5, 1)]
+    // 3. reranking (scoring):           group1 = [ (1, 9 + 5), (3, 2 + 100) ]  group2 [ (2, 8 + 50), (4, 3 + 74), (5, 1 + 1000)]
+    //      reordered:                   group2(max score = 1001) [ (5, 1001) (4, 77) (2, 58)) group1(max score = 102) = [ (3, 102) (1, 14)  ]
+    // (actual scores will be offset by 1.0 because of a constant boost (=1) added to the scores)
+
+    params = new ModifiableSolrParams();
+    params.add("q", "{!edismax bq=$bqq1}*:*");
+    params.add("bqq1", "{!func }field(test_tl)");
+    params.add("start", "0");
+    params.add("rows", "6");
+    params.add("fl", "id,score,[explain]");
+    params.add("group", "true");
     params.add("group.func", "log(test_tf)");
+    params.add("rq", "{!" + ReRankQParserPlugin.NAME + " " + ReRankQParserPlugin.RERANK_QUERY + "=$rqq "
+        + ReRankQParserPlugin.RERANK_DOCS + "=200 "+ ReRankQParserPlugin.RERANK_WEIGHT + "=1 }");
+    //rank query, rerank documents on the value of test_ti
+    params.add("rqq", "{!func }field(test_ti)");
+
     assertQ(req(params), "*[count(//doc)=2]",
-        "//arr/lst[1]/result/doc/float[@name='id'][.='1.0']", // should be 3.0
-        "//arr/lst[2]/result/doc/float[@name='id'][.='2.0']"  // should be 4.0
+        "//arr/lst[1]/result/doc/float[@name='id'][.='5.0']", // should be 3.0
+        "//arr/lst[2]/result/doc/float[@name='id'][.='3.0']"  // should be 4.0
     );
 
   }