GH-3466 Improve RunLengthBitPackingHybridDecoder.readNext to avoid per-call buffer allocation and DataInputStream wrapping

Author: arouel

parquet-column/src/main/java/org/apache/parquet/column/values/rle/RunLengthBitPackingHybridDecoder.java

@@ -18,7 +18,7 @@
  */
 package org.apache.parquet.column.values.rle;
 
-import java.io.DataInputStream;
+import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import org.apache.parquet.Preconditions;
@@ -48,6 +48,8 @@ private static enum MODE {
   private int currentCount;
   private int currentValue;
   private int[] currentBuffer;
+  private int currentBufferLength;
+  private byte[] packedBytes;
 
   public RunLengthBitPackingHybridDecoder(int bitWidth, InputStream in) {
     LOG.debug("decoding bitWidth {}", bitWidth);
@@ -69,7 +71,7 @@ public int readInt() throws IOException {
         result = currentValue;
         break;
       case PACKED:
-        result = currentBuffer[currentBuffer.length - 1 - currentCount];
+        result = currentBuffer[currentBufferLength - 1 - currentCount];
         break;
       default:
         throw new ParquetDecodingException("not a valid mode " + mode);
@@ -90,21 +92,39 @@ private void readNext() throws IOException {
       case PACKED:
         int numGroups = header >>> 1;
         currentCount = numGroups * 8;
+        currentBufferLength = currentCount;
         LOG.debug("reading {} values BIT PACKED", currentCount);
-        currentBuffer = new int[currentCount]; // TODO: reuse a buffer
-        byte[] bytes = new byte[numGroups * bitWidth];
+        if (currentBuffer == null || currentBuffer.length < currentCount) {
+          currentBuffer = new int[currentCount];
+        }
+        int bytesNeeded = numGroups * bitWidth;
+        if (packedBytes == null || packedBytes.length < bytesNeeded) {
+          packedBytes = new byte[bytesNeeded];
+        }
         // At the end of the file RLE data though, there might not be that many bytes left.
         int bytesToRead = (int) Math.ceil(currentCount * bitWidth / 8.0);
         bytesToRead = Math.min(bytesToRead, in.available());
-        new DataInputStream(in).readFully(bytes, 0, bytesToRead);
+        readFully(in, packedBytes, bytesToRead);
         for (int valueIndex = 0, byteIndex = 0;
             valueIndex < currentCount;
             valueIndex += 8, byteIndex += bitWidth) {
-          packer.unpack8Values(bytes, byteIndex, currentBuffer, valueIndex);
+          packer.unpack8Values(packedBytes, byteIndex, currentBuffer, valueIndex);
         }
         break;
       default:
         throw new ParquetDecodingException("not a valid mode " + mode);
     }
   }
+
+  private static void readFully(InputStream in, byte[] buf, int len) throws IOException {
+    int offset = 0;
+    while (offset < len) {
+      int read = in.read(buf, offset, len - offset);
+      if (read < 0) {
+        throw new EOFException(
+            "Unexpected end of stream: still needed " + (len - offset) + " bytes");
+      }
+      offset += read;
+    }
+  }
 }