Add native Zstd input and output streams

Author: dain

src/main/java/io/airlift/compress/v3/zstd/ZstdHadoopInputStream.java

@@ -24,45 +24,67 @@ class ZstdHadoopInputStream
         extends HadoopInputStream
 {
     private final InputStream in;
+    private final boolean useNative;
     private ZstdInputStream zstdInputStream;
 
-    public ZstdHadoopInputStream(InputStream in)
+    public ZstdHadoopInputStream(InputStream in, boolean useNative)
     {
         this.in = requireNonNull(in, "in is null");
-        zstdInputStream = new ZstdInputStream(in);
+        this.useNative = useNative;
+    }
+
+    private void createDecompressingStreamIfNecessary()
+            throws IOException
+    {
+        if (zstdInputStream == null) {
+            if (useNative) {
+                zstdInputStream = new ZstdNativeInputStream(in);
+            }
+            else {
+                zstdInputStream = new ZstdJavaInputStream(in);
+            }
+        }
     }
 
     @Override
     public int read()
             throws IOException
     {
+        createDecompressingStreamIfNecessary();
         return zstdInputStream.read();
     }
 
     @Override
     public int read(byte[] b)
             throws IOException
     {
+        createDecompressingStreamIfNecessary();
         return zstdInputStream.read(b);
     }
 
     @Override
     public int read(byte[] outputBuffer, int outputOffset, int outputLength)
             throws IOException
     {
+        createDecompressingStreamIfNecessary();
         return zstdInputStream.read(outputBuffer, outputOffset, outputLength);
     }
 
     @Override
     public void resetState()
     {
-        zstdInputStream = new ZstdInputStream(in);
+        zstdInputStream = null;
     }
 
     @Override
     public void close()
             throws IOException
     {
-        zstdInputStream.close();
+        if (zstdInputStream != null) {
+            zstdInputStream.close();
+        }
+        else {
+            in.close();
+        }
     }
 }

src/main/java/io/airlift/compress/v3/zstd/ZstdHadoopOutputStream.java

@@ -24,12 +24,13 @@ class ZstdHadoopOutputStream
         extends HadoopOutputStream
 {
     private final OutputStream out;
-    private boolean initialized;
+    private final boolean useNative;
     private ZstdOutputStream zstdOutputStream;
 
-    public ZstdHadoopOutputStream(OutputStream out)
+    public ZstdHadoopOutputStream(OutputStream out, boolean useNative)
     {
         this.out = requireNonNull(out, "out is null");
+        this.useNative = useNative;
     }
 
     @Override
@@ -70,10 +71,8 @@ public void close()
             throws IOException
     {
         try {
-            // it the stream has never been initialized, create a valid empty file
-            if (!initialized) {
-                openStreamIfNecessary();
-            }
+            // If the stream has never been initialized, create a valid empty file
+            openStreamIfNecessary();
             finish();
         }
         finally {
@@ -85,8 +84,12 @@ private void openStreamIfNecessary()
             throws IOException
     {
         if (zstdOutputStream == null) {
-            initialized = true;
-            zstdOutputStream = new ZstdOutputStream(out);
+            if (useNative) {
+                zstdOutputStream = new ZstdNativeOutputStream(out);
+            }
+            else {
+                zstdOutputStream = new ZstdJavaOutputStream(out);
+            }
         }
     }
 }

src/main/java/io/airlift/compress/v3/zstd/ZstdHadoopStreams.java

@@ -27,6 +27,18 @@
 public class ZstdHadoopStreams
         implements HadoopStreams
 {
+    private final boolean useNative;
+
+    public ZstdHadoopStreams()
+    {
+        this(true);  // Default to native when available
+    }
+
+    public ZstdHadoopStreams(boolean useNative)
+    {
+        this.useNative = useNative && ZstdNative.isEnabled();
+    }
+
     @Override
     public String getDefaultFileExtension()
     {
@@ -43,13 +55,13 @@ public List<String> getHadoopCodecName()
     public HadoopInputStream createInputStream(InputStream in)
             throws IOException
     {
-        return new ZstdHadoopInputStream(in);
+        return new ZstdHadoopInputStream(in, useNative);
     }
 
     @Override
     public HadoopOutputStream createOutputStream(OutputStream out)
             throws IOException
     {
-        return new ZstdHadoopOutputStream(out);
+        return new ZstdHadoopOutputStream(out, useNative);
     }
 }

src/main/java/io/airlift/compress/v3/zstd/ZstdInputStream.java

@@ -13,139 +13,14 @@
  */
 package io.airlift.compress.v3.zstd;
 
-import java.io.IOException;
 import java.io.InputStream;
-import java.util.Arrays;
 
-import static io.airlift.compress.v3.zstd.Util.checkPositionIndexes;
-import static io.airlift.compress.v3.zstd.Util.checkState;
-import static java.lang.Math.max;
-import static java.util.Objects.requireNonNull;
-import static sun.misc.Unsafe.ARRAY_BYTE_BASE_OFFSET;
-
-public class ZstdInputStream
+/**
+ * Abstract base class for zstd decompressing input streams.
+ * <p>
+ * This class provides a common type for both the pure Java implementation ({@link ZstdJavaInputStream})
+ * and the native implementation ({@link ZstdNativeInputStream}).
+ */
+public abstract sealed class ZstdInputStream
         extends InputStream
-{
-    private static final int MIN_BUFFER_SIZE = 4096;
-
-    private final InputStream inputStream;
-    private final ZstdIncrementalFrameDecompressor decompressor = new ZstdIncrementalFrameDecompressor();
-
-    private byte[] inputBuffer = new byte[decompressor.getInputRequired()];
-    private int inputBufferOffset;
-    private int inputBufferLimit;
-
-    private byte[] singleByteOutputBuffer;
-
-    private boolean closed;
-
-    public ZstdInputStream(InputStream inputStream)
-    {
-        this.inputStream = requireNonNull(inputStream, "inputStream is null");
-    }
-
-    @Override
-    public int read()
-            throws IOException
-    {
-        if (singleByteOutputBuffer == null) {
-            singleByteOutputBuffer = new byte[1];
-        }
-        int readSize = read(singleByteOutputBuffer, 0, 1);
-        checkState(readSize != 0, "A zero read size should never be returned");
-        if (readSize != 1) {
-            return -1;
-        }
-        return singleByteOutputBuffer[0] & 0xFF;
-    }
-
-    @Override
-    public int read(final byte[] outputBuffer, final int outputOffset, final int outputLength)
-            throws IOException
-    {
-        if (closed) {
-            throw new IOException("Stream is closed");
-        }
-
-        if (outputBuffer == null) {
-            throw new NullPointerException();
-        }
-        checkPositionIndexes(outputOffset, outputOffset + outputLength, outputBuffer.length);
-        if (outputLength == 0) {
-            return 0;
-        }
-
-        final int outputLimit = outputOffset + outputLength;
-        int outputUsed = 0;
-        while (outputUsed < outputLength) {
-            boolean enoughInput = fillInputBufferIfNecessary(decompressor.getInputRequired());
-            if (!enoughInput) {
-                if (decompressor.isAtStoppingPoint()) {
-                    return outputUsed > 0 ? outputUsed : -1;
-                }
-                throw new IOException("Not enough input bytes");
-            }
-
-            decompressor.partialDecompress(
-                    inputBuffer,
-                    inputBufferOffset + ARRAY_BYTE_BASE_OFFSET,
-                    inputBufferLimit + ARRAY_BYTE_BASE_OFFSET,
-                    outputBuffer,
-                    outputOffset + outputUsed,
-                    outputLimit);
-
-            inputBufferOffset += decompressor.getInputConsumed();
-            outputUsed += decompressor.getOutputBufferUsed();
-        }
-        return outputUsed;
-    }
-
-    private boolean fillInputBufferIfNecessary(int requiredSize)
-            throws IOException
-    {
-        if (inputBufferLimit - inputBufferOffset >= requiredSize) {
-            return true;
-        }
-
-        // compact existing buffered data to the front of the buffer
-        if (inputBufferOffset > 0) {
-            int copySize = inputBufferLimit - inputBufferOffset;
-            System.arraycopy(inputBuffer, inputBufferOffset, inputBuffer, 0, copySize);
-            inputBufferOffset = 0;
-            inputBufferLimit = copySize;
-        }
-
-        if (inputBuffer.length < requiredSize) {
-            inputBuffer = Arrays.copyOf(inputBuffer, max(requiredSize, MIN_BUFFER_SIZE));
-        }
-
-        while (inputBufferLimit < inputBuffer.length) {
-            int readSize = inputStream.read(inputBuffer, inputBufferLimit, inputBuffer.length - inputBufferLimit);
-            if (readSize < 0) {
-                break;
-            }
-            inputBufferLimit += readSize;
-        }
-        return inputBufferLimit >= requiredSize;
-    }
-
-    @Override
-    public int available()
-            throws IOException
-    {
-        if (closed) {
-            return 0;
-        }
-        return decompressor.getRequestedOutputSize();
-    }
-
-    @Override
-    public void close()
-            throws IOException
-    {
-        if (!closed) {
-            closed = true;
-            inputStream.close();
-        }
-    }
-}
+        permits ZstdJavaInputStream, ZstdNativeInputStream {}