Rewrite of the GroupVInt optimization without lambdas, varhandles and no code in subclasses (#15116)

* Rewrite of the GroupVInt optimization without lamdas, varhandles and no code in subclasses

* Remove optimized flag and duplicate code

* apply suggestions/fixes: fix maximum size; cleanup baseline method and hide internal details

* Apply suggestions from code review

Co-authored-by: Adrien Grand <[email protected]>

---------

Co-authored-by: Adrien Grand <[email protected]>

Author: uschindler

lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/GroupVIntBenchmark.java

@@ -23,9 +23,8 @@
 import java.util.concurrent.TimeUnit;
 import org.apache.lucene.store.ByteArrayDataInput;
 import org.apache.lucene.store.ByteArrayDataOutput;
-import org.apache.lucene.store.ByteBuffersDataInput;
 import org.apache.lucene.store.ByteBuffersDataOutput;
-import org.apache.lucene.store.DataInput;
+import org.apache.lucene.store.ByteBuffersDirectory;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
@@ -50,8 +49,8 @@
 @BenchmarkMode(Mode.Throughput)
 @OutputTimeUnit(TimeUnit.MICROSECONDS)
 @State(Scope.Benchmark)
-@Warmup(iterations = 3, time = 3)
-@Measurement(iterations = 5, time = 5)
+@Warmup(iterations = 4, time = 8)
+@Measurement(iterations = 5, time = 20)
 @Fork(
     value = 1,
     jvmArgsPrepend = {"--add-modules=jdk.unsupported"})
@@ -92,7 +91,7 @@ public class GroupVIntBenchmark {
   IndexInput mmapGVIntIn;
   IndexInput nioGVIntIn;
   IndexInput mmapVIntIn;
-  ByteBuffersDataInput byteBuffersGVIntIn;
+  IndexInput byteBuffersGVIntIn;
 
   ByteArrayDataInput byteArrayVIntIn;
   ByteArrayDataInput byteArrayGVIntIn;
@@ -125,9 +124,11 @@ void initNioInput(int[] docs) throws Exception {
   }
 
   void initByteBuffersInput(int[] docs) throws Exception {
-    ByteBuffersDataOutput buffer = new ByteBuffersDataOutput();
-    buffer.writeGroupVInts(docs, docs.length);
-    byteBuffersGVIntIn = buffer.toDataInput();
+    Directory dir = new ByteBuffersDirectory();
+    IndexOutput out = dir.createOutput("gvint", IOContext.DEFAULT);
+    out.writeGroupVInts(docs, docs.length);
+    out.close();
+    byteBuffersGVIntIn = dir.openInput("gvint", IOContext.DEFAULT);
   }
 
   void initMMapInput(int[] docs) throws Exception {
@@ -145,16 +146,6 @@ void initMMapInput(int[] docs) throws Exception {
     mmapVIntIn = dir.openInput("vint", IOContext.DEFAULT);
   }
 
-  private void readGroupVIntsBaseline(DataInput in, int[] dst, int limit) throws IOException {
-    int i;
-    for (i = 0; i <= limit - 4; i += 4) {
-      GroupVIntUtil.readGroupVInt(in, dst, i);
-    }
-    for (; i < limit; ++i) {
-      dst[i] = in.readVInt();
-    }
-  }
-
   @Setup(Level.Trial)
   public void init() throws Exception {
     Random r = new Random(0);
@@ -193,7 +184,7 @@ public void benchMMapDirectoryInputs_readGroupVInt(Blackhole bh) throws IOExcept
   @Benchmark
   public void benchMMapDirectoryInputs_readGroupVIntBaseline(Blackhole bh) throws IOException {
     mmapGVIntIn.seek(0);
-    this.readGroupVIntsBaseline(mmapGVIntIn, values, size);
+    GroupVIntUtil.readGroupVInts$Baseline(mmapGVIntIn, values, size);
     bh.consume(values);
   }
 
@@ -223,7 +214,7 @@ public void benchNIOFSDirectoryInputs_readGroupVInt(Blackhole bh) throws IOExcep
   @Benchmark
   public void benchNIOFSDirectoryInputs_readGroupVIntBaseline(Blackhole bh) throws IOException {
     nioGVIntIn.seek(0);
-    this.readGroupVIntsBaseline(nioGVIntIn, values, size);
+    GroupVIntUtil.readGroupVInts$Baseline(nioGVIntIn, values, size);
     bh.consume(values);
   }
 
@@ -237,7 +228,7 @@ public void benchByteBuffersIndexInput_readGroupVInt(Blackhole bh) throws IOExce
   @Benchmark
   public void benchByteBuffersIndexInput_readGroupVIntBaseline(Blackhole bh) throws IOException {
     byteBuffersGVIntIn.seek(0);
-    this.readGroupVIntsBaseline(byteBuffersGVIntIn, values, size);
+    GroupVIntUtil.readGroupVInts$Baseline(byteBuffersGVIntIn, values, size);
     bh.consume(values);
   }
 

lucene/core/src/java/org/apache/lucene/store/BufferedIndexInput.java

@@ -20,7 +20,6 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
-import org.apache.lucene.util.GroupVIntUtil;
 
 /** Base implementation class for buffered {@link IndexInput}. */
 public abstract class BufferedIndexInput extends IndexInput implements RandomAccessInput {
@@ -150,16 +149,6 @@ public final int readInt() throws IOException {
     }
   }
 
-  @Override
-  public void readGroupVInt(int[] dst, int offset) throws IOException {
-    final int len =
-        GroupVIntUtil.readGroupVInt(
-            this, buffer.remaining(), p -> buffer.getInt((int) p), buffer.position(), dst, offset);
-    if (len > 0) {
-      buffer.position(buffer.position() + len);
-    }
-  }
-
   @Override
   public final long readLong() throws IOException {
     if (Long.BYTES <= buffer.remaining()) {

lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataInput.java

@@ -29,7 +29,6 @@
 import java.util.Locale;
 import java.util.stream.Collectors;
 import org.apache.lucene.util.Accountable;
-import org.apache.lucene.util.GroupVIntUtil;
 import org.apache.lucene.util.RamUsageEstimator;
 
 /**
@@ -203,25 +202,6 @@ public long readLong() throws IOException {
     }
   }
 
-  @Override
-  public void readGroupVInt(int[] dst, int offset) throws IOException {
-    final ByteBuffer block = blocks[blockIndex(pos)];
-    final int blockOffset = blockOffset(pos);
-    // We MUST save the return value to local variable, could not use pos += readGroupVInt(...).
-    // because `pos +=` in java will move current value(not address) of pos to register first,
-    // then call the function, but we will update pos value in function via readByte(), then
-    // `pos +=` will use an old pos value plus return value, thereby missing 1 byte.
-    final int len =
-        GroupVIntUtil.readGroupVInt(
-            this,
-            block.limit() - blockOffset,
-            p -> block.getInt((int) p),
-            blockOffset,
-            dst,
-            offset);
-    pos += len;
-  }
-
   @Override
   public long length() {
     return length;

lucene/core/src/java/org/apache/lucene/store/ByteBuffersIndexInput.java

@@ -205,12 +205,6 @@ public void readLongs(long[] dst, int offset, int length) throws IOException {
     in.readLongs(dst, offset, length);
   }
 
-  @Override
-  public void readGroupVInt(int[] dst, int offset) throws IOException {
-    ensureOpen();
-    in.readGroupVInt(dst, offset);
-  }
-
   @Override
   public IndexInput clone() {
     ensureOpen();

lucene/core/src/java/org/apache/lucene/store/DataInput.java

@@ -27,7 +27,6 @@
 import java.util.TreeMap;
 import java.util.TreeSet;
 import org.apache.lucene.util.BitUtil;
-import org.apache.lucene.util.GroupVIntUtil;
 
 /**
  * Abstract base class for performing read operations of Lucene's low-level data types.
@@ -99,16 +98,6 @@ public int readInt() throws IOException {
     return ((b4 & 0xFF) << 24) | ((b3 & 0xFF) << 16) | ((b2 & 0xFF) << 8) | (b1 & 0xFF);
   }
 
-  /**
-   * Override if you have an efficient implementation. In general this is when the input supports
-   * random access.
-   *
-   * @lucene.experimental
-   */
-  public void readGroupVInt(int[] dst, int offset) throws IOException {
-    GroupVIntUtil.readGroupVInt(this, dst, offset);
-  }
-
   /**
    * Reads an int stored in variable-length format. Reads between one and five bytes. Smaller values
    * take fewer bytes. Negative numbers are supported, but should be avoided.

lucene/core/src/java/org/apache/lucene/store/MemorySegmentIndexInput.java

@@ -31,7 +31,6 @@
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BitUtil;
 import org.apache.lucene.util.Constants;
-import org.apache.lucene.util.GroupVIntUtil;
 import org.apache.lucene.util.IOConsumer;
 
 /**
@@ -429,23 +428,6 @@ public byte readByte(long pos) throws IOException {
     }
   }
 
-  @Override
-  public void readGroupVInt(int[] dst, int offset) throws IOException {
-    try {
-      final int len =
-          GroupVIntUtil.readGroupVInt(
-              this,
-              curSegment.byteSize() - curPosition,
-              p -> curSegment.get(LAYOUT_LE_INT, p),
-              curPosition,
-              dst,
-              offset);
-      curPosition += len;
-    } catch (NullPointerException | IllegalStateException e) {
-      throw alreadyClosed(e);
-    }
-  }
-
   @Override
   public void readBytes(long pos, byte[] b, int offset, int len) throws IOException {
     try {

lucene/core/src/java/org/apache/lucene/util/GroupVIntUtil.java

@@ -19,15 +19,17 @@
 import java.io.IOException;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.DataOutput;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.RandomAccessInput;
 
 /**
  * This class contains utility methods and constants for group varint
  *
  * @lucene.internal
  */
 public final class GroupVIntUtil {
-  // the maximum length of a single group-varint is 4 integers + 1 byte flag.
-  public static final int MAX_LENGTH_PER_GROUP = 17;
+  // the maximum length of a single group-varint is 1 byte flag and 4 integers.
+  public static final int MAX_LENGTH_PER_GROUP = Byte.BYTES + 4 * Integer.BYTES;
 
   private static final int[] INT_MASKS = new int[] {0xFF, 0xFFFF, 0xFFFFFF, ~0};
 
@@ -41,7 +43,7 @@ public final class GroupVIntUtil {
   public static void readGroupVInts(DataInput in, int[] dst, int limit) throws IOException {
     int i;
     for (i = 0; i <= limit - 4; i += 4) {
-      in.readGroupVInt(dst, i);
+      readGroupVInt(in, dst, i);
     }
     for (; i < limit; ++i) {
       dst[i] = in.readVInt();
@@ -56,7 +58,54 @@ public static void readGroupVInts(DataInput in, int[] dst, int limit) throws IOE
    * @param dst the array to read ints into.
    * @param offset the offset in the array to start storing ints.
    */
-  public static void readGroupVInt(DataInput in, int[] dst, int offset) throws IOException {
+  private static void readGroupVInt(DataInput in, int[] dst, int offset) throws IOException {
+    final int flag = in.readByte() & 0xFF;
+
+    final int n1Minus1 = flag >> 6;
+    final int n2Minus1 = (flag >> 4) & 0x03;
+    final int n3Minus1 = (flag >> 2) & 0x03;
+    final int n4Minus1 = flag & 0x03;
+
+    // if our DataInput implements RandomAccessInput for absolute access and IndexInput for seeking,
+    // we use a branch-less implementation:
+    if (in instanceof RandomAccessInput rin && in instanceof IndexInput iin) {
+      long pos = iin.getFilePointer();
+      if (iin.length() - pos >= 4 * Integer.BYTES) {
+        dst[offset] = rin.readInt(pos) & INT_MASKS[n1Minus1];
+        pos += 1 + n1Minus1;
+        dst[offset + 1] = rin.readInt(pos) & INT_MASKS[n2Minus1];
+        pos += 1 + n2Minus1;
+        dst[offset + 2] = rin.readInt(pos) & INT_MASKS[n3Minus1];
+        pos += 1 + n3Minus1;
+        dst[offset + 3] = rin.readInt(pos) & INT_MASKS[n4Minus1];
+        pos += 1 + n4Minus1;
+
+        iin.seek(pos);
+        return;
+      }
+    }
+
+    // fall-through: default impl
+    dst[offset] = readIntInGroup(in, n1Minus1);
+    dst[offset + 1] = readIntInGroup(in, n2Minus1);
+    dst[offset + 2] = readIntInGroup(in, n3Minus1);
+    dst[offset + 3] = readIntInGroup(in, n4Minus1);
+  }
+
+  /** DO not use! Only visible for benchmarking purposes! */
+  public static void readGroupVInts$Baseline(DataInput in, int[] dst, int limit)
+      throws IOException {
+    int i;
+    for (i = 0; i <= limit - 4; i += 4) {
+      readGroupVInt$Baseline(in, dst, i);
+    }
+    for (; i < limit; ++i) {
+      dst[i] = in.readVInt();
+    }
+  }
+
+  private static void readGroupVInt$Baseline(DataInput in, int[] dst, int offset)
+      throws IOException {
     final int flag = in.readByte() & 0xFF;
 
     final int n1Minus1 = flag >> 6;
@@ -83,55 +132,6 @@ private static int readIntInGroup(DataInput in, int numBytesMinus1) throws IOExc
     }
   }
 
-  /**
-   * Provides an abstraction for read int values, so that decoding logic can be reused in different
-   * DataInput.
-   */
-  @FunctionalInterface
-  public static interface IntReader {
-    int read(long v);
-  }
-
-  /**
-   * Faster implementation of read single group, It read values from the buffer that would not cross
-   * boundaries.
-   *
-   * @param in the input to use to read data.
-   * @param remaining the number of remaining bytes allowed to read for current block/segment.
-   * @param reader the supplier of read int.
-   * @param pos the start pos to read from the reader.
-   * @param dst the array to read ints into.
-   * @param offset the offset in the array to start storing ints.
-   * @return the number of bytes read excluding the flag. this indicates the number of positions
-   *     should to be increased for caller, it is 0 or positive number and less than {@link
-   *     #MAX_LENGTH_PER_GROUP}
-   */
-  public static int readGroupVInt(
-      DataInput in, long remaining, IntReader reader, long pos, int[] dst, int offset)
-      throws IOException {
-    if (remaining < MAX_LENGTH_PER_GROUP) {
-      readGroupVInt(in, dst, offset);
-      return 0;
-    }
-    final int flag = in.readByte() & 0xFF;
-    final long posStart = ++pos; // exclude the flag bytes, the position has updated via readByte().
-    final int n1Minus1 = flag >> 6;
-    final int n2Minus1 = (flag >> 4) & 0x03;
-    final int n3Minus1 = (flag >> 2) & 0x03;
-    final int n4Minus1 = flag & 0x03;
-
-    // This code path has fewer conditionals and tends to be significantly faster in benchmarks
-    dst[offset] = reader.read(pos) & INT_MASKS[n1Minus1];
-    pos += 1 + n1Minus1;
-    dst[offset + 1] = reader.read(pos) & INT_MASKS[n2Minus1];
-    pos += 1 + n2Minus1;
-    dst[offset + 2] = reader.read(pos) & INT_MASKS[n3Minus1];
-    pos += 1 + n3Minus1;
-    dst[offset + 3] = reader.read(pos) & INT_MASKS[n4Minus1];
-    pos += 1 + n4Minus1;
-    return (int) (pos - posStart);
-  }
-
   private static int numBytes(int v) {
     // | 1 to return 1 when v = 0
     return Integer.BYTES - (Integer.numberOfLeadingZeros(v | 1) >> 3);