SOLR-15437: ReRanking/LTR does not work in combination with custom sort and SolrCloud

solr/contrib/ltr/src/test/org/apache/solr/ltr/TestLTROnSolrCloud.java

@@ -16,11 +16,17 @@
 package org.apache.solr.ltr;
 
 import java.io.File;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.SortedMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import com.google.common.base.Splitter;
 import org.apache.commons.io.FileUtils;
 import org.apache.solr.client.solrj.SolrQuery;
 import org.apache.solr.client.solrj.embedded.JettyConfig;
@@ -29,6 +35,8 @@
 import org.apache.solr.client.solrj.response.CollectionAdminResponse;
 import org.apache.solr.client.solrj.response.QueryResponse;
 import org.apache.solr.cloud.MiniSolrCloudCluster;
+import org.apache.solr.common.SolrDocument;
+import org.apache.solr.common.SolrDocumentList;
 import org.apache.solr.common.SolrInputDocument;
 import org.apache.solr.common.cloud.ZkStateReader;
 import org.apache.solr.ltr.feature.FieldValueFeature;
@@ -50,6 +58,10 @@ public class TestLTROnSolrCloud extends TestRerankBase {
   String schema = "schema.xml";
 
   SortedMap<ServletHolder,String> extraServlets = null;
+
+  private final static String MODEL_WEIGHTS = "\"powpularityS\":1.0,\"c3\":1.0,\"original\":0.1," +
+            "\"dvIntFieldFeature\":0.1,\"dvLongFieldFeature\":0.1," +
+            "\"dvFloatFieldFeature\":0.1,\"dvDoubleFieldFeature\":0.1,\"dvStrNumFieldFeature\":0.1,\"dvStrBoolFieldFeature\":0.1";
 
   @Override
   public void setUp() throws Exception {
@@ -63,11 +75,8 @@ public void setUp() throws Exception {
     int numberOfNodes = numberOfShards * numberOfReplicas;
 
     setupSolrCluster(numberOfShards, numberOfReplicas, numberOfNodes);
-
-
   }
 
-
   @Override
   public void tearDown() throws Exception {
     restTestHarness.close();
@@ -76,6 +85,128 @@ public void tearDown() throws Exception {
     super.tearDown();
   }
 
+  @Test
+  public void testSimpleQueryCustomSort() throws Exception {
+    SolrQuery query = new SolrQuery("*:*");
+    query.setRequestHandler("/query");
+    query.setFields("*,[shard]");
+    query.setParam("rows", "8");
+    query.setParam("sort", "id asc");
+    query.add("rq", "{!ltr model=powpularityS-model reRankDocs=8}");
+
+    QueryResponse queryResponse =
+            solrCluster.getSolrClient().query(COLLECTION, query);
+    assertEquals(8, queryResponse.getResults().getNumFound());
+    assertEquals("8", queryResponse.getResults().get(0).get("id").toString());
+    assertEquals("7", queryResponse.getResults().get(1).get("id").toString());
+    assertEquals("6", queryResponse.getResults().get(2).get("id").toString());
+    assertEquals("5", queryResponse.getResults().get(3).get("id").toString());
+    assertEquals("4", queryResponse.getResults().get(4).get("id").toString());
+    assertEquals("3", queryResponse.getResults().get(5).get("id").toString());
+    assertEquals("2", queryResponse.getResults().get(6).get("id").toString());
+    assertEquals("1", queryResponse.getResults().get(7).get("id").toString());
+  }
+
+  @Test
+  public void testSimpleQueryCustomSortWithSubResultSet() throws Exception {
+    final int reRankDocs = 2;
+    SolrQuery query = new SolrQuery("*:*");
+    query.setRequestHandler("/query");
+    query.setFields("*,score,[shard],[fv]");
+    query.setParam("rows", "8");
+    query.setParam("sort", "id asc");
+    query.add("rq", "{!ltr model=powpularityS-model reRankDocs="+reRankDocs+"}");
+
+    QueryResponse queryResponse = solrCluster.getSolrClient().query(COLLECTION, query);
+    SolrDocumentList results = queryResponse.getResults();
+    assertEquals(8, results.getNumFound());
+
+    // save order to use it later
+    List<String> expectedDocIdOrder = new ArrayList<>();
+
+    int docCounter = 0;
+    float lastScore = Float.MAX_VALUE;
+    double lastId = 0d;
+    for(SolrDocument d : results){
+      float score = (float) d.getFieldValue("score");
+
+      double id = Double.parseDouble((String) d.getFirstValue("id"));
+      expectedDocIdOrder.add((String) d.getFirstValue("id"));
+      if(docCounter < reRankDocs){
+        final float calculatedScore = calculateLTRScoreForDoc(d);
+        assertEquals(calculatedScore, score, 0.0);
+        assertTrue(lastScore > score);
+      } else if(docCounter > reRankDocs + 1) {
+        assertTrue(lastId < id);
+      }
+      lastScore = score;
+      lastId = id;
+
+      docCounter++;
+    }
+
+    query.setFields("*,[shard],[fv]");
+
+    queryResponse = solrCluster.getSolrClient().query(COLLECTION, query);
+    results = queryResponse.getResults();
+    assertEquals(8, results.getNumFound());
+
+    List<String> docIdOrder = results.stream()
+        .map(document -> (String) document.getFirstValue("id"))
+        .collect(Collectors.toList());
+
+    // assert that sorting is correct when we do not return the score via fl param
+    assertEquals(expectedDocIdOrder, docIdOrder);
+  }
+
+  @Test
+  public void testSimpleQueryCustomSortWithSubResultSetAndRowsLessThanExistingDocs() throws Exception {
+    final int reRankDocs = 2;
+    SolrQuery query = new SolrQuery("*:*");
+    query.setRequestHandler("/query");
+    query.setFields("*,score,[shard],[fv]");  // score as fl is needed here to be able to use it for assertions
+    query.setParam("rows", "6");
+    query.setParam("sort", "id asc");
+    query.add("rq", "{!ltr model=powpularityS-model reRankDocs="+reRankDocs+"}");
+    QueryResponse queryResponse = solrCluster.getSolrClient().query(COLLECTION, query);
+    SolrDocumentList results = queryResponse.getResults();
+    assertEquals(6, results.size());
+    int docCounter = 0;
+    float lastScore = Float.MAX_VALUE;
+    double lastId = 0d;
+    for(SolrDocument d : results){
+      float score = (float) d.getFieldValue("score");
+      double id = Double.parseDouble((String) d.getFirstValue("id"));
+      if(docCounter < reRankDocs){
+        final float calculatedScore = calculateLTRScoreForDoc(d);
+        assertEquals(calculatedScore, score, 0.0);
+        assertTrue(lastScore > score);
+      } else if(docCounter > reRankDocs + 1) {
+        assertTrue(lastId < id);
+      }
+      lastScore = score;
+      lastId = id;
+      docCounter++;
+    }
+  }
+
+  private float calculateLTRScoreForDoc(SolrDocument d) {
+    Matcher matcher = Pattern.compile(",?(\\w+)=(-?[0-9]+\\.[0-9]+)").matcher((String) d.getFieldValue("[fv]"));
+    Map<String, Float> weights = Splitter.on(",")
+            .splitToList(MODEL_WEIGHTS)
+            .stream()
+            .map(fieldWithWeight -> fieldWithWeight.split(":"))
+            .collect(Collectors.toMap(fieldAndValue -> fieldAndValue[0].replaceAll("\"", ""),
+                fieldAndValue -> Float.parseFloat(fieldAndValue[1])));
+
+    float score = 0.0f;
+    while(matcher.find()) {
+        score += Float.parseFloat(matcher.group(2)) * weights.get(matcher.group(1));
+    }
+
+    return score;
+  }
+
   @Test
   // commented 4-Sep-2018 @LuceneTestCase.BadApple(bugUrl="https://issues.apache.org/jira/browse/SOLR-12028") // 2-Aug-2018
   // commented out on: 24-Dec-2018   @BadApple(bugUrl="https://issues.apache.org/jira/browse/SOLR-12028") // 14-Oct-2018
@@ -298,9 +429,7 @@ private void loadModelsAndFeatures() throws Exception {
     final String featureStore = "test";
     final String[] featureNames = new String[]{"powpularityS", "c3", "original", "dvIntFieldFeature",
         "dvLongFieldFeature", "dvFloatFieldFeature", "dvDoubleFieldFeature", "dvStrNumFieldFeature", "dvStrBoolFieldFeature"};
-    final String jsonModelParams = "{\"weights\":{\"powpularityS\":1.0,\"c3\":1.0,\"original\":0.1," +
-        "\"dvIntFieldFeature\":0.1,\"dvLongFieldFeature\":0.1," +
-        "\"dvFloatFieldFeature\":0.1,\"dvDoubleFieldFeature\":0.1,\"dvStrNumFieldFeature\":0.1,\"dvStrBoolFieldFeature\":0.1}}";
+    final String jsonModelParams = "{\"weights\":{" + MODEL_WEIGHTS + "}}";
 
     loadFeature(
         featureNames[0],

solr/core/src/java/org/apache/solr/handler/component/DefaultSortedHitQueueManager.java

@@ -0,0 +1,51 @@
+/*
+ * 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.handler.component;
+
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.SortField;
+
+/**
+ * This class is just a wrapper for the ShardFieldSortedHitQueue.
+ * Is is used when we do not use reRanking.
+ *
+ * All documents are added to the queue.
+ */
+public class DefaultSortedHitQueueManager extends SortedHitQueueManager {
+
+  private final ShardFieldSortedHitQueue queue;
+
+  public DefaultSortedHitQueueManager(SortField[] sortFields, int count, int offset, IndexSearcher searcher) {
+    queue = new ShardFieldSortedHitQueue(sortFields, count + offset, searcher);
+  }
+
+  @Override
+  public void addDocument(ShardDoc shardDoc) {
+    queue.insertWithOverflow(shardDoc);
+  }
+
+  @Override
+  public ShardDoc popDocument() {
+    return queue.pop();
+  }
+
+  @Override
+  public int size() {
+    return queue.size();
+  }
+
+}

solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java

@@ -73,6 +73,7 @@
 import org.apache.solr.schema.IndexSchema;
 import org.apache.solr.schema.SchemaField;
 import org.apache.solr.schema.SortableTextField;
+import org.apache.solr.search.AbstractReRankQuery;
 import org.apache.solr.search.CursorMark;
 import org.apache.solr.search.DocIterator;
 import org.apache.solr.search.DocList;
@@ -116,6 +117,7 @@
 import org.slf4j.LoggerFactory;
 
 import static org.apache.solr.common.params.CommonParams.QUERY_UUID;
+import static org.apache.solr.handler.component.SortedHitQueueManager.newSortedHitQueueManager;
 
 
 /**
@@ -148,14 +150,8 @@ public void prepare(ResponseBuilder rb) throws IOException
       }
     }
 
-    // Set field flags    
     ReturnFields returnFields = new SolrReturnFields( req );
     rsp.setReturnFields( returnFields );
-    int flags = 0;
-    if (returnFields.wantsScore()) {
-      flags |= SolrIndexSearcher.GET_SCORES;
-    }
-    rb.setFieldFlags( flags );
 
     String defType = params.get(QueryParsing.DEFTYPE, QParserPlugin.DEFAULT_QTYPE);
 
@@ -223,6 +219,14 @@ public void prepare(ResponseBuilder rb) throws IOException
       throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, e);
     }
 
+    // Set field flags. We presume to be first to set this flag.
+    // For this reason, we override any flags other components' prepare methods may have already set.
+    int flags = 0;
+    if (returnFields.wantsScore() || (rb.getRankQuery() instanceof AbstractReRankQuery)) {
+      flags |= SolrIndexSearcher.GET_SCORES;
+    }
+    rb.setFieldFlags( flags );
+
     if (params.getBool(GroupParams.GROUP, false)) {
       prepareGrouping(rb);
     } else {
@@ -865,7 +869,7 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
 
       // Merge the docs via a priority queue so we don't have to sort *all* of the
       // documents... we only need to order the top (rows+start)
-      final ShardFieldSortedHitQueue queue = new ShardFieldSortedHitQueue(sortFields, ss.getOffset() + ss.getCount(), rb.req.getSearcher());
+      final SortedHitQueueManager queueManager = newSortedHitQueueManager(sortFields, rb);
 
       NamedList<Object> shardInfo = null;
       if(rb.req.getParams().getBool(ShardParams.SHARDS_INFO, false)) {
@@ -999,19 +1003,19 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
 
           shardDoc.sortFieldValues = unmarshalledSortFieldValues;
 
-          queue.insertWithOverflow(shardDoc);
+          queueManager.addDocument(shardDoc);
         } // end for-each-doc-in-response
       } // end for-each-response
 
       // The queue now has 0 -> queuesize docs, where queuesize <= start + rows
       // So we want to pop the last documents off the queue to get
       // the docs offset -> queuesize
-      int resultSize = queue.size() - ss.getOffset();
+      int resultSize = queueManager.size() - ss.getOffset();
       resultSize = Math.max(0, resultSize);  // there may not be any docs in range
 
       Map<Object,ShardDoc> resultIds = new HashMap<>();
       for (int i=resultSize-1; i>=0; i--) {
-        ShardDoc shardDoc = queue.pop();
+        ShardDoc shardDoc = queueManager.popDocument();
         shardDoc.positionInResponse = i;
         // Need the toString() for correlation with other lists that must
         // be strings (like keys in highlighting, explain, etc)