diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt index 70253b88c3b..2f81d0a83df 100644 --- a/solr/CHANGES.txt +++ b/solr/CHANGES.txt @@ -203,6 +203,8 @@ Other Changes * SOLR-9547: Do not allow bin/solr start as root user, unless -force param specified (janhoy) +* SOLR-9567: Make ReRankQParserPlugin's private ReRankCollector a public class of its own. (Christine Poerschke) + ================== 6.2.1 ================== Bug Fixes diff --git a/solr/core/src/java/org/apache/solr/search/ReRankCollector.java b/solr/core/src/java/org/apache/solr/search/ReRankCollector.java new file mode 100644 index 00000000000..1ac1eaff436 --- /dev/null +++ b/solr/core/src/java/org/apache/solr/search/ReRankCollector.java @@ -0,0 +1,175 @@ +/* + * 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.Arrays; +import java.util.Comparator; +import java.util.Map; +import com.carrotsearch.hppc.IntFloatHashMap; +import com.carrotsearch.hppc.IntIntHashMap; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.LeafCollector; +import org.apache.lucene.search.Rescorer; +import org.apache.lucene.search.ScoreDoc; +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 { + + final private TopDocsCollector mainCollector; + final private IndexSearcher searcher; + final private int reRankDocs; + final private int length; + final private Map boostedPriority; + final private Rescorer reRankQueryRescorer; + + + public ReRankCollector(int reRankDocs, + int length, + Rescorer reRankQueryRescorer, + QueryCommand cmd, + IndexSearcher searcher, + Map boostedPriority) throws IOException { + super(null); + this.reRankDocs = reRankDocs; + this.length = length; + this.boostedPriority = boostedPriority; + Sort sort = cmd.getSort(); + if(sort == null) { + this.mainCollector = TopScoreDocCollector.create(Math.max(this.reRankDocs, length)); + } else { + sort = sort.rewrite(searcher); + this.mainCollector = TopFieldCollector.create(sort, Math.max(this.reRankDocs, length), false, true, true); + } + this.searcher = searcher; + this.reRankQueryRescorer = reRankQueryRescorer; + } + + public int getTotalHits() { + return mainCollector.getTotalHits(); + } + + @Override + public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { + return mainCollector.getLeafCollector(context); + } + + @Override + public boolean needsScores() { + return true; + } + + public TopDocs topDocs(int start, int howMany) { + + try { + + TopDocs mainDocs = mainCollector.topDocs(0, Math.max(reRankDocs, length)); + + if(mainDocs.totalHits == 0 || mainDocs.scoreDocs.length == 0) { + return mainDocs; + } + + ScoreDoc[] mainScoreDocs = mainDocs.scoreDocs; + 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); + + //Lower howMany to return if we've collected fewer documents. + howMany = Math.min(howMany, mainScoreDocs.length); + + if(boostedPriority != null) { + SolrRequestInfo info = SolrRequestInfo.getRequestInfo(); + Map requestContext = null; + if(info != null) { + requestContext = info.getReq().getContext(); + } + + IntIntHashMap boostedDocs = QueryElevationComponent.getBoostDocs((SolrIndexSearcher)searcher, boostedPriority, requestContext); + + Arrays.sort(rescoredDocs.scoreDocs, new BoostedComp(boostedDocs, mainDocs.scoreDocs, rescoredDocs.getMaxScore())); + } + + if(howMany == rescoredDocs.scoreDocs.length) { + return rescoredDocs; // Just return the rescoredDocs + } else if(howMany > rescoredDocs.scoreDocs.length) { + //We need to return more then we've reRanked, so create the combined page. + ScoreDoc[] scoreDocs = new ScoreDoc[howMany]; + System.arraycopy(mainScoreDocs, 0, scoreDocs, 0, scoreDocs.length); //lay down the initial docs + System.arraycopy(rescoredDocs.scoreDocs, 0, scoreDocs, 0, rescoredDocs.scoreDocs.length);//overlay the re-ranked docs. + rescoredDocs.scoreDocs = scoreDocs; + return rescoredDocs; + } else { + //We've rescored more then we need to return. + ScoreDoc[] scoreDocs = new ScoreDoc[howMany]; + System.arraycopy(rescoredDocs.scoreDocs, 0, scoreDocs, 0, howMany); + rescoredDocs.scoreDocs = scoreDocs; + return rescoredDocs; + } + } catch (Exception e) { + throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, e); + } + } + + public class BoostedComp implements Comparator { + IntFloatHashMap boostedMap; + + public BoostedComp(IntIntHashMap boostedDocs, ScoreDoc[] scoreDocs, float maxScore) { + this.boostedMap = new IntFloatHashMap(boostedDocs.size()*2); + + for(int i=0; i= 0) { + boostedMap.put(scoreDocs[i].doc, maxScore+boostedDocs.indexGet(idx)); + } else { + break; + } + } + } + + public int compare(Object o1, Object o2) { + ScoreDoc doc1 = (ScoreDoc) o1; + ScoreDoc doc2 = (ScoreDoc) o2; + float score1 = doc1.score; + float score2 = doc2.score; + int idx; + if((idx = boostedMap.indexOf(doc1.doc)) >= 0) { + score1 = boostedMap.indexGet(idx); + } + + if((idx = boostedMap.indexOf(doc2.doc)) >= 0) { + score2 = boostedMap.indexGet(idx); + } + + return -Float.compare(score1, score2); + } + } +} 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 a2301709b9f..fab39d658d2 100644 --- a/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java +++ b/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java @@ -17,26 +17,15 @@ package org.apache.solr.search; import java.io.IOException; -import java.util.Arrays; -import java.util.Comparator; import java.util.Map; -import com.carrotsearch.hppc.IntFloatHashMap; -import com.carrotsearch.hppc.IntIntHashMap; import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.LeafCollector; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.QueryRescorer; import org.apache.lucene.search.Rescorer; -import org.apache.lucene.search.ScoreDoc; -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.search.Weight; import org.apache.lucene.util.BytesRef; import org.apache.solr.common.SolrException; @@ -190,140 +179,4 @@ public class ReRankQParserPlugin extends QParserPlugin { return new ReRankWeight(mainQuery, reRankQueryRescorer, searcher, mainWeight); } } - - private class ReRankCollector extends TopDocsCollector { - - final private TopDocsCollector mainCollector; - final private IndexSearcher searcher; - final private int reRankDocs; - final private int length; - final private Map boostedPriority; - final private Rescorer reRankQueryRescorer; - - - public ReRankCollector(int reRankDocs, - int length, - Rescorer reRankQueryRescorer, - QueryCommand cmd, - IndexSearcher searcher, - Map boostedPriority) throws IOException { - super(null); - this.reRankDocs = reRankDocs; - this.length = length; - this.boostedPriority = boostedPriority; - Sort sort = cmd.getSort(); - if(sort == null) { - this.mainCollector = TopScoreDocCollector.create(Math.max(this.reRankDocs, length)); - } else { - sort = sort.rewrite(searcher); - this.mainCollector = TopFieldCollector.create(sort, Math.max(this.reRankDocs, length), false, true, true); - } - this.searcher = searcher; - this.reRankQueryRescorer = reRankQueryRescorer; - } - - public int getTotalHits() { - return mainCollector.getTotalHits(); - } - - @Override - public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { - return mainCollector.getLeafCollector(context); - } - - @Override - public boolean needsScores() { - return true; - } - - public TopDocs topDocs(int start, int howMany) { - - try { - - TopDocs mainDocs = mainCollector.topDocs(0, Math.max(reRankDocs, length)); - - if(mainDocs.totalHits == 0 || mainDocs.scoreDocs.length == 0) { - return mainDocs; - } - - ScoreDoc[] mainScoreDocs = mainDocs.scoreDocs; - 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); - - //Lower howMany to return if we've collected fewer documents. - howMany = Math.min(howMany, mainScoreDocs.length); - - if(boostedPriority != null) { - SolrRequestInfo info = SolrRequestInfo.getRequestInfo(); - Map requestContext = null; - if(info != null) { - requestContext = info.getReq().getContext(); - } - - IntIntHashMap boostedDocs = QueryElevationComponent.getBoostDocs((SolrIndexSearcher)searcher, boostedPriority, requestContext); - - Arrays.sort(rescoredDocs.scoreDocs, new BoostedComp(boostedDocs, mainDocs.scoreDocs, rescoredDocs.getMaxScore())); - } - - if(howMany == rescoredDocs.scoreDocs.length) { - return rescoredDocs; // Just return the rescoredDocs - } else if(howMany > rescoredDocs.scoreDocs.length) { - //We need to return more then we've reRanked, so create the combined page. - ScoreDoc[] scoreDocs = new ScoreDoc[howMany]; - System.arraycopy(mainScoreDocs, 0, scoreDocs, 0, scoreDocs.length); //lay down the initial docs - System.arraycopy(rescoredDocs.scoreDocs, 0, scoreDocs, 0, rescoredDocs.scoreDocs.length);//overlay the re-ranked docs. - rescoredDocs.scoreDocs = scoreDocs; - return rescoredDocs; - } else { - //We've rescored more then we need to return. - ScoreDoc[] scoreDocs = new ScoreDoc[howMany]; - System.arraycopy(rescoredDocs.scoreDocs, 0, scoreDocs, 0, howMany); - rescoredDocs.scoreDocs = scoreDocs; - return rescoredDocs; - } - } catch (Exception e) { - throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, e); - } - } - - } - - public class BoostedComp implements Comparator { - IntFloatHashMap boostedMap; - - public BoostedComp(IntIntHashMap boostedDocs, ScoreDoc[] scoreDocs, float maxScore) { - this.boostedMap = new IntFloatHashMap(boostedDocs.size()*2); - - for(int i=0; i= 0) { - boostedMap.put(scoreDocs[i].doc, maxScore+boostedDocs.indexGet(idx)); - } else { - break; - } - } - } - - public int compare(Object o1, Object o2) { - ScoreDoc doc1 = (ScoreDoc) o1; - ScoreDoc doc2 = (ScoreDoc) o2; - float score1 = doc1.score; - float score2 = doc2.score; - int idx; - if((idx = boostedMap.indexOf(doc1.doc)) >= 0) { - score1 = boostedMap.indexGet(idx); - } - - if((idx = boostedMap.indexOf(doc2.doc)) >= 0) { - score2 = boostedMap.indexGet(idx); - } - - return -Float.compare(score1, score2); - } - } }