lucene-java-3.4.0/lucene/contrib/join/src/java/org/apache/lucene/search/join/BlockJoinQuery.java

   1 package org.apache.lucene.search.join;
   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
  23 import org.apache.lucene.index.IndexReader;
  24 import org.apache.lucene.index.IndexWriter;       // javadocs
  25 import org.apache.lucene.index.Term;
  26 import org.apache.lucene.search.BooleanClause;
  27 import org.apache.lucene.search.DocIdSet;
  28 import org.apache.lucene.search.DocIdSetIterator;
  29 import org.apache.lucene.search.Explanation;
  30 import org.apache.lucene.search.Filter;
  31 import org.apache.lucene.search.Query;
  32 import org.apache.lucene.search.Scorer;
  33 import org.apache.lucene.search.Searcher;
  34 import org.apache.lucene.search.Weight;
  35 import org.apache.lucene.search.grouping.TopGroups;
  36 import org.apache.lucene.util.ArrayUtil;
  37 import org.apache.lucene.util.FixedBitSet;
  38
  39 /**
  40  * This query requires that you index
  41  * children and parent docs as a single block, using the
  42  * {@link IndexWriter#addDocuments} or {@link
  43  * IndexWriter#updateDocuments} API.  In each block, the
  44  * child documents must appear first, ending with the parent
  45  * document.  At search time you provide a Filter
  46  * identifying the parents, however this Filter must provide
  47  * an {@link FixedBitSet} per sub-reader.
  48  *
  49  * <p>Once the block index is built, use this query to wrap
  50  * any sub-query matching only child docs and join matches in that
  51  * child document space up to the parent document space.
  52  * You can then use this Query as a clause with
  53  * other queries in the parent document space.</p>
  54  *
  55  * <p>The child documents must be orthogonal to the parent
  56  * documents: the wrapped child query must never
  57  * return a parent document.</p>
  58  *
  59  * If you'd like to retrieve {@link TopGroups} for the
  60  * resulting query, use the {@link BlockJoinCollector}.
  61  * Note that this is not necessary, ie, if you simply want
  62  * to collect the parent documents and don't need to see
  63  * which child documents matched under that parent, then
  64  * you can use any collector.
  65  *
  66  * <p><b>NOTE</b>: If the overall query contains parent-only
  67  * matches, for example you OR a parent-only query with a
  68  * joined child-only query, then the resulting collected documents
  69  * will be correct, however the {@link TopGroups} you get
  70  * from {@link BlockJoinCollector} will not contain every
  71  * child for parents that had matched.
  72  *
  73  * <p>See {@link org.apache.lucene.search.join} for an
  74  * overview. </p>
  75  *
  76  * @lucene.experimental
  77  */
  78
  79 public class BlockJoinQuery extends Query {
  80
  81   public static enum ScoreMode {None, Avg, Max, Total};
  82
  83   private final Filter parentsFilter;
  84   private final Query childQuery;
  85
  86   // If we are rewritten, this is the original childQuery we
  87   // were passed; we use this for .equals() and
  88   // .hashCode().  This makes rewritten query equal the
  89   // original, so that user does not have to .rewrite() their
  90   // query before searching:
  91   private final Query origChildQuery;
  92   private final ScoreMode scoreMode;
  93
  94   public BlockJoinQuery(Query childQuery, Filter parentsFilter, ScoreMode scoreMode) {
  95     super();
  96     this.origChildQuery = childQuery;
  97     this.childQuery = childQuery;
  98     this.parentsFilter = parentsFilter;
  99     this.scoreMode = scoreMode;
 100   }
 101
 102   private BlockJoinQuery(Query origChildQuery, Query childQuery, Filter parentsFilter, ScoreMode scoreMode) {
 103     super();
 104     this.origChildQuery = origChildQuery;
 105     this.childQuery = childQuery;
 106     this.parentsFilter = parentsFilter;
 107     this.scoreMode = scoreMode;
 108   }
 109
 110   @Override
 111   public Weight createWeight(Searcher searcher) throws IOException {
 112     return new BlockJoinWeight(this, childQuery.createWeight(searcher), parentsFilter, scoreMode);
 113   }
 114
 115   private static class BlockJoinWeight extends Weight {
 116     private final Query joinQuery;
 117     private final Weight childWeight;
 118     private final Filter parentsFilter;
 119     private final ScoreMode scoreMode;
 120
 121     public BlockJoinWeight(Query joinQuery, Weight childWeight, Filter parentsFilter, ScoreMode scoreMode) {
 122       super();
 123       this.joinQuery = joinQuery;
 124       this.childWeight = childWeight;
 125       this.parentsFilter = parentsFilter;
 126       this.scoreMode = scoreMode;
 127     }
 128
 129     @Override
 130     public Query getQuery() {
 131       return joinQuery;
 132     }
 133
 134     @Override
 135     public float getValue() {
 136       return childWeight.getValue();
 137     }
 138
 139     @Override
 140     public float sumOfSquaredWeights() throws IOException {
 141       return childWeight.sumOfSquaredWeights();
 142     }
 143
 144     @Override
 145     public void normalize(float norm) {
 146       childWeight.normalize(norm);
 147     }
 148
 149     @Override
 150     public Scorer scorer(IndexReader reader, boolean scoreDocsInOrder, boolean topScorer) throws IOException {
 151       // Pass scoreDocsInOrder true, topScorer false to our sub:
 152       final Scorer childScorer = childWeight.scorer(reader, true, false);
 153
 154       if (childScorer == null) {
 155         // No matches
 156         return null;
 157       }
 158
 159       final int firstChildDoc = childScorer.nextDoc();
 160       if (firstChildDoc == DocIdSetIterator.NO_MORE_DOCS) {
 161         // No matches
 162         return null;
 163       }
 164
 165       final DocIdSet parents = parentsFilter.getDocIdSet(reader);
 166       // TODO: once we do random-access filters we can
 167       // generalize this:
 168       if (parents == null) {
 169         // No matches
 170         return null;
 171       }
 172       if (!(parents instanceof FixedBitSet)) {
 173         throw new IllegalStateException("parentFilter must return FixedBitSet; got " + parents);
 174       }
 175
 176       return new BlockJoinScorer(this, childScorer, (FixedBitSet) parents, firstChildDoc, scoreMode);
 177     }
 178
 179     @Override
 180     public Explanation explain(IndexReader reader, int doc) throws IOException {
 181       // TODO
 182       throw new UnsupportedOperationException(getClass().getName() +
 183                                               " cannot explain match on parent document");
 184     }
 185
 186     @Override
 187     public boolean scoresDocsOutOfOrder() {
 188       return false;
 189     }
 190   }
 191
 192   static class BlockJoinScorer extends Scorer {
 193     private final Scorer childScorer;
 194     private final FixedBitSet parentBits;
 195     private final ScoreMode scoreMode;
 196     private int parentDoc;
 197     private float parentScore;
 198     private int nextChildDoc;
 199
 200     private int[] pendingChildDocs = new int[5];
 201     private float[] pendingChildScores;
 202     private int childDocUpto;
 203
 204     public BlockJoinScorer(Weight weight, Scorer childScorer, FixedBitSet parentBits, int firstChildDoc, ScoreMode scoreMode) {
 205       super(weight);
 206       //System.out.println("Q.init firstChildDoc=" + firstChildDoc);
 207       this.parentBits = parentBits;
 208       this.childScorer = childScorer;
 209       this.scoreMode = scoreMode;
 210       if (scoreMode != ScoreMode.None) {
 211         pendingChildScores = new float[5];
 212       }
 213       nextChildDoc = firstChildDoc;
 214     }
 215
 216     @Override
 217     public void visitSubScorers(Query parent, BooleanClause.Occur relationship,
 218                                 ScorerVisitor<Query, Query, Scorer> visitor) {
 219       super.visitSubScorers(parent, relationship, visitor);
 220       //childScorer.visitSubScorers(weight.getQuery(), BooleanClause.Occur.MUST, visitor);
 221       childScorer.visitScorers(visitor);
 222     }
 223
 224     int getChildCount() {
 225       return childDocUpto;
 226     }
 227
 228     int[] swapChildDocs(int[] other) {
 229       final int[] ret = pendingChildDocs;
 230       if (other == null) {
 231         pendingChildDocs = new int[5];
 232       } else {
 233         pendingChildDocs = other;
 234       }
 235       return ret;
 236     }
 237
 238     float[] swapChildScores(float[] other) {
 239       if (scoreMode == ScoreMode.None) {
 240         throw new IllegalStateException("ScoreMode is None");
 241       }
 242       final float[] ret = pendingChildScores;
 243       if (other == null) {
 244         pendingChildScores = new float[5];
 245       } else {
 246         pendingChildScores = other;
 247       }
 248       return ret;
 249     }
 250
 251     @Override
 252     public int nextDoc() throws IOException {
 253       //System.out.println("Q.nextDoc() nextChildDoc=" + nextChildDoc);
 254
 255       if (nextChildDoc == NO_MORE_DOCS) {
 256         //System.out.println("  end");
 257         return parentDoc = NO_MORE_DOCS;
 258       }
 259
 260       // Gather all children sharing the same parent as nextChildDoc
 261       parentDoc = parentBits.nextSetBit(nextChildDoc);
 262       //System.out.println("  parentDoc=" + parentDoc);
 263       assert parentDoc != -1;
 264
 265       float totalScore = 0;
 266       float maxScore = Float.NEGATIVE_INFINITY;
 267
 268       childDocUpto = 0;
 269       do {
 270         //System.out.println("  c=" + nextChildDoc);
 271         if (pendingChildDocs.length == childDocUpto) {
 272           pendingChildDocs = ArrayUtil.grow(pendingChildDocs);
 273           if (scoreMode != ScoreMode.None) {
 274             pendingChildScores = ArrayUtil.grow(pendingChildScores);
 275           }
 276         }
 277         pendingChildDocs[childDocUpto] = nextChildDoc;
 278         if (scoreMode != ScoreMode.None) {
 279           // TODO: specialize this into dedicated classes per-scoreMode
 280           final float childScore = childScorer.score();
 281           pendingChildScores[childDocUpto] = childScore;
 282           maxScore = Math.max(childScore, maxScore);
 283           totalScore += childScore;
 284         }
 285         childDocUpto++;
 286         nextChildDoc = childScorer.nextDoc();
 287       } while (nextChildDoc < parentDoc);
 288       //System.out.println("  nextChildDoc=" + nextChildDoc);
 289
 290       // Parent & child docs are supposed to be orthogonal:
 291       assert nextChildDoc != parentDoc;
 292
 293       switch(scoreMode) {
 294       case Avg:
 295         parentScore = totalScore / childDocUpto;
 296         break;
 297       case Max:
 298         parentScore = maxScore;
 299         break;
 300       case Total:
 301         parentScore = totalScore;
 302         break;
 303       case None:
 304         break;
 305       }
 306
 307       //System.out.println("  return parentDoc=" + parentDoc);
 308       return parentDoc;
 309     }
 310
 311     @Override
 312     public int docID() {
 313       return parentDoc;
 314     }
 315
 316     @Override
 317     public float score() throws IOException {
 318       return parentScore;
 319     }
 320
 321     @Override
 322     public int advance(int parentTarget) throws IOException {
 323
 324       //System.out.println("Q.advance parentTarget=" + parentTarget);
 325       if (parentTarget == NO_MORE_DOCS) {
 326         return parentDoc = NO_MORE_DOCS;
 327       }
 328
 329       final int prevParentDoc = parentBits.prevSetBit(parentTarget-1);
 330
 331       //System.out.println("  rolled back to prevParentDoc=" + prevParentDoc + " vs parentDoc=" + parentDoc);
 332       assert prevParentDoc >= parentDoc;
 333       if (prevParentDoc > nextChildDoc) {
 334         nextChildDoc = childScorer.advance(prevParentDoc);
 335         // System.out.println("  childScorer advanced to child docID=" + nextChildDoc);
 336       //} else {
 337         //System.out.println("  skip childScorer advance");
 338       }
 339
 340       // Parent & child docs are supposed to be orthogonal:
 341       assert nextChildDoc != prevParentDoc;
 342
 343       final int nd = nextDoc();
 344       //System.out.println("  return nextParentDoc=" + nd);
 345       return nd;
 346     }
 347   }
 348
 349   @Override
 350   public void extractTerms(Set<Term> terms) {
 351     childQuery.extractTerms(terms);
 352   }
 353
 354   @Override
 355   public Query rewrite(IndexReader reader) throws IOException {
 356     final Query childRewrite = childQuery.rewrite(reader);
 357     if (childRewrite != childQuery) {
 358       return new BlockJoinQuery(childQuery,
 359                                 childRewrite,
 360                                 parentsFilter,
 361                                 scoreMode);
 362     } else {
 363       return this;
 364     }
 365   }
 366
 367   @Override
 368   public String toString(String field) {
 369     return "BlockJoinQuery ("+childQuery.toString()+")";
 370   }
 371
 372   @Override
 373   public void setBoost(float boost) {
 374     throw new UnsupportedOperationException("this query cannot support boosting; please use childQuery.setBoost instead");
 375   }
 376
 377   @Override
 378   public float getBoost() {
 379     throw new UnsupportedOperationException("this query cannot support boosting; please use childQuery.getBoost instead");
 380   }
 381
 382   @Override
 383   public boolean equals(Object _other) {
 384     if (_other instanceof BlockJoinQuery) {
 385       final BlockJoinQuery other = (BlockJoinQuery) _other;
 386       return origChildQuery.equals(other.origChildQuery) &&
 387         parentsFilter.equals(other.parentsFilter) &&
 388         scoreMode == other.scoreMode;
 389     } else {
 390       return false;
 391     }
 392   }
 393
 394   @Override
 395   public int hashCode() {
 396     final int prime = 31;
 397     int hash = 1;
 398     hash = prime * hash + origChildQuery.hashCode();
 399     hash = prime * hash + scoreMode.hashCode();
 400     hash = prime * hash + parentsFilter.hashCode();
 401     return hash;
 402   }
 403
 404   @Override
 405   public Object clone() {
 406     return new BlockJoinQuery((Query) origChildQuery.clone(),
 407                               parentsFilter,
 408                               scoreMode);
 409   }
 410 }