package org.apache.lucene.queries;
/*
* 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.
*/
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.Set;
import java.util.Arrays;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.queries.function.FunctionQuery;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.search.ComplexExplanation;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Weight;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.ToStringUtils;
/**
* Query that sets document score as a programmatic function of several (sub) scores:
* <ol>
* <li>the score of its subQuery (any query)</li>
* <li>(optional) the score of its {@link FunctionQuery} (or queries).</li>
* </ol>
* Subclasses can modify the computation by overriding {@link #getCustomScoreProvider}.
*
* @lucene.experimental
*/
public class CustomScoreQuery extends Query {
private Query subQuery;
private Query[] scoringQueries; // never null (empty array if there are no valSrcQueries).
private boolean strict = false; // if true, valueSource part of query does not take part in weights normalization.
/**
* Create a CustomScoreQuery over input subQuery.
* @param subQuery the sub query whose scored is being customized. Must not be null.
*/
public CustomScoreQuery(Query subQuery) {
this(subQuery, new FunctionQuery[0]);
}
/**
* Create a CustomScoreQuery over input subQuery and a {@link org.apache.lucene.queries.function.FunctionQuery}.
* @param subQuery the sub query whose score is being customized. Must not be null.
* @param scoringQuery a value source query whose scores are used in the custom score
* computation. This parameter is optional - it can be null.
*/
public CustomScoreQuery(Query subQuery, FunctionQuery scoringQuery) {
this(subQuery, scoringQuery!=null ? // don't want an array that contains a single null..
new FunctionQuery[] {scoringQuery} : new FunctionQuery[0]);
}
/**
* Create a CustomScoreQuery over input subQuery and a {@link org.apache.lucene.queries.function.FunctionQuery}.
* @param subQuery the sub query whose score is being customized. Must not be null.
* @param scoringQueries value source queries whose scores are used in the custom score
* computation. This parameter is optional - it can be null or even an empty array.
*/
public CustomScoreQuery(Query subQuery, FunctionQuery... scoringQueries) {
this.subQuery = subQuery;
this.scoringQueries = scoringQueries !=null?
scoringQueries : new Query[0];
if (subQuery == null) throw new IllegalArgumentException("<subquery> must not be null!");
}
/*(non-Javadoc) @see org.apache.lucene.search.Query#rewrite(org.apache.lucene.index.IndexReader) */
@Override
public Query rewrite(IndexReader reader) throws IOException {
CustomScoreQuery clone = null;
final Query sq = subQuery.rewrite(reader);
if (sq != subQuery) {
clone = clone();
clone.subQuery = sq;
}
for(int i = 0; i < scoringQueries.length; i++) {
final Query v = scoringQueries[i].rewrite(reader);
if (v != scoringQueries[i]) {
if (clone == null) clone = clone();
clone.scoringQueries[i] = v;
}
}
return (clone == null) ? this : clone;
}
/*(non-Javadoc) @see org.apache.lucene.search.Query#extractTerms(java.util.Set) */
@Override
public void extractTerms(Set<Term> terms) {
subQuery.extractTerms(terms);
for (Query scoringQuery : scoringQueries) {
scoringQuery.extractTerms(terms);
}
}
/*(non-Javadoc) @see org.apache.lucene.search.Query#clone() */
@Override
public CustomScoreQuery clone() {
CustomScoreQuery clone = (CustomScoreQuery)super.clone();
clone.subQuery = subQuery.clone();
clone.scoringQueries = new Query[scoringQueries.length];
for(int i = 0; i < scoringQueries.length; i++) {
clone.scoringQueries[i] = scoringQueries[i].clone();
}
return clone;
}
/* (non-Javadoc) @see org.apache.lucene.search.Query#toString(java.lang.String) */
@Override
public String toString(String field) {
StringBuilder sb = new StringBuilder(name()).append("(");
sb.append(subQuery.toString(field));
for (Query scoringQuery : scoringQueries) {
sb.append(", ").append(scoringQuery.toString(field));
}
sb.append(")");
sb.append(strict?" STRICT" : "");
return sb.toString() + ToStringUtils.boost(getBoost());
}
/** Returns true if <code>o</code> is equal to this. */
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!super.equals(o))
return false;
if (getClass() != o.getClass()) {
return false;
}
CustomScoreQuery other = (CustomScoreQuery)o;
if (this.getBoost() != other.getBoost() ||
!this.subQuery.equals(other.subQuery) ||
this.strict != other.strict ||
this.scoringQueries.length != other.scoringQueries.length) {
return false;
}
return Arrays.equals(scoringQueries, other.scoringQueries);
}
/** Returns a hash code value for this object. */
@Override
public int hashCode() {
return (getClass().hashCode() + subQuery.hashCode() + Arrays.hashCode(scoringQueries))
^ Float.floatToIntBits(getBoost()) ^ (strict ? 1234 : 4321);
}
/**
* Returns a {@link CustomScoreProvider} that calculates the custom scores
* for the given {@link IndexReader}. The default implementation returns a default
* implementation as specified in the docs of {@link CustomScoreProvider}.
* @since 2.9.2
*/
protected CustomScoreProvider getCustomScoreProvider(AtomicReaderContext context) throws IOException {
return new CustomScoreProvider(context);
}
//=========================== W E I G H T ============================
private class CustomWeight extends Weight {
Weight subQueryWeight;
Weight[] valSrcWeights;
boolean qStrict;
float queryWeight;
public CustomWeight(IndexSearcher searcher) throws IOException {
this.subQueryWeight = subQuery.createWeight(searcher);
this.valSrcWeights = new Weight[scoringQueries.length];
for(int i = 0; i < scoringQueries.length; i++) {
this.valSrcWeights[i] = scoringQueries[i].createWeight(searcher);
}
this.qStrict = strict;
}
/*(non-Javadoc) @see org.apache.lucene.search.Weight#getQuery() */
@Override
public Query getQuery() {
return CustomScoreQuery.this;
}
@Override
public float getValueForNormalization() throws IOException {
float sum = subQueryWeight.getValueForNormalization();
for (Weight valSrcWeight : valSrcWeights) {
if (qStrict) {
valSrcWeight.getValueForNormalization(); // do not include ValueSource part in the query normalization
} else {
sum += valSrcWeight.getValueForNormalization();
}
}
return sum;
}
/*(non-Javadoc) @see org.apache.lucene.search.Weight#normalize(float) */
@Override
public void normalize(float norm, float topLevelBoost) {
// note we DONT incorporate our boost, nor pass down any topLevelBoost
// (e.g. from outer BQ), as there is no guarantee that the CustomScoreProvider's
// function obeys the distributive law... it might call sqrt() on the subQuery score
// or some other arbitrary function other than multiplication.
// so, instead boosts are applied directly in score()
subQueryWeight.normalize(norm, 1f);
for (Weight valSrcWeight : valSrcWeights) {
if (qStrict) {
valSrcWeight.normalize(1, 1); // do not normalize the ValueSource part
} else {
valSrcWeight.normalize(norm, 1f);
}
}
queryWeight = topLevelBoost * getBoost();
}
@Override
public Scorer scorer(AtomicReaderContext context, Bits acceptDocs) throws IOException {
Scorer subQueryScorer = subQueryWeight.scorer(context, acceptDocs);
if (subQueryScorer == null) {
return null;
}
Scorer[] valSrcScorers = new Scorer[valSrcWeights.length];
for(int i = 0; i < valSrcScorers.length; i++) {
valSrcScorers[i] = valSrcWeights[i].scorer(context, acceptDocs);
}
return new CustomScorer(CustomScoreQuery.this.getCustomScoreProvider(context), this, queryWeight, subQueryScorer, valSrcScorers);
}
@Override
public Explanation explain(AtomicReaderContext context, int doc) throws IOException {
Explanation explain = doExplain(context, doc);
return explain == null ? new Explanation(0.0f, "no matching docs") : explain;
}
private Explanation doExplain(AtomicReaderContext info, int doc) throws IOException {
Explanation subQueryExpl = subQueryWeight.explain(info, doc);
if (!subQueryExpl.isMatch()) {
return subQueryExpl;
}
// match
Explanation[] valSrcExpls = new Explanation[valSrcWeights.length];
for(int i = 0; i < valSrcWeights.length; i++) {
valSrcExpls[i] = valSrcWeights[i].explain(info, doc);
}
Explanation customExp = CustomScoreQuery.this.getCustomScoreProvider(info).customExplain(doc,subQueryExpl,valSrcExpls);
float sc = getBoost() * customExp.getValue();
Explanation res = new ComplexExplanation(
true, sc, CustomScoreQuery.this.toString() + ", product of:");
res.addDetail(customExp);
res.addDetail(new Explanation(getBoost(), "queryBoost")); // actually using the q boost as q weight (== weight value)
return res;
}
@Override
public boolean scoresDocsOutOfOrder() {
return false;
}
}
//=========================== S C O R E R ============================
/**
* A scorer that applies a (callback) function on scores of the subQuery.
*/
private class CustomScorer extends Scorer {
private final float qWeight;
private final Scorer subQueryScorer;
private final Scorer[] valSrcScorers;
private final CustomScoreProvider provider;
private final float[] vScores; // reused in score() to avoid allocating this array for each doc
// constructor
private CustomScorer(CustomScoreProvider provider, CustomWeight w, float qWeight,
Scorer subQueryScorer, Scorer[] valSrcScorers) {
super(w);
this.qWeight = qWeight;
this.subQueryScorer = subQueryScorer;
this.valSrcScorers = valSrcScorers;
this.vScores = new float[valSrcScorers.length];
this.provider = provider;
}
@Override
public int nextDoc() throws IOException {
int doc = subQueryScorer.nextDoc();
if (doc != NO_MORE_DOCS) {
for (Scorer valSrcScorer : valSrcScorers) {
valSrcScorer.advance(doc);
}
}
return doc;
}
@Override
public int docID() {
return subQueryScorer.docID();
}
/*(non-Javadoc) @see org.apache.lucene.search.Scorer#score() */
@Override
public float score() throws IOException {
for (int i = 0; i < valSrcScorers.length; i++) {
vScores[i] = valSrcScorers[i].score();
}
return qWeight * provider.customScore(subQueryScorer.docID(), subQueryScorer.score(), vScores);
}
@Override
public int freq() throws IOException {
return subQueryScorer.freq();
}
@Override
public Collection<ChildScorer> getChildren() {
return Collections.singleton(new ChildScorer(subQueryScorer, "CUSTOM"));
}
@Override
public int advance(int target) throws IOException {
int doc = subQueryScorer.advance(target);
if (doc != NO_MORE_DOCS) {
for (Scorer valSrcScorer : valSrcScorers) {
valSrcScorer.advance(doc);
}
}
return doc;
}
@Override
public long cost() {
return subQueryScorer.cost();
}
}
@Override
public Weight createWeight(IndexSearcher searcher) throws IOException {
return new CustomWeight(searcher);
}
/**
* Checks if this is strict custom scoring.
* In strict custom scoring, the {@link ValueSource} part does not participate in weight normalization.
* This may be useful when one wants full control over how scores are modified, and does
* not care about normalizing by the {@link ValueSource} part.
* One particular case where this is useful if for testing this query.
* <P>
* Note: only has effect when the {@link ValueSource} part is not null.
*/
public boolean isStrict() {
return strict;
}
/**
* Set the strict mode of this query.
* @param strict The strict mode to set.
* @see #isStrict()
*/
public void setStrict(boolean strict) {
this.strict = strict;
}
/** The sub-query that CustomScoreQuery wraps, affecting both the score and which documents match. */
public Query getSubQuery() {
return subQuery;
}
/** The scoring queries that only affect the score of CustomScoreQuery. */
public Query[] getScoringQueries() {
return scoringQueries;
}
/**
* A short name of this query, used in {@link #toString(String)}.
*/
public String name() {
return "custom";
}
}