Add Bulk Scorer For ToParentBlockJoinQuery

lucene/join/src/java/org/apache/lucene/search/join/ToParentBlockJoinQuery.java

@@ -24,20 +24,25 @@
 import java.util.Locale;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.BulkScorer;
 import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.Explanation;
+import org.apache.lucene.search.FilterLeafCollector;
 import org.apache.lucene.search.FilterWeight;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Matches;
 import org.apache.lucene.search.MatchesUtils;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.TwoPhaseIterator;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.Bits;
 
 /**
  * This query requires that you index children and parent docs as a single block, using the {@link
@@ -156,6 +161,15 @@ public Scorer get(long leadCost) throws IOException {
           return new BlockJoinScorer(childScorerSupplier.get(leadCost), parents, scoreMode);
         }
 
+        @Override
+        public BulkScorer bulkScorer() throws IOException {
+          final BulkScorer innerBulkScorer = childScorerSupplier.bulkScorer();
+          if (innerBulkScorer == null) {
+            return null;
+          }
+          return new BlockJoinBulkScorer(innerBulkScorer, scoreMode, parents);
+        }
+
         @Override
         public long cost() {
           return childScorerSupplier.cost();
@@ -275,6 +289,51 @@ public float matchCost() {
     }
   }
 
+  private static class Score extends Scorable {
+    private final ScoreMode scoreMode;
+    private Float score;
+    private int freq;
+
+    public Score(ScoreMode scoreMode) {
+      this.scoreMode = scoreMode;
+      reset();
+    }
+
+    public void reset() {
+      score = null;
+      freq = 0;
+    }
+
+    public void addChildScore(float childScore) {
+      freq += 1;
+      switch (scoreMode) {
+        case Total:
+        case Avg:
+          score = score == null ? childScore : score + childScore;
+          break;
+        case Min:
+          score = score == null ? childScore : Math.min(score, childScore);
+          break;
+        case Max:
+          score = score == null ? childScore : Math.max(score, childScore);
+          break;
+        case None:
+          break;
+        default:
+          throw new AssertionError();
+      }
+    }
+
+    @Override
+    public float score() {
+      float score = this.score != null ? this.score : 0;
+      if (scoreMode == ScoreMode.Avg && freq > 0) {
+        score /= freq;
+      }
+      return score;
+    }
+  }
+
   static class BlockJoinScorer extends Scorer {
     private final Scorer childScorer;
     private final BitSet parentBits;
@@ -357,6 +416,7 @@ private void setScoreAndFreq() throws IOException {
       while (childApproximation.nextDoc() < parentApproximation.docID()) {
         if (childTwoPhase == null || childTwoPhase.matches()) {
           final float childScore = scoreMode == ScoreMode.None ? 0 : childScorer.score();
+          // TODO: Refactor
           freq += 1;
           switch (scoreMode) {
             case Total:
@@ -440,6 +500,83 @@ private String formatScoreExplanation(int matches, int start, int end, ScoreMode
     }
   }
 
+  private abstract static class BatchAwareLeafCollector extends FilterLeafCollector {
+    public BatchAwareLeafCollector(LeafCollector in) {
+      super(in);
+    }
+
+    public void endBatch(int doc) throws IOException {}
+  }
+
+  private static class BlockJoinBulkScorer extends BulkScorer {
+    private final BulkScorer childBulkScorer;
+    private final ScoreMode scoreMode;
+    private final BitSet parents;
+    private final Score currentParentScore;
+    private Integer currentParent;
+
+    public BlockJoinBulkScorer(BulkScorer childBulkScorer, ScoreMode scoreMode, BitSet parents) {
+      this.childBulkScorer = childBulkScorer;
+      this.scoreMode = scoreMode;
+      this.parents = parents;
+      this.currentParentScore = new Score(scoreMode);
+      this.currentParent = null;
+    }
+
+    @Override
+    public int score(LeafCollector collector, Bits acceptDocs, int min, int max)
+        throws IOException {
+      BatchAwareLeafCollector wrappedCollector = wrapCollector(collector);
+      childBulkScorer.score(wrappedCollector, acceptDocs, min, max);
+      wrappedCollector.endBatch(max);
+      return max;
+    }
+
+    @Override
+    public long cost() {
+      return childBulkScorer.cost();
+    }
+
+    // TODO: Need to resolve parent doc IDs in multi-reader space?
+    private BatchAwareLeafCollector wrapCollector(LeafCollector collector) {
+      return new BatchAwareLeafCollector(collector) {
+        private Scorable scorer = null;
+
+        @Override
+        public void setScorer(Scorable scorer) throws IOException {
+          this.scorer = scorer;
+          super.setScorer(scorer != null ? currentParentScore : null);
+        }
+
+        @Override
+        public void collect(int doc) throws IOException {
+          if (currentParent == null) {
+            currentParent = parents.nextSetBit(doc);
+          } else if (doc > currentParent) {
+            in.collect(currentParent); // Emit the current parent
+
+            // Get the next parent and reset the score
+            currentParent = parents.nextSetBit(doc);
+            currentParentScore.reset();
+          }
+
+          if (scorer != null && scoreMode != ScoreMode.None) {
+            currentParentScore.addChildScore(scorer.score());
+          }
+        }
+
+        @Override
+        public void endBatch(int doc) throws IOException {
+          if (currentParent != null && doc > currentParent) {
+            in.collect(currentParent);
+            currentParent = null;
+            currentParentScore.reset();
+          }
+        }
+      };
+    }
+  }
+
   @Override
   public Query rewrite(IndexSearcher indexSearcher) throws IOException {
     final Query childRewrite = childQuery.rewrite(indexSearcher);