Decode doc ids in BKD leaves with auto-vectorized loops (#14203)

Author: gf2121

lucene/CHANGES.txt

@@ -178,6 +178,8 @@ Optimizations
 
 * GITHUB#14301: Avoid unnecessary evaluations and skipping documents. (hanbj)
 
+# GITHUB#14203: Decode doc ids in BKD leaves with auto-vectorized loops when using DEFAULT_MAX_POINTS_IN_LEAF_NODE. (Guo Feng)
+
 Bug Fixes
 ---------------------
 

lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90PointsFormat.java

@@ -17,11 +17,13 @@
 package org.apache.lucene.codecs.lucene90;
 
 import java.io.IOException;
+import java.util.Map;
 import org.apache.lucene.codecs.PointsFormat;
 import org.apache.lucene.codecs.PointsReader;
 import org.apache.lucene.codecs.PointsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.util.bkd.BKDWriter;
 
 /**
  * Lucene 9.0 point format, which encodes dimensional values in a block KD-tree structure for fast
@@ -59,18 +61,40 @@ public final class Lucene90PointsFormat extends PointsFormat {
   public static final String META_EXTENSION = "kdm";
 
   static final int VERSION_START = 0;
-  static final int VERSION_CURRENT = VERSION_START;
+  static final int VERSION_BKD_VECTORIZED_BPV24 = 1;
+  static final int VERSION_CURRENT = VERSION_BKD_VECTORIZED_BPV24;
+
+  private static final Map<Integer, Integer> VERSION_TO_BKD_VERSION =
+      Map.of(
+          VERSION_START, BKDWriter.VERSION_META_FILE,
+          VERSION_BKD_VECTORIZED_BPV24, BKDWriter.VERSION_VECTORIZED_DOCID);
+
+  private final int version;
 
   /** Sole constructor */
-  public Lucene90PointsFormat() {}
+  public Lucene90PointsFormat() {
+    this(VERSION_CURRENT);
+  }
+
+  /** Constructor that takes a version. This is used for testing with older versions. */
+  Lucene90PointsFormat(int version) {
+    if (VERSION_TO_BKD_VERSION.containsKey(version) == false) {
+      throw new IllegalArgumentException("Invalid version: " + version);
+    }
+    this.version = version;
+  }
 
   @Override
   public PointsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-    return new Lucene90PointsWriter(state);
+    return new Lucene90PointsWriter(state, version);
   }
 
   @Override
   public PointsReader fieldsReader(SegmentReadState state) throws IOException {
     return new Lucene90PointsReader(state);
   }
+
+  static int bkdVersion(int version) {
+    return VERSION_TO_BKD_VERSION.get(version);
+  }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90PointsWriter.java

@@ -46,16 +46,18 @@ public class Lucene90PointsWriter extends PointsWriter {
   final SegmentWriteState writeState;
   final int maxPointsInLeafNode;
   final double maxMBSortInHeap;
+  final int version;
   private boolean finished;
 
   /** Full constructor */
   public Lucene90PointsWriter(
-      SegmentWriteState writeState, int maxPointsInLeafNode, double maxMBSortInHeap)
+      SegmentWriteState writeState, int maxPointsInLeafNode, double maxMBSortInHeap, int version)
       throws IOException {
     assert writeState.fieldInfos.hasPointValues();
     this.writeState = writeState;
     this.maxPointsInLeafNode = maxPointsInLeafNode;
     this.maxMBSortInHeap = maxMBSortInHeap;
+    this.version = version;
     String dataFileName =
         IndexFileNames.segmentFileName(
             writeState.segmentInfo.name,
@@ -105,6 +107,12 @@ public Lucene90PointsWriter(
     }
   }
 
+  public Lucene90PointsWriter(
+      SegmentWriteState writeState, int maxPointsInLeafNode, double maxMBSortInHeap)
+      throws IOException {
+    this(writeState, maxPointsInLeafNode, maxMBSortInHeap, Lucene90PointsFormat.VERSION_CURRENT);
+  }
+
   /**
    * Uses the defaults values for {@code maxPointsInLeafNode} (512) and {@code maxMBSortInHeap}
    * (16.0)
@@ -113,7 +121,17 @@ public Lucene90PointsWriter(SegmentWriteState writeState) throws IOException {
     this(
         writeState,
         BKDConfig.DEFAULT_MAX_POINTS_IN_LEAF_NODE,
-        BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP);
+        BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP,
+        Lucene90PointsFormat.VERSION_CURRENT);
+  }
+
+  /** Constructor that takes a version. This is used for testing with older versions. */
+  Lucene90PointsWriter(SegmentWriteState writeState, int version) throws IOException {
+    this(
+        writeState,
+        BKDConfig.DEFAULT_MAX_POINTS_IN_LEAF_NODE,
+        BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP,
+        version);
   }
 
   @Override
@@ -135,7 +153,8 @@ public void writeField(FieldInfo fieldInfo, PointsReader reader) throws IOExcept
             writeState.segmentInfo.name,
             config,
             maxMBSortInHeap,
-            values.size())) {
+            values.size(),
+            Lucene90PointsFormat.bkdVersion(version))) {
 
       if (values instanceof MutablePointTree) {
         IORunnable finalizer =
@@ -233,7 +252,8 @@ public void merge(MergeState mergeState) throws IOException {
                   writeState.segmentInfo.name,
                   config,
                   maxMBSortInHeap,
-                  totMaxSize)) {
+                  totMaxSize,
+                  Lucene90PointsFormat.bkdVersion(version))) {
             List<PointValues> pointValues = new ArrayList<>();
             List<MergeState.DocMap> docMaps = new ArrayList<>();
             for (int i = 0; i < mergeState.pointsReaders.length; i++) {

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

@@ -34,7 +34,6 @@
  * @lucene.experimental
  */
 public class BKDReader extends PointValues {
-
   final BKDConfig config;
   final int numLeaves;
   final IndexInput in;
@@ -269,7 +268,7 @@ private BKDPointTree(
           1,
           minPackedValue,
           maxPackedValue,
-          new BKDReaderDocIDSetIterator(config.maxPointsInLeafNode()),
+          new BKDReaderDocIDSetIterator(config.maxPointsInLeafNode(), version),
           new byte[config.packedBytesLength()],
           new byte[config.packedIndexBytesLength()],
           new byte[config.packedIndexBytesLength()],
@@ -598,7 +597,8 @@ public void addAll(PointValues.IntersectVisitor visitor, boolean grown) throws I
         // How many points are stored in this leaf cell:
         int count = leafNodes.readVInt();
         // No need to call grow(), it has been called up-front
-        docIdsWriter.readInts(leafNodes, count, visitor);
+        // Borrow scratchIterator.docIds as decoding buffer
+        docIdsWriter.readInts(leafNodes, count, visitor, scratchIterator.docIDs);
       } else {
         pushLeft();
         addAll(visitor, grown);
@@ -1036,9 +1036,9 @@ private static class BKDReaderDocIDSetIterator extends DocIdSetIterator {
     final int[] docIDs;
     private final DocIdsWriter docIdsWriter;
 
-    public BKDReaderDocIDSetIterator(int maxPointsInLeafNode) {
+    public BKDReaderDocIDSetIterator(int maxPointsInLeafNode, int version) {
       this.docIDs = new int[maxPointsInLeafNode];
-      this.docIdsWriter = new DocIdsWriter(maxPointsInLeafNode);
+      this.docIdsWriter = new DocIdsWriter(maxPointsInLeafNode, version);
     }
 
     @Override

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

@@ -85,7 +85,8 @@ public class BKDWriter implements Closeable {
   public static final int VERSION_SELECTIVE_INDEXING = 6;
   public static final int VERSION_LOW_CARDINALITY_LEAVES = 7;
   public static final int VERSION_META_FILE = 9;
-  public static final int VERSION_CURRENT = VERSION_META_FILE;
+  public static final int VERSION_VECTORIZED_DOCID = 10;
+  public static final int VERSION_CURRENT = VERSION_VECTORIZED_DOCID;
 
   /** Number of splits before we compute the exact bounding box of an inner node. */
   private static final int SPLITS_BEFORE_EXACT_BOUNDS = 4;
@@ -103,6 +104,7 @@ public class BKDWriter implements Closeable {
   final TrackingDirectoryWrapper tempDir;
   final String tempFileNamePrefix;
   final double maxMBSortInHeap;
+  final int version;
 
   final byte[] scratchDiff;
   final byte[] scratch;
@@ -139,6 +141,29 @@ public BKDWriter(
       BKDConfig config,
       double maxMBSortInHeap,
       long totalPointCount) {
+    this(
+        maxDoc,
+        tempDir,
+        tempFileNamePrefix,
+        config,
+        maxMBSortInHeap,
+        totalPointCount,
+        BKDWriter.VERSION_CURRENT);
+  }
+
+  /** This ctor should be only used for testing with older versions. */
+  public BKDWriter(
+      int maxDoc,
+      Directory tempDir,
+      String tempFileNamePrefix,
+      BKDConfig config,
+      double maxMBSortInHeap,
+      long totalPointCount,
+      int version) {
+    if (version < VERSION_START || version > VERSION_CURRENT) {
+      throw new IllegalArgumentException("Version out of range: " + version);
+    }
+    this.version = version;
     verifyParams(maxMBSortInHeap, totalPointCount);
     // We use tracking dir to deal with removing files on exception, so each place that
     // creates temp files doesn't need crazy try/finally/sucess logic:
@@ -165,7 +190,7 @@ public BKDWriter(
 
     // Maximum number of points we hold in memory at any time
     maxPointsSortInHeap = (int) ((maxMBSortInHeap * 1024 * 1024) / (config.bytesPerDoc()));
-    docIdsWriter = new DocIdsWriter(config.maxPointsInLeafNode());
+    docIdsWriter = new DocIdsWriter(config.maxPointsInLeafNode(), version);
     // Finally, we must be able to hold at least the leaf node in heap during build:
     if (maxPointsSortInHeap < config.maxPointsInLeafNode()) {
       throw new IllegalArgumentException(
@@ -1245,7 +1270,7 @@ private void writeIndex(
       byte[] packedIndex,
       long dataStartFP)
       throws IOException {
-    CodecUtil.writeHeader(metaOut, CODEC_NAME, VERSION_CURRENT);
+    CodecUtil.writeHeader(metaOut, CODEC_NAME, version);
     metaOut.writeVInt(config.numDims());
     metaOut.writeVInt(config.numIndexDims());
     metaOut.writeVInt(countPerLeaf);