From 2c65a616a88fa1da3020be99bcd656bf79f9c8d3 Mon Sep 17 00:00:00 2001
From: easyice <80152403@qq.com>
Date: Fri, 11 Oct 2024 12:07:06 +0800
Subject: [PATCH] Reduce allocations in BKDReaderDocIDSetIterator

---
 .../org/apache/lucene/util/bkd/BKDReader.java   |  6 +++---
 .../apache/lucene/util/bkd/DocIdsWriter.java    | 17 +++++++++--------
 .../lucene/util/bkd/TestDocIdsWriter.java       |  3 +--
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java
index 0efcc2ef465..c55b5b30da7 100644
--- a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java
+++ b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDReader.java
@@ -639,7 +639,7 @@ public class BKDReader extends PointValues {
       // How many points are stored in this leaf cell:
       int count = in.readVInt();
 
-      docIdsWriter.readInts(in, count, iterator.docIDs);
+      iterator.docIDs = docIdsWriter.readInts(in, count);
 
       return count;
     }
@@ -1011,11 +1011,10 @@ public class BKDReader extends PointValues {
     private int length;
     private int offset;
     private int docID;
-    final int[] docIDs;
+    private int[] docIDs;
     private final DocIdsWriter docIdsWriter;
 
     public BKDReaderDocIDSetIterator(int maxPointsInLeafNode) {
-      this.docIDs = new int[maxPointsInLeafNode];
       this.docIdsWriter = new DocIdsWriter(maxPointsInLeafNode);
     }
 
@@ -1037,6 +1036,7 @@ public class BKDReader extends PointValues {
       if (idx == length) {
         docID = DocIdSetIterator.NO_MORE_DOCS;
       } else {
+        assert docIDs != null;
         docID = docIDs[offset + idx];
         idx++;
       }
diff --git a/lucene/core/src/java/org/apache/lucene/util/bkd/DocIdsWriter.java b/lucene/core/src/java/org/apache/lucene/util/bkd/DocIdsWriter.java
index b9ea0d9aa08..31f5d3ee982 100644
--- a/lucene/core/src/java/org/apache/lucene/util/bkd/DocIdsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/util/bkd/DocIdsWriter.java
@@ -182,31 +182,32 @@ final class DocIdsWriter {
     assert currentWordIndex + 1 == totalWordCount;
   }
 
-  /** Read {@code count} integers into {@code docIDs}. */
-  void readInts(IndexInput in, int count, int[] docIDs) throws IOException {
+  /** Read {@code count} integers. */
+  int[] readInts(IndexInput in, int count) throws IOException {
     final int bpv = in.readByte();
     switch (bpv) {
       case CONTINUOUS_IDS:
-        readContinuousIds(in, count, docIDs);
+        readContinuousIds(in, count, scratch);
         break;
       case BITSET_IDS:
-        readBitSet(in, count, docIDs);
+        readBitSet(in, count, scratch);
         break;
       case DELTA_BPV_16:
-        readDelta16(in, count, docIDs);
+        readDelta16(in, count, scratch);
         break;
       case BPV_24:
-        readInts24(in, count, docIDs);
+        readInts24(in, count, scratch);
         break;
       case BPV_32:
-        readInts32(in, count, docIDs);
+        readInts32(in, count, scratch);
         break;
       case LEGACY_DELTA_VINT:
-        readLegacyDeltaVInts(in, count, docIDs);
+        readLegacyDeltaVInts(in, count, scratch);
         break;
       default:
         throw new IOException("Unsupported number of bits per value: " + bpv);
     }
+    return scratch;
   }
 
   private DocIdSetIterator readBitSetIterator(IndexInput in, int count) throws IOException {
diff --git a/lucene/core/src/test/org/apache/lucene/util/bkd/TestDocIdsWriter.java b/lucene/core/src/test/org/apache/lucene/util/bkd/TestDocIdsWriter.java
index 086eb460e56..0238b35ce0f 100644
--- a/lucene/core/src/test/org/apache/lucene/util/bkd/TestDocIdsWriter.java
+++ b/lucene/core/src/test/org/apache/lucene/util/bkd/TestDocIdsWriter.java
@@ -122,8 +122,7 @@ public class TestDocIdsWriter extends LuceneTestCase {
       }
     }
     try (IndexInput in = dir.openInput("tmp", IOContext.READONCE)) {
-      int[] read = new int[ints.length];
-      docIdsWriter.readInts(in, ints.length, read);
+      int[] read = docIdsWriter.readInts(in, ints.length);
       assertArrayEquals(ints, read);
       assertEquals(len, in.getFilePointer());
     }