SOLR-9574: Factor out AbstractReRankQuery from ReRankQParserPlugin's private ReRankQuery.

This commit is contained in:
Christine Poerschke 2016-09-28 18:49:37 +01:00
parent e8ec99015f
commit 031de301c2
2 changed files with 87 additions and 52 deletions

View File

@ -0,0 +1,83 @@
/*
* 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.Map;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Rescorer;
import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.Weight;
import org.apache.lucene.util.BytesRef;
import org.apache.solr.handler.component.MergeStrategy;
import org.apache.solr.handler.component.QueryElevationComponent;
import org.apache.solr.request.SolrRequestInfo;
public abstract class AbstractReRankQuery extends RankQuery {
protected Query mainQuery;
final protected int reRankDocs;
final protected Rescorer reRankQueryRescorer;
protected Map<BytesRef, Integer> boostedPriority;
public AbstractReRankQuery(Query mainQuery, int reRankDocs, Rescorer reRankQueryRescorer) {
this.mainQuery = mainQuery;
this.reRankDocs = reRankDocs;
this.reRankQueryRescorer = reRankQueryRescorer;
}
public RankQuery wrap(Query _mainQuery) {
if(_mainQuery != null){
this.mainQuery = _mainQuery;
}
return this;
}
public MergeStrategy getMergeStrategy() {
return null;
}
public TopDocsCollector getTopDocsCollector(int len, QueryCommand cmd, IndexSearcher searcher) throws IOException {
if(this.boostedPriority == null) {
SolrRequestInfo info = SolrRequestInfo.getRequestInfo();
if(info != null) {
Map context = info.getReq().getContext();
this.boostedPriority = (Map<BytesRef, Integer>)context.get(QueryElevationComponent.BOOSTED_PRIORITY);
}
}
return new ReRankCollector(reRankDocs, len, reRankQueryRescorer, cmd, searcher, boostedPriority);
}
public Query rewrite(IndexReader reader) throws IOException {
Query q = mainQuery.rewrite(reader);
if (q != mainQuery) {
return rewrite(q);
}
return super.rewrite(reader);
}
protected abstract Query rewrite(Query rewrittenMainQuery) throws IOException;
public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException{
final Weight mainWeight = mainQuery.createWeight(searcher, needsScores);
return new ReRankWeight(mainQuery, reRankQueryRescorer, searcher, mainWeight);
}
}

View File

@ -17,23 +17,13 @@
package org.apache.solr.search;
import java.io.IOException;
import java.util.Map;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.IndexSearcher;
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.TopDocsCollector;
import org.apache.lucene.search.Weight;
import org.apache.lucene.util.BytesRef;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.handler.component.MergeStrategy;
import org.apache.solr.handler.component.QueryElevationComponent;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.SolrRequestInfo;
/*
*
@ -100,13 +90,9 @@ public class ReRankQParserPlugin extends QParserPlugin {
}
}
private final class ReRankQuery extends RankQuery {
private Query mainQuery = defaultQuery;
private final class ReRankQuery extends AbstractReRankQuery {
final private Query reRankQuery;
final private int reRankDocs;
final private double reRankWeight;
final private Rescorer reRankQueryRescorer;
private Map<BytesRef, Integer> boostedPriority;
public int hashCode() {
return 31 * classHash() + mainQuery.hashCode()+reRankQuery.hashCode()+(int)reRankWeight+reRankDocs;
@ -125,34 +111,9 @@ public class ReRankQParserPlugin extends QParserPlugin {
}
public ReRankQuery(Query reRankQuery, int reRankDocs, double reRankWeight) {
super(defaultQuery, reRankDocs, new ReRankQueryRescorer(reRankQuery, reRankWeight));
this.reRankQuery = reRankQuery;
this.reRankDocs = reRankDocs;
this.reRankWeight = reRankWeight;
this.reRankQueryRescorer = new ReRankQueryRescorer(reRankQuery, reRankWeight);
}
public RankQuery wrap(Query _mainQuery) {
if(_mainQuery != null){
this.mainQuery = _mainQuery;
}
return this;
}
public MergeStrategy getMergeStrategy() {
return null;
}
public TopDocsCollector getTopDocsCollector(int len, QueryCommand cmd, IndexSearcher searcher) throws IOException {
if(this.boostedPriority == null) {
SolrRequestInfo info = SolrRequestInfo.getRequestInfo();
if(info != null) {
Map context = info.getReq().getContext();
this.boostedPriority = (Map<BytesRef, Integer>)context.get(QueryElevationComponent.BOOSTED_PRIORITY);
}
}
return new ReRankCollector(reRankDocs, len, reRankQueryRescorer, cmd, searcher, boostedPriority);
}
@Override
@ -166,17 +127,8 @@ public class ReRankQParserPlugin extends QParserPlugin {
return sb.toString();
}
public Query rewrite(IndexReader reader) throws IOException {
Query q = mainQuery.rewrite(reader);
if (q != mainQuery) {
return new ReRankQuery(reRankQuery, reRankDocs, reRankWeight).wrap(q);
}
return super.rewrite(reader);
}
public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException{
final Weight mainWeight = mainQuery.createWeight(searcher, needsScores);
return new ReRankWeight(mainQuery, reRankQueryRescorer, searcher, mainWeight);
protected Query rewrite(Query rewrittenMainQuery) throws IOException {
return new ReRankQuery(reRankQuery, reRankDocs, reRankWeight).wrap(rewrittenMainQuery);
}
}
}