feature: Add bit-shuffling interfaces for unshuffle with provided output array (#608)

tianjiqx · web-flow · commit 4277fbcdc86a · 2025-02-03T21:27:43.000-08:00
diff --git a/src/main/java/org/xerial/snappy/BitShuffle.java b/src/main/java/org/xerial/snappy/BitShuffle.java
@@ -208,6 +208,35 @@ public static int unshuffle(ByteBuffer shuffled, BitShuffleType type, ByteBuffer
         return numProcessed;
     }
 
+    /**
+     * Convert the input bit-shuffled byte array into an original byte array.
+     *
+     * @param input
+     * @return a byte array
+     * @throws IOException
+     */
+    public static byte[] unshuffleByteArray(byte[] input) throws IOException {
+        byte[] output = new byte[input.length];
+        int numProcessed = impl.unshuffle(input, 0, 1, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return output;
+    }
+
+    /**
+     * Convert the input bit-shuffled byte array into an original byte array.
+     *
+     * @param input
+     * @param output
+     * @return byte size of the unshuffled data.
+     * @throws IOException
+     */
+    public static int unshuffleByteArray(byte[] input, byte[] output) throws IOException {
+        assert(input.length == output.length);
+        int numProcessed = impl.unshuffle(input, 0, 1, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return numProcessed;
+    }
+
     /**
      * Convert the input bit-shuffled byte array into an original short array.
      *
@@ -222,6 +251,21 @@ public static short[] unshuffleShortArray(byte[] input) throws IOException {
         return output;
     }
 
+    /**
+     * Convert the input bit-shuffled byte array into an original short array.
+     *
+     * @param input
+     * @param output
+     * @return byte size of the unshuffled data.
+     * @throws IOException
+     */
+    public static int unshuffleShortArray(byte[] input, short[] output) throws IOException {
+        assert(input.length == output.length * 2);
+        int numProcessed = impl.unshuffle(input, 0, 2, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return numProcessed;
+    }
+
     /**
      * Convert the input bit-shuffled byte array into an original int array.
      *
@@ -236,6 +280,21 @@ public static int[] unshuffleIntArray(byte[] input) throws IOException {
         return output;
     }
 
+    /**
+     * Convert the input bit-shuffled byte array into an original int array.
+     *
+     * @param input
+     * @param output
+     * @return byte size of the unshuffled data.
+     * @throws IOException
+     */
+    public static int unshuffleIntArray(byte[] input, int[] output) throws IOException {
+        assert(input.length == output.length * 4);
+        int numProcessed = impl.unshuffle(input, 0, 4, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return numProcessed;
+    }
+
     /**
      * Convert the input bit-shuffled byte array into an original long array.
      *
@@ -250,6 +309,21 @@ public static long[] unshuffleLongArray(byte[] input) throws IOException {
         return output;
     }
 
+    /**
+     * Convert the input bit-shuffled byte array into an original long array.
+     *
+     * @param input
+     * @param output
+     * @return byte size of the unshuffled data.
+     * @throws IOException
+     */
+    public static int unshuffleLongArray(byte[] input, long[] output) throws IOException {
+        assert(input.length == output.length * 8);
+        int numProcessed = impl.unshuffle(input, 0, 8, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return numProcessed;
+    }
+
     /**
      * Convert the input bit-shuffled byte array into an original float array.
      *
@@ -264,6 +338,21 @@ public static float[] unshuffleFloatArray(byte[] input) throws IOException {
         return output;
     }
 
+    /**
+     * Convert the input bit-shuffled byte array into an original float array.
+     *
+     * @param input
+     * @param output
+     * @return byte size of the unshuffled data.
+     * @throws IOException
+     */
+    public static int unshuffleFloatArray(byte[] input, float[] output) throws IOException {
+        assert(input.length == output.length * 4);
+        int numProcessed = impl.unshuffle(input, 0, 4, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return numProcessed;
+    }
+
     /**
      * Convert the input bit-shuffled byte array into an original double array.
      *
@@ -277,4 +366,19 @@ public static double[] unshuffleDoubleArray(byte[] input) throws IOException {
         assert(numProcessed == input.length);
         return output;
     }
+
+    /**
+     * Convert the input bit-shuffled byte array into an original double array.
+     *
+     * @param input
+     * @param output
+     * @return byte size of the unshuffled data.
+     * @throws IOException
+     */
+    public static int unshuffleDoubleArray(byte[] input, double[] output) throws IOException {
+        assert(input.length == output.length * 8);
+        int numProcessed = impl.unshuffle(input, 0, 8, input.length, output, 0);
+        assert(numProcessed == input.length);
+        return numProcessed;
+    }
 }
diff --git a/src/test/java/org/xerial/snappy/BitShuffleTest.java b/src/test/java/org/xerial/snappy/BitShuffleTest.java
@@ -276,6 +276,22 @@ public void shuffleLongArray()
         byte[] shuffledData = BitShuffle.shuffle(data);
         long[] result = BitShuffle.unshuffleLongArray(shuffledData);
         assertArrayEquals(data, result);
+        long[] output = new long[data.length];
+        BitShuffle.unshuffleLongArray(shuffledData, output);
+        assertArrayEquals(data, output);
+    }
+
+    @Test
+    public void shuffleByteArray()
+            throws Exception
+    {
+        byte[] data = new byte[] {43, -32, 1, 3, 34, 43, 34, Byte.MAX_VALUE, -1};
+        byte[] shuffledData = BitShuffle.shuffle(data);
+        byte[] result = BitShuffle.unshuffleByteArray(shuffledData);
+        assertArrayEquals(data, result);
+        byte[] output = new byte[data.length];
+        BitShuffle.unshuffleByteArray(shuffledData, output);
+        assertArrayEquals(data, output);
     }
 
     @Test
@@ -286,6 +302,9 @@ public void shuffleShortArray()
         byte[] shuffledData = BitShuffle.shuffle(data);
         short[] result = BitShuffle.unshuffleShortArray(shuffledData);
         assertArrayEquals(data, result);
+        short[] output = new short[data.length];
+        BitShuffle.unshuffleShortArray(shuffledData, output);
+        assertArrayEquals(data, output);
     }
 
     @Test
@@ -296,6 +315,9 @@ public void shuffleIntArray()
         byte[] shuffledData = BitShuffle.shuffle(data);
         int[] result = BitShuffle.unshuffleIntArray(shuffledData);
         assertArrayEquals(data, result);
+        int[] output = new int[data.length];
+        BitShuffle.unshuffleIntArray(shuffledData, output);
+        assertArrayEquals(data, output);
     }
 
     @Test
@@ -306,6 +328,9 @@ public void shuffleFloatArray()
         byte[] shuffledData = BitShuffle.shuffle(data);
         float[] result = BitShuffle.unshuffleFloatArray(shuffledData);
         assertArrayEquals(data, result, 0.0000001f);
+        float[] output = new float[data.length];
+        BitShuffle.unshuffleFloatArray(shuffledData, output);
+        assertArrayEquals(data, output);
     }
 
     @Test
@@ -316,5 +341,8 @@ public void shuffleDoubleArray()
         byte[] shuffledData = BitShuffle.shuffle(data);
         double[] result = BitShuffle.unshuffleDoubleArray(shuffledData);
         assertArrayEquals(data, result, 0.0000001f);
+        double[] output = new double[data.length];
+        BitShuffle.unshuffleDoubleArray(shuffledData, output);
+        assertArrayEquals(data, output);
     }
 }
