Support for traversing BKD tree with prefetching

Author: Saurabh Singh

lucene/CHANGES.txt

@@ -984,6 +984,8 @@ API Changes
 * GITHUB#13820, GITHUB#13825, GITHUB#13830: Corrects DataInput.readGroupVInts to be public and not-final, removes the protected
   DataInput.readGroupVInt method. (Zhang Chao, Robert Muir, Uwe Schindler, Dawid Weiss)
 
+* GITHUB#15376, GITHUB#15197: Added prefetching in bkd tree traversal, couple of new api in PointValues visitDocIDs from a position and prepareOrVisitDocIDs to prefetch the IO before visiting docIds (Saurabh Singh)
+
 New Features
 ---------------------
 

lucene/core/src/java/org/apache/lucene/index/PointValues.java

@@ -20,6 +20,8 @@
 import java.io.UncheckedIOException;
 import java.math.BigInteger;
 import java.net.InetAddress;
+import java.util.ArrayList;
+import java.util.List;
 import org.apache.lucene.document.BinaryPoint;
 import org.apache.lucene.document.DoublePoint;
 import org.apache.lucene.document.Field;
@@ -274,6 +276,19 @@ public interface PointTree extends Cloneable {
 
     /** Visit all the docs and values below the current node. */
     void visitDocValues(IntersectVisitor visitor) throws IOException;
+
+    /** Visit all the docs below the node at position pos */
+    default void visitDocIDs(long pos, IntersectVisitor visitor) throws IOException {}
+    ;
+
+    /**
+     * call prefetch for docs below the current node if vistor supports prefetching otherwise visit
+     * docIds
+     */
+    default void prepareOrVisitDocIDs(IntersectVisitor visitor) throws IOException {
+      visitDocIDs(visitor);
+    }
+    ;
   }
 
   /**
@@ -341,13 +356,91 @@ default void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOExcep
     default void grow(int count) {}
   }
 
+  /*
+   * We can recurse the {@link PointTree} using {@link TwoPhaseIntersectVisitor}. This visitor
+   * travere {@link PointTree} in two phases. In the first phase, it recurses over the {@link PointTree} optionally
+   * triggering IO for some of the blocks and caching them. In the second phase, once the recursion is over it visits
+   * the cached blocks one by one.
+   *
+   * @lucene.experimental
+   */
+  public interface TwoPhaseIntersectVisitor extends IntersectVisitor {
+    /** return the last deferred block ordinal during recursion. */
+    public int lastDeferredBlockOrdinal();
+
+    /** set last deferred block ordinal */
+    public void setLastDeferredBlockOrdinal(int leafNodeOrdinal);
+
+    /** Defer this block for processing in the second phase. */
+    public void deferBlock(long leafFp);
+
+    /** Returns a snapshot of the currently deferred blocks. */
+    public List<Long> deferredBlocks();
+
+    /** Mark the given block as processed and remove it from the deferred set. */
+    public void onProcessingDeferredBlock(long leafFp);
+  }
+
+  /**
+   * Base implementation of {@link TwoPhaseIntersectVisitor} that maintains a list of deferred
+   * blocks from first phase of traversal and visits them in the second phase.
+   *
+   * @lucene.experimental
+   */
+  public abstract static class BaseTwoPhaseIntersectVisitor implements TwoPhaseIntersectVisitor {
+
+    int lastDeferredBlockOrdinal = -1;
+    List<Long> deferredBlocks = new ArrayList<>();
+
+    /**
+     * return the last deferred block ordinal - this is used to avoid prefetching call for
+     * contiguous ordinals assuming contiguous ordinals prefetching can be taken care by readaheads.
+     */
+    @Override
+    public int lastDeferredBlockOrdinal() {
+      return lastDeferredBlockOrdinal;
+    }
+
+    /** set last deferred block ordinal * */
+    @Override
+    public void setLastDeferredBlockOrdinal(int leafNodeOrdinal) {
+      lastDeferredBlockOrdinal = leafNodeOrdinal;
+    }
+
+    /** Defer this block for processing in the second phase. */
+    @Override
+    public void deferBlock(long leafFp) {
+      deferredBlocks.add(leafFp);
+    }
+
+    /** Returns a snapshot of the currently deferred blocks. */
+    @Override
+    public List<Long> deferredBlocks() {
+      return new ArrayList<>(deferredBlocks);
+    }
+
+    /** Mark the given block as processed and remove it from the deferred set. */
+    @Override
+    public void onProcessingDeferredBlock(long leafFp) {
+      deferredBlocks.remove(leafFp);
+    }
+  }
+
   /**
    * Finds all documents and points matching the provided visitor. This method does not enforce live
    * documents, so it's up to the caller to test whether each document is deleted, if necessary.
    */
   public final void intersect(IntersectVisitor visitor) throws IOException {
     final PointTree pointTree = getPointTree();
     intersect(visitor, pointTree);
+    if (visitor instanceof TwoPhaseIntersectVisitor twoPhaseVisitor) {
+      List<Long> fps = twoPhaseVisitor.deferredBlocks();
+      for (int i = 0; i < fps.size(); ++i) {
+        long fp = fps.get(i);
+        pointTree.visitDocIDs(fp, visitor);
+        twoPhaseVisitor.onProcessingDeferredBlock(fp);
+      }
+    }
     assert pointTree.moveToParent() == false;
   }
 
@@ -358,7 +451,7 @@ private static void intersect(IntersectVisitor visitor, PointTree pointTree) thr
       if (compare == Relation.CELL_INSIDE_QUERY) {
         // This cell is fully inside the query shape: recursively add all points in this cell
         // without filtering
-        pointTree.visitDocIDs(visitor);
+        pointTree.prepareOrVisitDocIDs(visitor);
       } else if (compare == Relation.CELL_CROSSES_QUERY) {
         // The cell crosses the shape boundary, or the cell fully contains the query, so we fall
         // through and do full filtering:

lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java

@@ -589,6 +589,69 @@ public void visitDocIDs(PointValues.IntersectVisitor visitor) throws IOException
       addAll(visitor, false);
     }
 
+    @Override
+    public void prepareOrVisitDocIDs(IntersectVisitor visitor) throws IOException {
+      resetNodeDataPosition();
+      prefetchAll(visitor, false);
+    }
+
+    @Override
+    public void visitDocIDs(long position, IntersectVisitor visitor) throws IOException {
+      visitDocIDs(position, visitor, false);
+    }
+
+    private void visitDocIDs(long position, IntersectVisitor visitor, boolean grown)
+        throws IOException {
+      leafNodes.seek(position);
+      int count = leafNodes.readVInt();
+      if (!grown) {
+        visitor.grow(count);
+      }
+      docIdsWriter.readInts(leafNodes, count, visitor, scratchIterator.docIDs);
+    }
+
+    private int getLeafNodeOrdinal() {
+      assert isLeafNode() : "nodeID=" + nodeID + " is not a leaf";
+      return nodeID - leafNodeOffset;
+    }
+
+    public void prefetchAll(IntersectVisitor visitor, boolean grown) throws IOException {
+      if (grown == false) {
+        final long size = size();
+        if (size <= Integer.MAX_VALUE) {
+          visitor.grow((int) size);
+          grown = true;
+        }
+      }
+      if (isLeafNode()) {
+        // int count = isLastLeaf() ? config.maxPointsInLeafNode() : lastLeafNodePointCount;
+        long leafFp = getLeafBlockFP();
+        int leafNodeOrdinal = getLeafNodeOrdinal();
+        if (visitor instanceof BaseTwoPhaseIntersectVisitor twoPhaseIntersectVisitor) {
+          // Only call prefetch is this is the first leaf node ordinal or the first match in
+          // contigiuous sequence of matches for leaf nodes
+          // boolean prefetched = false;
+          if (twoPhaseIntersectVisitor.lastDeferredBlockOrdinal() == -1
+              || twoPhaseIntersectVisitor.lastDeferredBlockOrdinal() + 1 < leafNodeOrdinal) {
+            // System.out.println("Prefetched called on " + leafNodeOrdinal);
+            leafNodes.prefetch(leafFp, 1);
+            // prefetched = true;
+          }
+          twoPhaseIntersectVisitor.setLastDeferredBlockOrdinal(leafNodeOrdinal);
+          twoPhaseIntersectVisitor.deferBlock(leafFp);
+        } else {
+          visitDocIDs(getLeafBlockFP(), visitor, true);
+        }
+      } else {
+        pushLeft();
+        prefetchAll(visitor, grown);
+        pop();
+        pushRight();
+        prefetchAll(visitor, grown);
+        pop();
+      }
+    }
+
     public void addAll(PointValues.IntersectVisitor visitor, boolean grown) throws IOException {
       if (grown == false) {
         final long size = size();

lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestLucene90PointsFormat.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.List;
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.codecs.FilterCodec;
 import org.apache.lucene.codecs.PointsFormat;
@@ -355,4 +356,67 @@ public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
     r.close();
     dir.close();
   }
+
+  public void testBasicWithPrefetchVisitor() throws Exception {
+    Directory dir = newDirectory();
+    IndexWriterConfig iwc = newIndexWriterConfig();
+    // Avoid mockRandomMP since it may cause non-optimal merges that make the
+    // number of points per leaf hard to predict
+    while (iwc.getMergePolicy() instanceof MockRandomMergePolicy) {
+      iwc.setMergePolicy(newMergePolicy());
+    }
+    IndexWriter w = new IndexWriter(dir, iwc);
+    byte[] pointValue = new byte[3];
+    byte[] uniquePointValue = new byte[3];
+    random().nextBytes(uniquePointValue);
+    final int numDocs =
+        TEST_NIGHTLY ? atLeast(10000) : atLeast(500); // at night, make sure we have several leaves
+    final boolean multiValues = random().nextBoolean();
+    int totalValues = 0;
+    for (int i = 0; i < numDocs; ++i) {
+      Document doc = new Document();
+      if (i == numDocs / 2) {
+        totalValues++;
+        doc.add(new BinaryPoint("f", uniquePointValue));
+      } else {
+        final int numValues = (multiValues) ? TestUtil.nextInt(random(), 2, 100) : 1;
+        for (int j = 0; j < numValues; j++) {
+          do {
+            random().nextBytes(pointValue);
+          } while (Arrays.equals(pointValue, uniquePointValue));
+          doc.add(new BinaryPoint("f", pointValue));
+          totalValues++;
+        }
+      }
+      w.addDocument(doc);
+    }
+    w.forceMerge(1);
+    final IndexReader r = DirectoryReader.open(w);
+    w.close();
+
+    final LeafReader lr = getOnlyLeafReader(r);
+    PointValues points = lr.getPointValues("f");
+
+    PointValues.BaseTwoPhaseIntersectVisitor allPointsVisitor =
+        new PointValues.BaseTwoPhaseIntersectVisitor() {
+          @Override
+          public void visit(int docID, byte[] packedValue) throws IOException {}
+
+          @Override
+          public void visit(int docID) throws IOException {}
+
+          @Override
+          public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
+            return Relation.CELL_INSIDE_QUERY;
+          }
+        };
+
+    List<Long> savedBlocks = allPointsVisitor.deferredBlocks();
+    assertEquals(0, savedBlocks.size()); // Test that all deferred blocks were processed
+    assertEquals(totalValues, points.estimatePointCount(allPointsVisitor));
+    assertEquals(numDocs, points.estimateDocCount(allPointsVisitor));
+
+    r.close();
+    dir.close();
+  }
 }

lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java

@@ -17,8 +17,10 @@
 package org.apache.lucene.tests.index;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Objects;
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.DocValues;
@@ -1581,7 +1583,7 @@ public long size() {
     @Override
     public void visitDocIDs(IntersectVisitor visitor) throws IOException {
       in.visitDocIDs(
-          new AssertingIntersectVisitor(
+          new AssertingIntersectVisitorBase(
               pointValues.getNumDimensions(),
               pointValues.getNumIndexDimensions(),
               pointValues.getBytesPerDimension(),
@@ -1591,19 +1593,29 @@ public void visitDocIDs(IntersectVisitor visitor) throws IOException {
     @Override
     public void visitDocValues(IntersectVisitor visitor) throws IOException {
       in.visitDocValues(
-          new AssertingIntersectVisitor(
+          new AssertingIntersectVisitorBase(
               pointValues.getNumDimensions(),
               pointValues.getNumIndexDimensions(),
               pointValues.getBytesPerDimension(),
               visitor));
     }
+
+    @Override
+    public void visitDocIDs(long pos, IntersectVisitor visitor) throws IOException {
+      in.visitDocIDs(pos, visitor);
+    }
+
+    @Override
+    public void prepareOrVisitDocIDs(IntersectVisitor visitor) throws IOException {
+      in.prepareOrVisitDocIDs(visitor);
+    }
   }
 
   /**
    * Validates in the 1D case that all points are visited in order, and point values are in bounds
    * of the last cell checked
    */
-  static class AssertingIntersectVisitor implements IntersectVisitor {
+  static class AssertingIntersectVisitorBase extends PointValues.BaseTwoPhaseIntersectVisitor {
     final IntersectVisitor in;
     final int numDataDims;
     final int numIndexDims;
@@ -1614,8 +1626,10 @@ static class AssertingIntersectVisitor implements IntersectVisitor {
     private Relation lastCompareResult;
     private int lastDocID = -1;
     private int docBudget;
+    int lastMatchedBlock;
+    private List<Long> prefetchedBlocks;
 
-    AssertingIntersectVisitor(
+    AssertingIntersectVisitorBase(
         int numDataDims, int numIndexDims, int bytesPerDim, IntersectVisitor in) {
       this.in = in;
       this.numDataDims = numDataDims;
@@ -1716,6 +1730,26 @@ public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
       lastCompareResult = in.compare(minPackedValue, maxPackedValue);
       return lastCompareResult;
     }
+
+    @Override
+    public int lastDeferredBlockOrdinal() {
+      return lastMatchedBlock;
+    }
+
+    @Override
+    public void setLastDeferredBlockOrdinal(int leafNodeOrdinal) {
+      lastMatchedBlock = leafNodeOrdinal;
+    }
+
+    @Override
+    public void deferBlock(long leafFp) {
+      prefetchedBlocks.add(leafFp);
+    }
+
+    @Override
+    public List<Long> deferredBlocks() {
+      return new ArrayList<>(prefetchedBlocks);
+    }
   }
 
   @Override