scoredoc2
diff --git a/solr/core/src/java/org/apache/solr/search/QueryRescorer.java b/solr/core/src/java/org/apache/solr/search/QueryRescorer.java
new file mode 100644
index 0000000..3502bf4
--- /dev/null
+++ b/solr/core/src/java/org/apache/solr/search/QueryRescorer.java
@@ -0,0 +1,207 @@
+/*
+ * 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;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.*;
+import org.apache.lucene.util.ArrayUtil;
+
+/**
+ * A {@link Rescorer} that uses a provided Query to assign scores to the first-pass hits.
+ *
+ * @lucene.experimental
+ */
+public abstract class QueryRescorer extends Rescorer {
+
+ private final Query query;
+
+ /** Sole constructor, passing the 2nd pass query to assign scores to the 1st pass hits. */
+ public QueryRescorer(Query query) {
+ this.query = query;
+ }
+
+ /**
+ * Implement this in a subclass to combine the first pass and second pass scores. If
+ * secondPassMatches is false then the second pass query failed to match a hit from the first pass
+ * query, and you should ignore the secondPassScore.
+ */
+ protected abstract float combine(
+ float firstPassScore, boolean secondPassMatches, float secondPassScore);
+
+ @Override
+ public TopDocs rescore(IndexSearcher searcher, TopDocs firstPassTopDocs, int topN)
+ throws IOException {
+
+ // convert Lucene ScoreDoc list to Solr ScoreDoc2 list (the latter supports index attribute)
+ // and record original order in each hit's index attribute
+ // NB code below relies on hits being in docId order so the order is about to be lost
+ ArrayList<ScoreDoc2> solrHits = new ArrayList<ScoreDoc2>();
+ for(int i=0; i<firstPassTopDocs.scoreDocs.length; i++) {
+ ScoreDoc2 hit = (ScoreDoc2) firstPassTopDocs.scoreDocs[i];
+ hit.index=i;
+ solrHits.add(hit);
+ }
+
+ ScoreDoc2[] hits = new ScoreDoc2[solrHits.size()];
+ hits = solrHits.toArray(hits);
+
+ Arrays.sort(
+ hits,
+ new Comparator<ScoreDoc2>() {
+ @Override
+ public int compare(ScoreDoc2 a, ScoreDoc2 b) {
+ return a.doc - b.doc;
+ }
+ });
+
+ List<LeafReaderContext> leaves = searcher.getIndexReader().leaves();
+
+ Query rewritten = searcher.rewrite(query);
+ Weight weight = searcher.createWeight(rewritten, ScoreMode.COMPLETE, 1);
+
+ // Now merge sort docIDs from hits, with reader's leaves:
+ int hitUpto = 0;
+ int readerUpto = -1;
+ int endDoc = 0;
+ int docBase = 0;
+ Scorer scorer = null;
+
+ while (hitUpto < hits.length) {
+ ScoreDoc2 hit = hits[hitUpto];
+ int docID = hit.doc;
+ LeafReaderContext readerContext = null;
+ while (docID >= endDoc) {
+ readerUpto++;
+ readerContext = leaves.get(readerUpto);
+ endDoc = readerContext.docBase + readerContext.reader().maxDoc();
+ }
+
+ if (readerContext != null) {
+ // We advanced to another segment:
+ docBase = readerContext.docBase;
+ scorer = weight.scorer(readerContext);
+ }
+
+ if (scorer != null) {
+ int targetDoc = docID - docBase;
+ int actualDoc = scorer.docID();
+ if (actualDoc < targetDoc) {
+ actualDoc = scorer.iterator().advance(targetDoc);
+ }
+
+ if (actualDoc == targetDoc) {
+ // Query did match this doc:
+ hit.score = combine(hit.score, true, scorer.score());
+ } else {
+ // Query did not match this doc:
+ assert actualDoc > targetDoc;
+ hit.score = combine(hit.score, false, 0.0f);
+ }
+ } else {
+ // Query did not match this doc:
+ hit.score = combine(hit.score, false, 0.0f);
+ }
+
+ hitUpto++;
+ }
+
+ Comparator<ScoreDoc2> sortDocComparator =
+ new Comparator<ScoreDoc2>() {
+ @Override
+ public int compare(ScoreDoc2 a, ScoreDoc2 b) {
+ // Sort by score descending, then original order as
+ // recorded in index attribute above
+ if (a.score > b.score) {
+ return -1;
+ } else if (a.score < b.score) {
+ return 1;
+ } else if (a.index > b.index) {
+ return 1;
+ } else {
+ return -1;
+ }
+ }
+ };
+
+ if (topN < hits.length) {
+ ArrayUtil.select(hits, 0, hits.length, topN, sortDocComparator);
+ ScoreDoc2[] subset = new ScoreDoc2[topN];
+ System.arraycopy(hits, 0, subset, 0, topN);
+ hits = subset;
+ }
+
+ Arrays.sort(hits, sortDocComparator);
+
+ return new TopDocs(firstPassTopDocs.totalHits, hits);
+ }
+
+ @Override
+ public Explanation explain(IndexSearcher searcher, Explanation firstPassExplanation, int docID)
+ throws IOException {
+ Explanation secondPassExplanation = searcher.explain(query, docID);
+
+ Number secondPassScore =
+ secondPassExplanation.isMatch() ? secondPassExplanation.getValue() : null;
+
+ float score;
+ if (secondPassScore == null) {
+ score = combine(firstPassExplanation.getValue().floatValue(), false, 0.0f);
+ } else {
+ score =
+ combine(firstPassExplanation.getValue().floatValue(), true, secondPassScore.floatValue());
+ }
+
+ Explanation first =
+ Explanation.match(
+ firstPassExplanation.getValue(), "first pass score", firstPassExplanation);
+
+ Explanation second;
+ if (secondPassScore == null) {
+ second = Explanation.noMatch("no second pass score");
+ } else {
+ second = Explanation.match(secondPassScore, "second pass score", secondPassExplanation);
+ }
+
+ return Explanation.match(
+ score, "combined first and second pass score using " + getClass(), first, second);
+ }
+
+ /**
+ * Sugar API, calling {#rescore} using a simple linear combination of firstPassScore + weight *
+ * secondPassScore
+ */
+ public static TopDocs rescore(
+ IndexSearcher searcher, TopDocs topDocs, Query query, final double weight, int topN)
+ throws IOException {
+ return new QueryRescorer(query) {
+ @Override
+ protected float combine(
+ float firstPassScore, boolean secondPassMatches, float secondPassScore) {
+ float score = firstPassScore;
+ if (secondPassMatches) {
+ score += weight * secondPassScore;
+ }
+ return score;
+ }
+ }.rescore(searcher, topDocs, topN);
+ }
+}
diff --git a/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java b/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java
index 5bf5ce4..5294cf7 100644
--- a/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java
+++ b/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java
@@ -20,7 +20,6 @@
import java.lang.invoke.MethodHandles;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
-import org.apache.lucene.search.QueryRescorer;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.SolrParams;
diff --git a/solr/core/src/java/org/apache/solr/search/ReRankScaler.java b/solr/core/src/java/org/apache/solr/search/ReRankScaler.java
index 8410467..319cb48 100644
--- a/solr/core/src/java/org/apache/solr/search/ReRankScaler.java
+++ b/solr/core/src/java/org/apache/solr/search/ReRankScaler.java
@@ -23,7 +23,6 @@
import java.util.Map;
import java.util.Set;
import org.apache.lucene.search.Explanation;
-import org.apache.lucene.search.QueryRescorer;
import org.apache.lucene.search.ScoreDoc;
public class ReRankScaler {
diff --git a/solr/core/src/java/org/apache/solr/search/ScoreDoc2.java b/solr/core/src/java/org/apache/solr/search/ScoreDoc2.java
new file mode 100644
index 0000000..97d258b
--- /dev/null
+++ b/solr/core/src/java/org/apache/solr/search/ScoreDoc2.java
@@ -0,0 +1,39 @@
+/*
+ * 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 ScoreDoc2 extends org.apache.lucene.search.ScoreDoc {
+
+ public ScoreDoc2(int doc, float score) {
+ super(doc, score);
+ }
+
+ /**
+ * Original index of the doc in the result set.
+ *
+ * Only set/used by {@link QueryRescorer#rescore} and Solr's LTRRescorer.
+ *
+ * Could shardIndex be used for this instead?
+ */
+ public int index = -1;
+
+ // A convenience method for debugging.
+ @Override
+ public String toString() {
+ return "doc=" + doc + " score=" + score + " shardIndex=" + shardIndex + " index=" + index;
+ }
+}
