ru.yandex.clickhouse.response.ClickHouseLZ4Stream Maven / Gradle / Ivy
package ru.yandex.clickhouse.response;
import net.jpountz.lz4.LZ4Factory;
import net.jpountz.lz4.LZ4FastDecompressor;
import ru.yandex.clickhouse.util.ClickHouseBlockChecksum;
import ru.yandex.clickhouse.util.Utils;
import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
/**
* Reader from clickhouse in lz4
*/
public class ClickHouseLZ4Stream extends InputStream {
private static final LZ4Factory factory = LZ4Factory.fastestInstance();
public static final int MAGIC = 0x82;
private final InputStream stream;
private final DataInputStream dataWrapper;
private byte[] currentBlock;
private int pointer;
public ClickHouseLZ4Stream(InputStream stream) {
this.stream = stream;
dataWrapper = new DataInputStream(stream);
}
@Override
public int read() throws IOException {
if (!checkNext()) return -1;
byte b = currentBlock[pointer];
pointer += 1;
return b & 0xFF;
}
@Override
public int read(byte[] b, int off, int len) throws IOException {
if (b == null) {
throw new NullPointerException();
} else if (off < 0 || len < 0 || len > b.length - off) {
throw new IndexOutOfBoundsException();
} else if (len == 0) {
return 0;
}
if (!checkNext()) return -1;
int copied = 0;
int targetPointer = off;
while(copied != len) {
int toCopy = Math.min(currentBlock.length - pointer, len - copied);
System.arraycopy(currentBlock, pointer, b, targetPointer, toCopy);
targetPointer += toCopy;
pointer += toCopy;
copied += toCopy;
if (!checkNext()) { //finished
return copied;
}
}
return copied;
}
@Override
public void close() throws IOException {
stream.close();
}
private boolean checkNext() throws IOException {
if (currentBlock == null || pointer == currentBlock.length) {
currentBlock = readNextBlock();
pointer = 0;
}
return currentBlock != null;
}
// every block is:
private byte[] readNextBlock() throws IOException {
int read = stream.read();
if (read < 0) return null;
byte[] checksum = new byte[16];
checksum[0] = (byte)read;
// checksum - 16 bytes.
Utils.readFully(dataWrapper, checksum, 1, 15);
ClickHouseBlockChecksum expected = ClickHouseBlockChecksum.fromBytes(checksum);
// header:
// 1 byte - 0x82 (shows this is LZ4)
int magic = dataWrapper.readUnsignedByte();
if (magic != MAGIC) throw new IOException("Magic is not correct: " + magic);
// 4 bytes - size of the compressed data including 9 bytes of the header
int compressedSizeWithHeader = Utils.readInt(dataWrapper);
// 4 bytes - size of uncompressed data
int uncompressedSize = Utils.readInt(dataWrapper);
int compressedSize = compressedSizeWithHeader - 9; //header
byte[] block = new byte[compressedSize];
// compressed data: compressed_size - 9 байт.
Utils.readFully(dataWrapper, block);
ClickHouseBlockChecksum real = ClickHouseBlockChecksum.calculateForBlock((byte)magic, compressedSizeWithHeader, uncompressedSize, block, compressedSize);
if (!real.equals(expected)) {
throw new IllegalArgumentException("Checksum doesn't match: corrupted data.");
}
byte[] decompressed = new byte[uncompressedSize];
LZ4FastDecompressor decompressor = factory.fastDecompressor();
decompressor.decompress(block, 0, decompressed, 0, uncompressedSize);
return decompressed;
}
}