Fix/solr 17018 branch 9 1 (#2358)
* SOLR-17018: Query limit timeout is now supported by Learning To Rank rescoring.
diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt
index 04697b7..f9e8446 100644
--- a/solr/CHANGES.txt
+++ b/solr/CHANGES.txt
@@ -7,7 +7,8 @@
================== 9.1.2 ==================
Bug Fixes
---------------------
-(No changes)
+
+* SOLR-17018: Query limit timeout is now supported by Learning To Rank rescoring. (Alessandro Benedetti)
================== 9.1.1 ==================
diff --git a/solr/core/src/java/org/apache/solr/search/IncompleteRerankingException.java b/solr/core/src/java/org/apache/solr/search/IncompleteRerankingException.java
new file mode 100644
index 0000000..c3f7190
--- /dev/null
+++ b/solr/core/src/java/org/apache/solr/search/IncompleteRerankingException.java
@@ -0,0 +1,24 @@
+/*
+ * 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;
+
+public class IncompleteRerankingException extends RuntimeException {
+
+ public IncompleteRerankingException() {
+ super();
+ }
+}
diff --git a/solr/core/src/java/org/apache/solr/search/ReRankCollector.java b/solr/core/src/java/org/apache/solr/search/ReRankCollector.java
index c4bb7c1..2cfc35b 100644
--- a/solr/core/src/java/org/apache/solr/search/ReRankCollector.java
+++ b/solr/core/src/java/org/apache/solr/search/ReRankCollector.java
@@ -23,6 +23,7 @@
import java.util.Comparator;
import java.util.Map;
import java.util.Set;
+import org.apache.lucene.index.ExitableDirectoryReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafCollector;
@@ -39,6 +40,7 @@
import org.apache.solr.common.SolrException;
import org.apache.solr.handler.component.QueryElevationComponent;
import org.apache.solr.request.SolrRequestInfo;
+import org.apache.solr.response.SolrQueryResponse;
/* A TopDocsCollector used by reranking queries. */
public class ReRankCollector extends TopDocsCollector<ScoreDoc> {
@@ -109,13 +111,25 @@
}
ScoreDoc[] mainScoreDocs = mainDocs.scoreDocs;
+ ScoreDoc[] mainScoreDocsClone = deepClone(mainScoreDocs);
ScoreDoc[] reRankScoreDocs = new ScoreDoc[Math.min(mainScoreDocs.length, reRankDocs)];
System.arraycopy(mainScoreDocs, 0, reRankScoreDocs, 0, reRankScoreDocs.length);
mainDocs.scoreDocs = reRankScoreDocs;
- TopDocs rescoredDocs =
- reRankQueryRescorer.rescore(searcher, mainDocs, mainDocs.scoreDocs.length);
+ TopDocs rescoredDocs;
+ try {
+ rescoredDocs = reRankQueryRescorer.rescore(searcher, mainDocs, mainDocs.scoreDocs.length);
+ } catch (IncompleteRerankingException | ExitableDirectoryReader.ExitingReaderException ex) {
+ mainDocs.scoreDocs = mainScoreDocsClone;
+ rescoredDocs = mainDocs;
+ SolrRequestInfo.getRequestInfo()
+ .getRsp()
+ .getResponseHeader()
+ .asShallowMap()
+ .put(SolrQueryResponse.RESPONSE_HEADER_PARTIAL_RESULTS_KEY, Boolean.TRUE);
+ SolrQueryTimeoutImpl.reset();
+ }
// Lower howMany to return if we've collected fewer documents.
howMany = Math.min(howMany, mainScoreDocs.length);
@@ -164,6 +178,17 @@
}
}
+ private ScoreDoc[] deepClone(ScoreDoc[] scoreDocs) {
+ ScoreDoc[] scoreDocs1 = new ScoreDoc[scoreDocs.length];
+ for (int i = 0; i < scoreDocs.length; i++) {
+ ScoreDoc scoreDoc = scoreDocs[i];
+ if (scoreDoc != null) {
+ scoreDocs1[i] = new ScoreDoc(scoreDoc.doc, scoreDoc.score);
+ }
+ }
+ return scoreDocs1;
+ }
+
public static class BoostedComp implements Comparator<ScoreDoc> {
IntFloatHashMap boostedMap;
diff --git a/solr/modules/ltr/src/java/org/apache/solr/ltr/LTRRescorer.java b/solr/modules/ltr/src/java/org/apache/solr/ltr/LTRRescorer.java
index 6407034..ccee096 100644
--- a/solr/modules/ltr/src/java/org/apache/solr/ltr/LTRRescorer.java
+++ b/solr/modules/ltr/src/java/org/apache/solr/ltr/LTRRescorer.java
@@ -31,7 +31,9 @@
import org.apache.lucene.search.TotalHits;
import org.apache.lucene.search.Weight;
import org.apache.solr.ltr.interleaving.OriginalRankingLTRScoringQuery;
+import org.apache.solr.search.IncompleteRerankingException;
import org.apache.solr.search.SolrIndexSearcher;
+import org.apache.solr.search.SolrQueryTimeoutImpl;
/**
* Implements the rescoring logic. The top documents returned by solr with their original scores,
@@ -234,6 +236,12 @@
scorer.getDocInfo().setOriginalDocScore(hit.score);
hit.score = scorer.score();
+
+ if (SolrQueryTimeoutImpl.getInstance().isTimeoutEnabled()
+ && SolrQueryTimeoutImpl.getInstance().shouldExit()) {
+ throw new IncompleteRerankingException();
+ }
+
if (hitUpto < topN) {
reranked[hitUpto] = hit;
// if the heap is not full, maybe I want to log the features for this
diff --git a/solr/modules/ltr/src/test-files/featureExamples/features-slow.json b/solr/modules/ltr/src/test-files/featureExamples/features-slow.json
new file mode 100644
index 0000000..a60c47d
--- /dev/null
+++ b/solr/modules/ltr/src/test-files/featureExamples/features-slow.json
@@ -0,0 +1,7 @@
+[
+ {
+ "name" : "slow",
+ "class" : "org.apache.solr.ltr.feature.SolrFeature",
+ "params" : { "q" : "{!func}sleep(1000,999)" }
+ }
+]
diff --git a/solr/modules/ltr/src/test-files/modelExamples/linear-slow-model.json b/solr/modules/ltr/src/test-files/modelExamples/linear-slow-model.json
new file mode 100644
index 0000000..824b9c4
--- /dev/null
+++ b/solr/modules/ltr/src/test-files/modelExamples/linear-slow-model.json
@@ -0,0 +1,14 @@
+{
+ "class": "org.apache.solr.ltr.model.LinearModel",
+ "name": "slowModel",
+ "features": [
+ {
+ "name": "slow"
+ }
+ ],
+ "params": {
+ "weights": {
+ "slow": 1
+ }
+ }
+}
diff --git a/solr/modules/ltr/src/test/org/apache/solr/ltr/TestLTRQParserPlugin.java b/solr/modules/ltr/src/test/org/apache/solr/ltr/TestLTRQParserPlugin.java
index 1bdca68..976725a 100644
--- a/solr/modules/ltr/src/test/org/apache/solr/ltr/TestLTRQParserPlugin.java
+++ b/solr/modules/ltr/src/test/org/apache/solr/ltr/TestLTRQParserPlugin.java
@@ -29,6 +29,9 @@
loadFeatures("features-linear.json");
loadModels("linear-model.json");
+
+ loadFeatures("features-slow.json");
+ loadModels("linear-slow-model.json"); // just a linear model with one feature
}
@AfterClass
@@ -137,4 +140,63 @@
query.add("rq", "{!ltr reRankDocs=3 model=6029760550880411648}");
assertJQ("/query" + query.toQueryString(), "/response/numFound/==0");
}
+
+ @Test
+ public void ltr_expensiveFeatureRescoring_shouldTimeOutAndReturnPartialResults()
+ throws Exception {
+ /* One SolrFeature is defined: {!func}sleep(1000,999)
+ * It simulates a slow feature extraction, sleeping for 1000ms and returning 999 as a score when finished
+ * */
+
+ final String solrQuery = "_query_:{!edismax qf='id' v='8^=10 9^=5 7^=3 6^=1'}";
+ final SolrQuery query = new SolrQuery();
+ query.setQuery(solrQuery);
+ query.add("fl", "*, score");
+ query.add("rows", "4");
+ query.add("fv", "true");
+ query.add("rq", "{!ltr model=slowModel reRankDocs=3}");
+ query.add("timeAllowed", "300");
+
+ assertJQ(
+ "/query" + query.toQueryString(),
+ "/response/numFound/==4",
+ "/responseHeader/partialResults/==true",
+ "/response/docs/[0]/id=='8'",
+ "/response/docs/[0]/score==10.0",
+ "/response/docs/[1]/id=='9'",
+ "/response/docs/[1]/score==5.0",
+ "/response/docs/[2]/id=='7'",
+ "/response/docs/[2]/score==3.0",
+ "/response/docs/[3]/id=='6'",
+ "/response/docs/[3]/score==1.0");
+ }
+
+ @Test
+ public void ltr_expensiveFeatureRescoringWithinTimeAllowed_shouldReturnRerankedResults()
+ throws Exception {
+ /* One SolrFeature is defined: {!func}sleep(1000,999)
+ * It simulates a slow feature extraction, sleeping for 1000ms and returning 999 as a score when finished
+ * */
+
+ final String solrQuery = "_query_:{!edismax qf='id' v='8^=10 9^=5 7^=3 6^=1'}";
+ final SolrQuery query = new SolrQuery();
+ query.setQuery(solrQuery);
+ query.add("fl", "*, score");
+ query.add("rows", "4");
+ query.add("fv", "true");
+ query.add("rq", "{!ltr model=slowModel reRankDocs=3}");
+ query.add("timeAllowed", "2000");
+
+ assertJQ(
+ "/query" + query.toQueryString(),
+ "/response/numFound/==4",
+ "/response/docs/[0]/id=='7'",
+ "/response/docs/[0]/score==999.0",
+ "/response/docs/[1]/id=='8'",
+ "/response/docs/[1]/score==999.0",
+ "/response/docs/[2]/id=='9'",
+ "/response/docs/[2]/score==999.0",
+ "/response/docs/[3]/id=='6'",
+ "/response/docs/[3]/score==1.0");
+ }
}
