org.apache.solr.search.ReRankCollector Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of solr-core Show documentation
Show all versions of solr-core Show documentation
Apache Solr (module: core)
/*
* 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 com.carrotsearch.hppc.IntFloatHashMap;
import com.carrotsearch.hppc.IntIntHashMap;
import java.io.IOException;
import java.util.Arrays;
import java.util.Comparator;
import java.util.Map;
import java.util.Set;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Rescorer;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.TopFieldCollector;
import org.apache.lucene.search.TopScoreDocCollector;
import org.apache.lucene.util.BytesRef;
import org.apache.solr.common.SolrException;
import org.apache.solr.handler.component.QueryElevationComponent;
import org.apache.solr.request.SolrRequestInfo;
/* A TopDocsCollector used by reranking queries. */
public class ReRankCollector extends TopDocsCollector {
private final TopDocsCollector mainCollector;
private final IndexSearcher searcher;
private final int reRankDocs;
private final int length;
private final Set boostedPriority; // order is the "priority"
private final Rescorer reRankQueryRescorer;
private final Sort sort;
private final Query query;
private ReRankScaler reRankScaler;
private ReRankOperator reRankOperator;
public ReRankCollector(
int reRankDocs,
int length,
Rescorer reRankQueryRescorer,
QueryCommand cmd,
IndexSearcher searcher,
Set boostedPriority,
ReRankScaler reRankScaler,
ReRankOperator reRankOperator)
throws IOException {
this(reRankDocs, length, reRankQueryRescorer, cmd, searcher, boostedPriority);
this.reRankScaler = reRankScaler;
this.reRankOperator = reRankOperator;
}
public ReRankCollector(
int reRankDocs,
int length,
Rescorer reRankQueryRescorer,
QueryCommand cmd,
IndexSearcher searcher,
Set boostedPriority)
throws IOException {
super(null);
this.reRankDocs = reRankDocs;
this.length = length;
this.boostedPriority = boostedPriority;
this.query = cmd.getQuery();
Sort sort = cmd.getSort();
if (sort == null) {
this.sort = null;
this.mainCollector =
TopScoreDocCollector.create(Math.max(this.reRankDocs, length), cmd.getMinExactCount());
} else {
this.sort = sort = sort.rewrite(searcher);
// scores are needed for Rescorer (regardless of whether sort needs it)
this.mainCollector =
TopFieldCollector.create(sort, Math.max(this.reRankDocs, length), cmd.getMinExactCount());
}
this.searcher = searcher;
this.reRankQueryRescorer = reRankQueryRescorer;
}
@Override
public int getTotalHits() {
return mainCollector.getTotalHits();
}
@Override
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
return mainCollector.getLeafCollector(context);
}
@Override
public ScoreMode scoreMode() {
return this.mainCollector.scoreMode();
}
@Override
public TopDocs topDocs(int start, int howMany) {
try {
TopDocs mainDocs = mainCollector.topDocs(0, Math.max(reRankDocs, length));
if (mainDocs.totalHits.value == 0 || mainDocs.scoreDocs.length == 0) {
return mainDocs;
}
if (sort != null) {
TopFieldCollector.populateScores(mainDocs.scoreDocs, searcher, query);
}
ScoreDoc[] mainScoreDocs = mainDocs.scoreDocs;
boolean zeroOutScores = reRankScaler != null && reRankScaler.scaleScores();
ScoreDoc[] mainScoreDocsClone = deepClone(mainScoreDocs, zeroOutScores);
ScoreDoc[] reRankScoreDocs = new ScoreDoc[Math.min(mainScoreDocs.length, reRankDocs)];
System.arraycopy(mainScoreDocs, 0, reRankScoreDocs, 0, reRankScoreDocs.length);
mainDocs.scoreDocs = reRankScoreDocs;
// If we're scaling scores use the replace rescorer because we just want the re-rank score.
TopDocs rescoredDocs;
try {
rescoredDocs =
zeroOutScores // previously zero-ed out scores are to be replaced
? reRankScaler
.getReplaceRescorer()
.rescore(searcher, mainDocs, mainDocs.scoreDocs.length)
: reRankQueryRescorer.rescore(searcher, mainDocs, mainDocs.scoreDocs.length);
} catch (IncompleteRerankingException ex) {
mainDocs.scoreDocs = mainScoreDocsClone;
rescoredDocs = mainDocs;
}
// Lower howMany to return if we've collected fewer documents.
howMany = Math.min(howMany, mainScoreDocs.length);
if (boostedPriority != null) {
SolrRequestInfo info = SolrRequestInfo.getRequestInfo();
Map