Examples with ByteBufferAllocator org.apache.parquet.bytes.ByteBufferAllocator used on opensource projects

Example 1 with ByteBufferAllocator

use of org.apache.parquet.bytes.ByteBufferAllocator in project parquet-mr by apache.

the class TestDirectCodecFactory method test.

private void test(int size, CompressionCodecName codec, boolean useOnHeapCompression, Decompression decomp) {
    ByteBuffer rawBuf = null;
    ByteBuffer outBuf = null;
    ByteBufferAllocator allocator = null;
    try {
        allocator = new DirectByteBufferAllocator();
        final CodecFactory codecFactory = CodecFactory.createDirectCodecFactory(new Configuration(), allocator, pageSize);
        rawBuf = allocator.allocate(size);
        final byte[] rawArr = new byte[size];
        outBuf = allocator.allocate(size * 2);
        final Random r = new Random();
        final byte[] random = new byte[1024];
        int pos = 0;
        while (pos < size) {
            r.nextBytes(random);
            rawBuf.put(random);
            System.arraycopy(random, 0, rawArr, pos, random.length);
            pos += random.length;
        }
        rawBuf.flip();
        final DirectCodecFactory.BytesCompressor c = codecFactory.getCompressor(codec);
        final CodecFactory.BytesDecompressor d = codecFactory.getDecompressor(codec);
        final BytesInput compressed;
        if (useOnHeapCompression) {
            compressed = c.compress(BytesInput.from(rawArr));
        } else {
            compressed = c.compress(BytesInput.from(rawBuf));
        }
        switch(decomp) {
            case OFF_HEAP:
                {
                    final ByteBuffer buf = compressed.toByteBuffer();
                    final ByteBuffer b = allocator.allocate(buf.capacity());
                    try {
                        b.put(buf);
                        b.flip();
                        d.decompress(b, (int) compressed.size(), outBuf, size);
                        for (int i = 0; i < size; i++) {
                            Assert.assertTrue("Data didn't match at " + i, outBuf.get(i) == rawBuf.get(i));
                        }
                    } finally {
                        allocator.release(b);
                    }
                    break;
                }
            case OFF_HEAP_BYTES_INPUT:
                {
                    final ByteBuffer buf = compressed.toByteBuffer();
                    final ByteBuffer b = allocator.allocate(buf.limit());
                    try {
                        b.put(buf);
                        b.flip();
                        final BytesInput input = d.decompress(BytesInput.from(b), size);
                        Assert.assertArrayEquals(String.format("While testing codec %s", codec), input.toByteArray(), rawArr);
                    } finally {
                        allocator.release(b);
                    }
                    break;
                }
            case ON_HEAP:
                {
                    final byte[] buf = compressed.toByteArray();
                    final BytesInput input = d.decompress(BytesInput.from(buf), size);
                    Assert.assertArrayEquals(input.toByteArray(), rawArr);
                    break;
                }
        }
    } catch (Exception e) {
        final String msg = String.format("Failure while testing Codec: %s, OnHeapCompressionInput: %s, Decompression Mode: %s, Data Size: %d", codec.name(), useOnHeapCompression, decomp.name(), size);
        System.out.println(msg);
        throw new RuntimeException(msg, e);
    } finally {
        if (rawBuf != null) {
            allocator.release(rawBuf);
        }
        if (outBuf != null) {
            allocator.release(rawBuf);
        }
    }
}