com.github.luben.zstd.Zstd Java Examples

@Test
public void testMerge() {
    final String origin1 = "Benedict";
    final String origin2 = "Jin";

    final byte[] originBytes1 = origin1.getBytes(StandardCharsets.UTF_8);
    info(originBytes1);
    final byte[] originBytes2 = origin2.getBytes(StandardCharsets.UTF_8);
    info(originBytes2);

    final byte[] compressedBytes1 = Zstd.compress(originBytes1);
    info(compressedBytes1);
    final byte[] compressedBytes2 = Zstd.compress(originBytes2);
    info(compressedBytes2);

    final byte[] compressedBytes = ByteUtils.merge(compressedBytes1, compressedBytes2);
    info(compressedBytes);

    final long size = Zstd.decompressedSize(compressedBytes);
    System.out.println(size);
    System.out.println(compressedBytes.length);
    final byte[] decompress = Zstd.decompress(compressedBytes, (int) size * 10);
    Assert.assertEquals(origin1 + origin2, new String(decompress, StandardCharsets.UTF_8));
}
 

/**
 * Write the genotype probability data from the byte buffer in compressed form to the output channel.
 *
 * @param bgenOutputByteChannel The channel that is able to pass byte buffers to the output stream.
 * @param genotypeDataBlockBuffer The byte buffer that holds the genotype probability data.
 * @throws IOException If an I/O exception has occurred.
 * @return The number of bytes written.
 */
private long writeCompressedBgenGenotypeDataBlock(WritableByteChannel bgenOutputByteChannel,
												  ByteBuffer genotypeDataBlockBuffer) throws IOException {

	// Compress the genotype data using the Zstandard compression method.
	// Compression level 3 is used since this was the defualt when writing this.
	// (The byte buffer compression method does not have have an overloaded method with default compression level)
	ByteBuffer compressedGenotypeDataBuffer = Zstd.compress(genotypeDataBlockBuffer, 3);
	// If this was successful, the limit is equal to the position of the original buffer.
	// Check if the alternative condition applies.
	if (genotypeDataBlockBuffer.position() != genotypeDataBlockBuffer.limit()) {
		throw new GenotypeDataException("Zstd compression of genotype data failed, exiting");
	}

	// Get the length of the compressed variant probability data.
	int lengthOfVariantProbabilityData =
			compressedGenotypeDataBuffer.limit() + 4; // Add 4 since the D field takes 4 bytes
	// Get the decompressed length of variant probability data.
	int decompressedLengthOfVariantProbabilityData = genotypeDataBlockBuffer.capacity();

	// Write the (decompressed) length to a bytebuffer.
	ByteBuffer genotypeBlockLengths = ByteBuffer.allocate(8)
			.order(ByteOrder.LITTLE_ENDIAN);

	genotypeBlockLengths.putInt(lengthOfVariantProbabilityData);
	genotypeBlockLengths.putInt(decompressedLengthOfVariantProbabilityData);

	genotypeBlockLengths.flip(); // reset pointer

	// First write the lengths.
	int numberOfBytesWritten = bgenOutputByteChannel.write(genotypeBlockLengths);
	// Now write the actual data.
	return numberOfBytesWritten + bgenOutputByteChannel.write(compressedGenotypeDataBuffer);
}
 

private ChannelFuture finishEncode(final ChannelHandlerContext ctx, ChannelPromise promise) {
    if (finished) {
        promise.setSuccess();
        return promise;
    }
    finished = true;

    final ByteBuf footer = ctx.alloc().ioBuffer(
            (int) Zstd.compressBound(buffer.readableBytes()) + HEADER_LENGTH);
    flushBufferedData(footer);

    final int idx = footer.writerIndex();
    footer.setInt(idx, MAGIC_NUMBER);
    footer.setByte(idx + TOKEN_OFFSET, (byte) (BLOCK_TYPE_NON_COMPRESSED | compressionLevel));
    footer.setInt(idx + COMPRESSED_LENGTH_OFFSET, 0);
    footer.setInt(idx + DECOMPRESSED_LENGTH_OFFSET, 0);
    footer.setInt(idx + CHECKSUM_OFFSET, 0);

    footer.writerIndex(idx + HEADER_LENGTH);

    return ctx.writeAndFlush(footer, promise);
}
 

@Override
public ByteBuf encode(ByteBuf source) {
    int uncompressedLength = source.readableBytes();
    int maxLength = (int) Zstd.compressBound(uncompressedLength);

    ByteBuf target = PulsarByteBufAllocator.DEFAULT.directBuffer(maxLength, maxLength);
    int compressedLength;

    if (source.hasMemoryAddress()) {
        compressedLength = (int) Zstd.compressUnsafe(target.memoryAddress(), maxLength,
                source.memoryAddress() + source.readerIndex(),
                uncompressedLength, ZSTD_COMPRESSION_LEVEL);
    } else {
        ByteBuffer sourceNio = source.nioBuffer(source.readerIndex(), source.readableBytes());
        ByteBuffer targetNio = target.nioBuffer(0, maxLength);

        compressedLength = Zstd.compress(targetNio, sourceNio, ZSTD_COMPRESSION_LEVEL);
    }

    target.writerIndex(compressedLength);
    return target;
}
 

@Override
public ByteBuf decode(ByteBuf encoded, int uncompressedLength) throws IOException {
    ByteBuf uncompressed = PulsarByteBufAllocator.DEFAULT.directBuffer(uncompressedLength, uncompressedLength);

    if (encoded.hasMemoryAddress()) {
        Zstd.decompressUnsafe(uncompressed.memoryAddress(), uncompressedLength,
                encoded.memoryAddress() + encoded.readerIndex(),
                encoded.readableBytes());
    } else {
        ByteBuffer uncompressedNio = uncompressed.nioBuffer(0, uncompressedLength);
        ByteBuffer encodedNio = encoded.nioBuffer(encoded.readerIndex(), encoded.readableBytes());

        Zstd.decompress(uncompressedNio, encodedNio);
    }

    uncompressed.writerIndex(uncompressedLength);
    return uncompressed;
}
 

public ZstdInputStream(InputStream inStream) throws IOException {
    // FilterInputStream constructor
    super(inStream);

    // allocate input buffer with max frame header size
    src = new byte[srcBuffSize];
    if (src == null) {
        throw new IOException("Error allocating the input buffer of size " + srcBuffSize);
    }
    stream = createDStream();
    int size = initDStream(stream);
    if (Zstd.isError(size)) {
        throw new IOException("Decompression error: " + Zstd.getErrorName(size));
    }
}
 

@Test
public void testMixed() {
    final String origin1 = "compress";
    final String origin2 = "origin";

    final byte[] originBytes1 = origin1.getBytes(StandardCharsets.UTF_8);
    info(originBytes1);
    final byte[] originBytes2 = origin2.getBytes(StandardCharsets.UTF_8);
    info(originBytes2);

    final byte[] compressedBytes1 = Zstd.compress(originBytes1);
    info(compressedBytes1);

    final byte[] fixedBytes = ByteUtils.merge(compressedBytes1, originBytes2);
    info(fixedBytes);

    final long size = Zstd.decompressedSize(fixedBytes);
    System.out.println(size);
    System.out.println(fixedBytes.length);

    final ZstdException e = Assert.assertThrows(ZstdException.class, () -> {
        final byte[] decompress = Zstd.decompress(fixedBytes, compressSize(fixedBytes));
        System.out.println(new String(decompress, StandardCharsets.UTF_8));
    });
    Assert.assertEquals("Src size is incorrect", e.getMessage());
    System.err.println("Compressed and uncompressed arrays should not be mixed together.");
}
 

@Test
public void testSimple() {
    final String origin = "yuzhouwan.com, www.yuzhouwan.com, unknown.yuzhouwan.com";

    final byte[] originBytes = origin.getBytes(StandardCharsets.UTF_8);
    info(originBytes);

    final byte[] compressedBytes = Zstd.compress(originBytes);
    info(compressedBytes);

    final int size = compressSize(compressedBytes);
    final byte[] decompress = Zstd.decompress(compressedBytes, size);
    Assert.assertEquals(origin, new String(decompress, StandardCharsets.UTF_8));
}
 

/**
 * @param compactPage Compacted page.
 * @param compactSize Compacted page size.
 * @param compressLevel Compression level.
 * @return Compressed page.
 */
private ByteBuffer compressPageZstd(ByteBuffer compactPage, int compactSize, int compressLevel) {
    ByteBuffer compressedPage = compressBuf.get();

    copyPageHeader(compactPage, compressedPage, compactSize);
    Zstd.compress(compressedPage, compactPage, compressLevel);

    compactPage.flip();
    compressedPage.flip();

    return compressedPage;
}
 

/**
 * Reads a block of zstd-compressed data. This method
 * expects the following ints to be the compressed size,
 * and uncompressed size respectively.
 *
 * @return the uncompressed data
 * @throws IOException if the bytes cannot be read
 * @throws IllegalArgumentException if the uncompressed length doesn't match
 */
public byte[] readCompressed() throws IOException {
    int compressedLength = readInt();
    int uncompressedLength = readInt();

    byte[] compressed = readByteArray(compressedLength);
    byte[] data = Zstd.decompress(compressed, uncompressedLength);

    if (data.length != uncompressedLength) {
        throw new IllegalArgumentException("Uncompressed length doesn't match");
    }

    return data;
}
 

@Test
public void testZstdCompressingDecompressingStream() throws IOException {

    String sample = randomString(DEFAULT_BLOCK_SIZE);
    ByteBuf source = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    source.writeBytes(sample.getBytes());

    ByteBuf target = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    int idx = target.writerIndex();
    ByteBuffer buffer = target.internalNioBuffer(idx + HEADER_LENGTH, target.writableBytes() - HEADER_LENGTH);

    int compressLength = Zstd.compress(
            buffer,
            source.internalNioBuffer(source.readerIndex(), source.readableBytes()),
            0);
    target.writerIndex(idx + HEADER_LENGTH + compressLength);
    logger.info("{}", target.readerIndex());
    logger.info("{}", target.readableBytes());
    logger.info("{}", target.writerIndex());

    ByteBuf output = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    ByteBuffer outputBuffer = output.internalNioBuffer(output.writerIndex(), output.writableBytes());

    ByteBuffer inputBuffer = target.internalNioBuffer(target.readerIndex() + HEADER_LENGTH, target.readableBytes() - HEADER_LENGTH);

    logger.info("{}", inputBuffer.position());
    logger.info("{}", inputBuffer.limit());
    Zstd.decompress(outputBuffer, inputBuffer);

    outputBuffer.flip();
    output = Unpooled.wrappedBuffer(outputBuffer);

    String val = output.toString(Charset.defaultCharset());

    Assert.assertEquals(sample, val);
}
 

private void flushBufferedData(ByteBuf out) {
    final int flushableBytes = buffer.readableBytes();
    if (flushableBytes == 0) {
        return;
    }
    checksum.reset();
    checksum.update(buffer.internalNioBuffer(buffer.readerIndex(), flushableBytes));
    final int check = (int) checksum.getValue();

    final int bufSize = (int) Zstd.compressBound(flushableBytes) + HEADER_LENGTH;
    out.ensureWritable(bufSize);
    final int idx = out.writerIndex();
    int compressedLength;
    try {
        ByteBuffer outNioBuffer = out.internalNioBuffer(idx + HEADER_LENGTH, out.writableBytes() - HEADER_LENGTH);
        compressedLength = Zstd.compress(
                outNioBuffer,
                buffer.internalNioBuffer(buffer.readerIndex(), flushableBytes),
                DEFAULT_COMPRESS_LEVEL);
    } catch (Exception e) {
        throw new CompressionException(e);
    }
    final int blockType;
    if (compressedLength >= flushableBytes) {
        blockType = BLOCK_TYPE_NON_COMPRESSED;
        compressedLength = flushableBytes;
        out.setBytes(idx + HEADER_LENGTH, buffer, 0, flushableBytes);
    } else {
        blockType = BLOCK_TYPE_COMPRESSED;
    }

    out.setInt(idx, MAGIC_NUMBER);
    out.setByte(idx + TOKEN_OFFSET, (byte) (blockType | compressionLevel));
    out.setIntLE(idx + COMPRESSED_LENGTH_OFFSET, compressedLength);
    out.setIntLE(idx + DECOMPRESSED_LENGTH_OFFSET, flushableBytes);
    out.setIntLE(idx + CHECKSUM_OFFSET, check);
    out.writerIndex(idx + HEADER_LENGTH + compressedLength);
    buffer.clear();
}
 

private ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect,
                               boolean allowEmptyReturn) {
    int targetBufSize = 0;
    int remaining = msg.readableBytes() + buffer.readableBytes();

    // quick overflow check
    if (remaining < 0) {
        throw new EncoderException("too much data to allocate a buffer for compression");
    }

    while (remaining > 0) {
        int curSize = Math.min(blockSize, remaining);
        remaining -= curSize;
        if(curSize <= MIN_BLOCK_SIZE) {
            targetBufSize += curSize + HEADER_LENGTH;
        } else {
            targetBufSize += Zstd.compressBound(curSize) + HEADER_LENGTH;
        }
    }

    if (targetBufSize > maxEncodeSize || 0 > targetBufSize) {
        throw new EncoderException(String.format("requested encode buffer size (%d bytes) exceeds the maximum " +
                "allowable size (%d bytes)", targetBufSize, maxEncodeSize));
    }

    if (preferDirect) {
        return ctx.alloc().ioBuffer(targetBufSize, targetBufSize);
    } else {
        return ctx.alloc().heapBuffer(targetBufSize, targetBufSize);
    }
}
 

public int read(byte[] dst, int offset, int len) throws IOException {

        if (isClosed) {
            throw new IOException("Stream closed");
        }

        // guard agains buffer overflows
        if (offset < 0 || len > dst.length - offset) {
            throw new IndexOutOfBoundsException("Requested lenght " + len
                    + " from offset " + offset + " in buffer of size " + dst.length);
        }
        int dstSize = offset + len;
        dstPos = offset;

        while (dstPos < dstSize) {
            if (srcSize - srcPos == 0) {
                srcSize = in.read(src, 0, srcBuffSize);
                srcPos = 0;
                if (srcSize < 0) {
                    srcSize = 0;
                    if (frameFinished) {
                        return -1;
                    } else if (isContinuous) {
                        return (int)(dstPos - offset);
                    } else {
                        throw new IOException("Read error or truncated source");
                    }
                }
                frameFinished = false;
            }

            int size = decompressStream(stream, dst, dstSize, src, (int) srcSize);

            if (Zstd.isError(size)) {
                throw new IOException("Decompression error: " + Zstd.getErrorName(size));
            }

            // we have completed a frame
            if (size == 0) {
                frameFinished = true;
                // re-init the codec so it can start decoding next frame
                size = initDStream(stream);
                if (Zstd.isError(size)) {
                    throw new IOException("Decompression error: " + Zstd.getErrorName(size));
                }
                return (int)(dstPos - offset);
            }
        }
        return len;
    }
 

private static int compressSize(byte[] compressedBytes) {
    return (int) Math.max(compressedBytes.length, Zstd.decompressedSize(compressedBytes));
}
 

/**
 * Method for obtaining decompressed data for layout 2.
 */
private byte[] getDecompressedBlockData(ReadOnlyGeneticVariantBgen variant) throws IOException {
	// First extend the genetic variant with ids, alleles, etc.
	// This makes sure the file pointer is in the correct location for the next step.
	variant.extendWithAdditionalVariantData();

	// Extract the variant genotype data block info starting from the current position of the
	// file pointer, which should be right after all alleles for a specific variant.
	VariantGenotypeBlockInfo variantGenotypeBlockInfo = extractVariantGenotypeDataBlockInfo();

	long decompressedVariantBlockLength = variantGenotypeBlockInfo.getDecompressedBlockLength();
	long variantBlockLength = variantGenotypeBlockInfo.getBlockLength();
	long variantProbabilitiesStartPosition = variantGenotypeBlockInfo.getVariantProbabilitiesStartPosition();

	if (decompressedVariantBlockLength > Integer.MAX_VALUE - 5) {
		throw new GenotypeDataException(String.format(
				"Length of decompressed genotype data exceeds maximum supported value of (2^31)-6 (%d)",
				decompressedVariantBlockLength));
	}

	if (variantBlockLength > Integer.MAX_VALUE - 5) {
		throw new GenotypeDataException(String.format(
				"Length of compressed genotype data exceeds maximum supported value of (2^31)-6 (%d)", variantBlockLength));
	}

	// Initialize byte arrays.
	byte[] compressedBlockData = new byte[(int) variantBlockLength];
	byte[] decompressedBlockData = new byte[(int) decompressedVariantBlockLength];

	// Read the compressed / uncompressed data starting from the correct location.
	this.bgenFile.seek(variantProbabilitiesStartPosition);
	this.bgenFile.read(compressedBlockData, 0, (int) variantBlockLength);

	switch (snpBlockRepresentation) {

		case compression_1:
			decompressVariantBlockGzip(
					compressedBlockData,
					decompressedBlockData);
			break;

		case compression_2:
			Zstd.decompress(
					decompressedBlockData,
					compressedBlockData);
			break;

		default:
			decompressedBlockData = compressedBlockData;
			break;
	}
	return decompressedBlockData;
}
 

@Test
public void testEncodeWithDefaultBlockSize() {
    String sample = randomString(DEFAULT_BLOCK_SIZE);

    EmbeddedChannel channel = new EmbeddedChannel(encoder);

    ByteBuf byteBuf = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    byte[] sampleBytes = sample.getBytes();
    byteBuf.ensureWritable(sampleBytes.length);
    byteBuf.writeBytes(sampleBytes);

    Assert.assertTrue(channel.writeOutbound(byteBuf));

    ByteBuf target = channel.readOutbound();
    Assert.assertNotNull(target);

    ByteBuffer buffer = ByteBuffer.allocateDirect(sampleBytes.length);
    buffer.clear();

    target.readerIndex(target.readerIndex() + HEADER_LENGTH);
    ByteBuffer source = target.internalNioBuffer(target.readerIndex(), target.readableBytes());

    Zstd.decompress(buffer, source);

    buffer.flip();
    ByteBuf buf = Unpooled.wrappedBuffer(buffer);

    String val = buf.toString(Charset.defaultCharset());

    Assert.assertEquals(sample, val);
}