lucene-java-3.4.0/lucene/src/java/org/apache/lucene/search/function/CustomScoreQuery.java

   1 package org.apache.lucene.search.function;
   2
   3 /**
   4  * Licensed to the Apache Software Foundation (ASF) under one or more
   5  * contributor license agreements.  See the NOTICE file distributed with
   6  * this work for additional information regarding copyright ownership.
   7  * The ASF licenses this file to You under the Apache License, Version 2.0
   8  * (the "License"); you may not use this file except in compliance with
   9  * the License.  You may obtain a copy of the License at
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  13  * Unless required by applicable law or agreed to in writing, software
  14  * distributed under the License is distributed on an "AS IS" BASIS,
  15  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16  * See the License for the specific language governing permissions and
  17  * limitations under the License.
  18  */
  19
  20 import java.io.IOException;
  21 import java.util.Set;
  22 import java.util.Arrays;
  23
  24 import org.apache.lucene.index.IndexReader;
  25 import org.apache.lucene.index.Term;
  26 import org.apache.lucene.search.ComplexExplanation;
  27 import org.apache.lucene.search.Explanation;
  28 import org.apache.lucene.search.Query;
  29 import org.apache.lucene.search.Weight;
  30 import org.apache.lucene.search.Scorer;
  31 import org.apache.lucene.search.Searcher;
  32 import org.apache.lucene.search.Similarity;
  33 import org.apache.lucene.util.ToStringUtils;
  34
  35 /**
  36  * Query that sets document score as a programmatic function of several (sub) scores:
  37  * <ol>
  38  *    <li>the score of its subQuery (any query)</li>
  39  *    <li>(optional) the score of its ValueSourceQuery (or queries).
  40  *        For most simple/convenient use cases this query is likely to be a
  41  *        {@link org.apache.lucene.search.function.FieldScoreQuery FieldScoreQuery}</li>
  42  * </ol>
  43  * Subclasses can modify the computation by overriding {@link #getCustomScoreProvider}.
  44  *
  45  * @lucene.experimental
  46  */
  47 public class CustomScoreQuery extends Query {
  48
  49   private Query subQuery;
  50   private ValueSourceQuery[] valSrcQueries; // never null (empty array if there are no valSrcQueries).
  51   private boolean strict = false; // if true, valueSource part of query does not take part in weights normalization.
  52
  53   /**
  54    * Create a CustomScoreQuery over input subQuery.
  55    * @param subQuery the sub query whose scored is being customized. Must not be null.
  56    */
  57   public CustomScoreQuery(Query subQuery) {
  58     this(subQuery, new ValueSourceQuery[0]);
  59   }
  60
  61   /**
  62    * Create a CustomScoreQuery over input subQuery and a {@link ValueSourceQuery}.
  63    * @param subQuery the sub query whose score is being customized. Must not be null.
  64    * @param valSrcQuery a value source query whose scores are used in the custom score
  65    * computation. For most simple/convenient use case this would be a
  66    * {@link org.apache.lucene.search.function.FieldScoreQuery FieldScoreQuery}.
  67    * This parameter is optional - it can be null.
  68    */
  69   public CustomScoreQuery(Query subQuery, ValueSourceQuery valSrcQuery) {
  70           this(subQuery, valSrcQuery!=null ? // don't want an array that contains a single null..
  71         new ValueSourceQuery[] {valSrcQuery} : new ValueSourceQuery[0]);
  72   }
  73
  74   /**
  75    * Create a CustomScoreQuery over input subQuery and a {@link ValueSourceQuery}.
  76    * @param subQuery the sub query whose score is being customized. Must not be null.
  77    * @param valSrcQueries value source queries whose scores are used in the custom score
  78    * computation. For most simple/convenient use case these would be
  79    * {@link org.apache.lucene.search.function.FieldScoreQuery FieldScoreQueries}.
  80    * This parameter is optional - it can be null or even an empty array.
  81    */
  82   public CustomScoreQuery(Query subQuery, ValueSourceQuery... valSrcQueries) {
  83     this.subQuery = subQuery;
  84     this.valSrcQueries = valSrcQueries!=null?
  85         valSrcQueries : new ValueSourceQuery[0];
  86     if (subQuery == null) throw new IllegalArgumentException("<subquery> must not be null!");
  87   }
  88
  89   /*(non-Javadoc) @see org.apache.lucene.search.Query#rewrite(org.apache.lucene.index.IndexReader) */
  90   @Override
  91   public Query rewrite(IndexReader reader) throws IOException {
  92     CustomScoreQuery clone = null;
  93
  94     final Query sq = subQuery.rewrite(reader);
  95     if (sq != subQuery) {
  96       clone = (CustomScoreQuery) clone();
  97       clone.subQuery = sq;
  98     }
  99
 100     for(int i = 0; i < valSrcQueries.length; i++) {
 101       final ValueSourceQuery v = (ValueSourceQuery) valSrcQueries[i].rewrite(reader);
 102       if (v != valSrcQueries[i]) {
 103         if (clone == null) clone = (CustomScoreQuery) clone();
 104         clone.valSrcQueries[i] = v;
 105       }
 106     }
 107
 108     return (clone == null) ? this : clone;
 109   }
 110
 111   /*(non-Javadoc) @see org.apache.lucene.search.Query#extractTerms(java.util.Set) */
 112   @Override
 113   public void extractTerms(Set<Term> terms) {
 114     subQuery.extractTerms(terms);
 115     for(int i = 0; i < valSrcQueries.length; i++) {
 116       valSrcQueries[i].extractTerms(terms);
 117     }
 118   }
 119
 120   /*(non-Javadoc) @see org.apache.lucene.search.Query#clone() */
 121   @Override
 122   public Object clone() {
 123     CustomScoreQuery clone = (CustomScoreQuery)super.clone();
 124     clone.subQuery = (Query) subQuery.clone();
 125     clone.valSrcQueries = new ValueSourceQuery[valSrcQueries.length];
 126     for(int i = 0; i < valSrcQueries.length; i++) {
 127       clone.valSrcQueries[i] = (ValueSourceQuery) valSrcQueries[i].clone();
 128     }
 129     return clone;
 130   }
 131
 132   /* (non-Javadoc) @see org.apache.lucene.search.Query#toString(java.lang.String) */
 133   @Override
 134   public String toString(String field) {
 135     StringBuilder sb = new StringBuilder(name()).append("(");
 136     sb.append(subQuery.toString(field));
 137     for(int i = 0; i < valSrcQueries.length; i++) {
 138       sb.append(", ").append(valSrcQueries[i].toString(field));
 139     }
 140     sb.append(")");
 141     sb.append(strict?" STRICT" : "");
 142     return sb.toString() + ToStringUtils.boost(getBoost());
 143   }
 144
 145   /** Returns true if <code>o</code> is equal to this. */
 146   @Override
 147   public boolean equals(Object o) {
 148     if (this == o)
 149       return true;
 150     if (!super.equals(o))
 151       return false;
 152     if (getClass() != o.getClass()) {
 153       return false;
 154     }
 155     CustomScoreQuery other = (CustomScoreQuery)o;
 156     if (this.getBoost() != other.getBoost() ||
 157         !this.subQuery.equals(other.subQuery) ||
 158         this.strict != other.strict ||
 159         this.valSrcQueries.length != other.valSrcQueries.length) {
 160       return false;
 161     }
 162     return Arrays.equals(valSrcQueries, other.valSrcQueries);
 163   }
 164
 165   /** Returns a hash code value for this object. */
 166   @Override
 167   public int hashCode() {
 168     return (getClass().hashCode() + subQuery.hashCode() + Arrays.hashCode(valSrcQueries))
 169       ^ Float.floatToIntBits(getBoost()) ^ (strict ? 1234 : 4321);
 170   }
 171
 172   /**
 173    * Returns a {@link CustomScoreProvider} that calculates the custom scores
 174    * for the given {@link IndexReader}. The default implementation returns a default
 175    * implementation as specified in the docs of {@link CustomScoreProvider}.
 176    * @since 2.9.2
 177    */
 178   protected CustomScoreProvider getCustomScoreProvider(IndexReader reader) throws IOException {
 179     return new CustomScoreProvider(reader);
 180   }
 181
 182   //=========================== W E I G H T ============================
 183
 184   private class CustomWeight extends Weight {
 185     Similarity similarity;
 186     Weight subQueryWeight;
 187     Weight[] valSrcWeights;
 188     boolean qStrict;
 189
 190     public CustomWeight(Searcher searcher) throws IOException {
 191       this.similarity = getSimilarity(searcher);
 192       this.subQueryWeight = subQuery.createWeight(searcher);
 193       this.valSrcWeights = new Weight[valSrcQueries.length];
 194       for(int i = 0; i < valSrcQueries.length; i++) {
 195         this.valSrcWeights[i] = valSrcQueries[i].createWeight(searcher);
 196       }
 197       this.qStrict = strict;
 198     }
 199
 200     /*(non-Javadoc) @see org.apache.lucene.search.Weight#getQuery() */
 201     @Override
 202     public Query getQuery() {
 203       return CustomScoreQuery.this;
 204     }
 205
 206     /*(non-Javadoc) @see org.apache.lucene.search.Weight#getValue() */
 207     @Override
 208     public float getValue() {
 209       return getBoost();
 210     }
 211
 212     /*(non-Javadoc) @see org.apache.lucene.search.Weight#sumOfSquaredWeights() */
 213     @Override
 214     public float sumOfSquaredWeights() throws IOException {
 215       float sum = subQueryWeight.sumOfSquaredWeights();
 216       for(int i = 0; i < valSrcWeights.length; i++) {
 217         if (qStrict) {
 218           valSrcWeights[i].sumOfSquaredWeights(); // do not include ValueSource part in the query normalization
 219         } else {
 220           sum += valSrcWeights[i].sumOfSquaredWeights();
 221         }
 222       }
 223       sum *= getBoost() * getBoost(); // boost each sub-weight
 224       return sum ;
 225     }
 226
 227     /*(non-Javadoc) @see org.apache.lucene.search.Weight#normalize(float) */
 228     @Override
 229     public void normalize(float norm) {
 230       norm *= getBoost(); // incorporate boost
 231       subQueryWeight.normalize(norm);
 232       for(int i = 0; i < valSrcWeights.length; i++) {
 233         if (qStrict) {
 234           valSrcWeights[i].normalize(1); // do not normalize the ValueSource part
 235         } else {
 236           valSrcWeights[i].normalize(norm);
 237         }
 238       }
 239     }
 240
 241     @Override
 242     public Scorer scorer(IndexReader reader, boolean scoreDocsInOrder, boolean topScorer) throws IOException {
 243       // Pass true for "scoresDocsInOrder", because we
 244       // require in-order scoring, even if caller does not,
 245       // since we call advance on the valSrcScorers.  Pass
 246       // false for "topScorer" because we will not invoke
 247       // score(Collector) on these scorers:
 248       Scorer subQueryScorer = subQueryWeight.scorer(reader, true, false);
 249       if (subQueryScorer == null) {
 250         return null;
 251       }
 252       Scorer[] valSrcScorers = new Scorer[valSrcWeights.length];
 253       for(int i = 0; i < valSrcScorers.length; i++) {
 254          valSrcScorers[i] = valSrcWeights[i].scorer(reader, true, topScorer);
 255       }
 256       return new CustomScorer(similarity, reader, this, subQueryScorer, valSrcScorers);
 257     }
 258
 259     @Override
 260     public Explanation explain(IndexReader reader, int doc) throws IOException {
 261       Explanation explain = doExplain(reader, doc);
 262       return explain == null ? new Explanation(0.0f, "no matching docs") : explain;
 263     }
 264
 265     private Explanation doExplain(IndexReader reader, int doc) throws IOException {
 266       Explanation subQueryExpl = subQueryWeight.explain(reader, doc);
 267       if (!subQueryExpl.isMatch()) {
 268         return subQueryExpl;
 269       }
 270       // match
 271       Explanation[] valSrcExpls = new Explanation[valSrcWeights.length];
 272       for(int i = 0; i < valSrcWeights.length; i++) {
 273         valSrcExpls[i] = valSrcWeights[i].explain(reader, doc);
 274       }
 275       Explanation customExp = CustomScoreQuery.this.getCustomScoreProvider(reader).customExplain(doc,subQueryExpl,valSrcExpls);
 276       float sc = getValue() * customExp.getValue();
 277       Explanation res = new ComplexExplanation(
 278         true, sc, CustomScoreQuery.this.toString() + ", product of:");
 279       res.addDetail(customExp);
 280       res.addDetail(new Explanation(getValue(), "queryBoost")); // actually using the q boost as q weight (== weight value)
 281       return res;
 282     }
 283
 284     @Override
 285     public boolean scoresDocsOutOfOrder() {
 286       return false;
 287     }
 288
 289   }
 290
 291
 292   //=========================== S C O R E R ============================
 293
 294   /**
 295    * A scorer that applies a (callback) function on scores of the subQuery.
 296    */
 297   private class CustomScorer extends Scorer {
 298     private final float qWeight;
 299     private Scorer subQueryScorer;
 300     private Scorer[] valSrcScorers;
 301     private final CustomScoreProvider provider;
 302     private float vScores[]; // reused in score() to avoid allocating this array for each doc
 303
 304     // constructor
 305     private CustomScorer(Similarity similarity, IndexReader reader, CustomWeight w,
 306         Scorer subQueryScorer, Scorer[] valSrcScorers) throws IOException {
 307       super(similarity,w);
 308       this.qWeight = w.getValue();
 309       this.subQueryScorer = subQueryScorer;
 310       this.valSrcScorers = valSrcScorers;
 311       this.vScores = new float[valSrcScorers.length];
 312       this.provider = CustomScoreQuery.this.getCustomScoreProvider(reader);
 313     }
 314
 315     @Override
 316     public int nextDoc() throws IOException {
 317       int doc = subQueryScorer.nextDoc();
 318       if (doc != NO_MORE_DOCS) {
 319         for (int i = 0; i < valSrcScorers.length; i++) {
 320           valSrcScorers[i].advance(doc);
 321         }
 322       }
 323       return doc;
 324     }
 325
 326     @Override
 327     public int docID() {
 328       return subQueryScorer.docID();
 329     }
 330
 331     /*(non-Javadoc) @see org.apache.lucene.search.Scorer#score() */
 332     @Override
 333     public float score() throws IOException {
 334       for (int i = 0; i < valSrcScorers.length; i++) {
 335         vScores[i] = valSrcScorers[i].score();
 336       }
 337       return qWeight * provider.customScore(subQueryScorer.docID(), subQueryScorer.score(), vScores);
 338     }
 339
 340     @Override
 341     public int advance(int target) throws IOException {
 342       int doc = subQueryScorer.advance(target);
 343       if (doc != NO_MORE_DOCS) {
 344         for (int i = 0; i < valSrcScorers.length; i++) {
 345           valSrcScorers[i].advance(doc);
 346         }
 347       }
 348       return doc;
 349     }
 350   }
 351
 352   @Override
 353   public Weight createWeight(Searcher searcher) throws IOException {
 354     return new CustomWeight(searcher);
 355   }
 356
 357   /**
 358    * Checks if this is strict custom scoring.
 359    * In strict custom scoring, the ValueSource part does not participate in weight normalization.
 360    * This may be useful when one wants full control over how scores are modified, and does
 361    * not care about normalizing by the ValueSource part.
 362    * One particular case where this is useful if for testing this query.
 363    * <P>
 364    * Note: only has effect when the ValueSource part is not null.
 365    */
 366   public boolean isStrict() {
 367     return strict;
 368   }
 369
 370   /**
 371    * Set the strict mode of this query.
 372    * @param strict The strict mode to set.
 373    * @see #isStrict()
 374    */
 375   public void setStrict(boolean strict) {
 376     this.strict = strict;
 377   }
 378
 379   /**
 380    * A short name of this query, used in {@link #toString(String)}.
 381    */
 382   public String name() {
 383     return "custom";
 384   }
 385
 386 }