Optimize PForUtil.encode() with histogram-based bit selection (#15165)

Author: RamakrishnaChilaka

lucene/CHANGES.txt

@@ -131,6 +131,8 @@ Optimizations
 
 * GITHUB#14998: Speed up flushing of live docs. (Adrien Grand)
 
+* GITHUB#15165: Optimize PForUtil.encode() with histogram-based bit selection. (Ramakrishna Chilaka)
+
 * GITHUB#15151: Use `SimScorer#score` bulk API to compute impact scores per
   block of postings. (Adrien Grand)
 

lucene/core/src/java/org/apache/lucene/codecs/lucene103/PForUtil.java

@@ -21,7 +21,6 @@
 import org.apache.lucene.internal.vectorization.PostingDecodingUtil;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.DataOutput;
-import org.apache.lucene.util.LongHeap;
 import org.apache.lucene.util.packed.PackedInts;
 
 /** Utility class to encode sequences of 128 small positive integers. */
@@ -46,34 +45,32 @@ static boolean allEqual(int[] l) {
 
   /** Encode 128 integers from {@code ints} into {@code out}. */
   void encode(int[] ints, DataOutput out) throws IOException {
-    // Determine the top MAX_EXCEPTIONS + 1 values
-    final LongHeap top = new LongHeap(MAX_EXCEPTIONS + 1);
-    for (int i = 0; i <= MAX_EXCEPTIONS; ++i) {
-      top.push(ints[i]);
-    }
-    long topValue = top.top();
-    for (int i = MAX_EXCEPTIONS + 1; i < ForUtil.BLOCK_SIZE; ++i) {
-      if (ints[i] > topValue) {
-        topValue = top.updateTop(ints[i]);
-      }
-    }
-
-    long max = 0L;
-    for (int i = 1; i <= top.size(); ++i) {
-      max = Math.max(max, top.get(i));
+    // histogram of bit widths
+    final int[] histogram = new int[32];
+    int maxBitsRequired = 0;
+    for (int i = 0; i < ForUtil.BLOCK_SIZE; ++i) {
+      final int v = ints[i];
+      final int bits = PackedInts.bitsRequired(v);
+      histogram[bits]++;
+      maxBitsRequired = Math.max(maxBitsRequired, bits);
     }
 
-    final int maxBitsRequired = PackedInts.bitsRequired(max);
-    // We store the patch on a byte, so we can't decrease the number of bits required by more than 8
-    final int patchedBitsRequired =
-        Math.max(PackedInts.bitsRequired(topValue), maxBitsRequired - 8);
+    // We store patch on a byte, so we can't decrease bits by more than 8
+    final int minBits = Math.max(0, maxBitsRequired - 8);
+    int cumulativeExceptions = 0;
+    int patchedBitsRequired = maxBitsRequired;
     int numExceptions = 0;
-    final long maxUnpatchedValue = (1L << patchedBitsRequired) - 1;
-    for (int i = 2; i <= top.size(); ++i) {
-      if (top.get(i) > maxUnpatchedValue) {
-        numExceptions++;
+
+    for (int b = maxBitsRequired; b >= minBits; --b) {
+      if (cumulativeExceptions > MAX_EXCEPTIONS) {
+        break;
       }
+      patchedBitsRequired = b;
+      numExceptions = cumulativeExceptions;
+      cumulativeExceptions += histogram[b];
     }
+
+    final int maxUnpatchedValue = (1 << patchedBitsRequired) - 1;
     final byte[] exceptions = new byte[numExceptions * 2];
     if (numExceptions > 0) {
       int exceptionCount = 0;
@@ -91,7 +88,7 @@ void encode(int[] ints, DataOutput out) throws IOException {
     if (allEqual(ints) && maxBitsRequired <= 8) {
       for (int i = 0; i < numExceptions; ++i) {
         exceptions[2 * i + 1] =
-            (byte) (Byte.toUnsignedLong(exceptions[2 * i + 1]) << patchedBitsRequired);
+            (byte) (Byte.toUnsignedInt(exceptions[2 * i + 1]) << patchedBitsRequired);
       }
       out.writeByte((byte) (numExceptions << 5));
       out.writeVInt(ints[0]);
@@ -115,7 +112,7 @@ void decode(PostingDecodingUtil pdu, int[] ints) throws IOException {
     }
     final int numExceptions = token >>> 5;
     for (int i = 0; i < numExceptions; ++i) {
-      ints[Byte.toUnsignedInt(in.readByte())] |= Byte.toUnsignedLong(in.readByte()) << bitsPerValue;
+      ints[Byte.toUnsignedInt(in.readByte())] |= Byte.toUnsignedInt(in.readByte()) << bitsPerValue;
     }
   }
 

lucene/core/src/java/org/apache/lucene/util/packed/PackedInts.java

@@ -792,6 +792,33 @@ public static int unsignedBitsRequired(long bits) {
     return Math.max(1, 64 - Long.numberOfLeadingZeros(bits));
   }
 
+  /**
+   * Returns how many bits are required to hold values up to and including maxValue NOTE: This
+   * method returns at least 1.
+   *
+   * @param maxValue the maximum int value that should be representable.
+   * @return the amount of bits needed to represent values from 0 to maxValue.
+   * @lucene.internal
+   */
+  public static int bitsRequired(int maxValue) {
+    if (maxValue < 0) {
+      throw new IllegalArgumentException("maxValue must be non-negative (got: " + maxValue + ")");
+    }
+    return unsignedBitsRequired(maxValue);
+  }
+
+  /**
+   * Returns how many bits are required to store <code>bits</code>, interpreted as an unsigned
+   * value. NOTE: This method returns at least 1.
+   *
+   * @param bits the int value to be stored, interpreted as unsigned.
+   * @return the amount of bits needed to represent the unsigned value.
+   * @lucene.internal
+   */
+  public static int unsignedBitsRequired(int bits) {
+    return Math.max(1, 32 - Integer.numberOfLeadingZeros(bits));
+  }
+
   /**
    * Calculates the maximum unsigned long that can be expressed with the given number of bits.
    *

lucene/core/src/test/org/apache/lucene/util/packed/TestPackedInts.java

@@ -75,8 +75,17 @@ public void testBitsRequired() {
     assertEquals(61, PackedInts.bitsRequired(0x1FFFFFFFFFFFFFFFL));
     assertEquals(62, PackedInts.bitsRequired(0x3FFFFFFFFFFFFFFFL));
     assertEquals(63, PackedInts.bitsRequired(0x7FFFFFFFFFFFFFFFL));
-    assertEquals(64, PackedInts.unsignedBitsRequired(-1));
+    assertEquals(64, PackedInts.unsignedBitsRequired(-1L));
     assertEquals(64, PackedInts.unsignedBitsRequired(Long.MIN_VALUE));
+    assertEquals(1, PackedInts.bitsRequired(0L));
+  }
+
+  public void testBitsRequiredInt() {
+    assertEquals(29, PackedInts.bitsRequired((int) Math.pow(2, 29) - 1));
+    assertEquals(30, PackedInts.bitsRequired(0x3FFFFFFF));
+    assertEquals(31, PackedInts.bitsRequired(0x7FFFFFFF));
+    assertEquals(32, PackedInts.unsignedBitsRequired(-1));
+    assertEquals(32, PackedInts.unsignedBitsRequired(Integer.MIN_VALUE));
     assertEquals(1, PackedInts.bitsRequired(0));
   }