ClickHouse/dbms/include/DB/IO/CompressedWriteBuffer.h

171 lines
4.3 KiB
C
Raw Normal View History

2012-01-05 10:22:02 +00:00
#pragma once
2010-06-04 18:25:25 +00:00
2012-01-05 18:35:22 +00:00
#include <math.h>
#include <vector>
#include <city.h>
2011-06-17 21:19:39 +00:00
#include <quicklz/quicklz_level1.h>
2012-01-05 18:35:22 +00:00
#include <lz4/lz4.h>
#include <lz4/lz4hc.h>
#include <zstd/zstd.h>
2012-01-05 18:35:22 +00:00
2013-12-09 00:23:17 +00:00
#include <DB/Common/PODArray.h>
2012-01-05 18:35:22 +00:00
#include <DB/Core/Types.h>
2010-06-04 18:25:25 +00:00
#include <DB/IO/WriteBuffer.h>
#include <DB/IO/BufferWithOwnMemory.h>
2011-06-17 21:19:39 +00:00
#include <DB/IO/CompressedStream.h>
2010-06-04 18:25:25 +00:00
namespace DB
{
class CompressedWriteBuffer : public BufferWithOwnMemory<WriteBuffer>
2010-06-04 18:25:25 +00:00
{
private:
WriteBuffer & out;
CompressionMethod method;
2010-06-04 18:25:25 +00:00
2013-12-09 00:23:17 +00:00
PODArray<char> compressed_buffer;
qlz_state_compress * qlz_state;
2011-06-17 21:19:39 +00:00
void nextImpl()
2010-06-04 18:25:25 +00:00
{
2011-06-27 19:34:03 +00:00
if (!offset())
return;
size_t uncompressed_size = offset();
2012-01-05 18:35:22 +00:00
size_t compressed_size = 0;
2014-04-08 07:31:51 +00:00
char * compressed_buffer_ptr = nullptr;
2012-01-05 18:35:22 +00:00
/** Формат сжатого блока - см. CompressedStream.h
2012-01-05 18:35:22 +00:00
*/
2011-06-26 21:30:59 +00:00
2012-01-05 18:35:22 +00:00
switch (method)
{
case CompressionMethod::QuickLZ:
{
compressed_buffer.resize(uncompressed_size + QUICKLZ_ADDITIONAL_SPACE);
compressed_size = qlz_compress(
working_buffer.begin(),
&compressed_buffer[0],
uncompressed_size,
qlz_state);
2012-01-05 18:35:22 +00:00
compressed_buffer[0] &= 3;
2012-01-05 18:35:22 +00:00
compressed_buffer_ptr = &compressed_buffer[0];
break;
}
case CompressionMethod::LZ4:
case CompressionMethod::LZ4HC:
2012-01-05 18:35:22 +00:00
{
static constexpr size_t header_size = 1 + sizeof(UInt32) + sizeof(UInt32);
compressed_buffer.resize(header_size + LZ4_COMPRESSBOUND(uncompressed_size));
2012-01-05 18:35:22 +00:00
compressed_buffer[0] = static_cast<UInt8>(CompressionMethodByte::LZ4);
2012-01-05 18:35:22 +00:00
if (method == CompressionMethod::LZ4)
compressed_size = header_size + LZ4_compress(
working_buffer.begin(),
&compressed_buffer[header_size],
uncompressed_size);
else
compressed_size = header_size + LZ4_compressHC(
working_buffer.begin(),
&compressed_buffer[header_size],
uncompressed_size);
2012-01-05 18:35:22 +00:00
UInt32 compressed_size_32 = compressed_size;
UInt32 uncompressed_size_32 = uncompressed_size;
memcpy(&compressed_buffer[1], reinterpret_cast<const char *>(&compressed_size_32), sizeof(compressed_size_32));
memcpy(&compressed_buffer[5], reinterpret_cast<const char *>(&uncompressed_size_32), sizeof(uncompressed_size_32));
compressed_buffer_ptr = &compressed_buffer[0];
break;
}
case CompressionMethod::ZSTD:
{
static constexpr size_t header_size = 1 + sizeof(UInt32) + sizeof(UInt32);
compressed_buffer.resize(header_size + ZSTD_compressBound(uncompressed_size));
compressed_buffer[0] = static_cast<UInt8>(CompressionMethodByte::ZSTD);
size_t res = ZSTD_compress(
&compressed_buffer[header_size],
compressed_buffer.size(),
working_buffer.begin(),
uncompressed_size);
if (ZSTD_isError(res))
throw Exception("Cannot compress block with ZSTD: " + std::string(ZSTD_getErrorName(res)), ErrorCodes::CANNOT_COMPRESS);
compressed_size = header_size + res;
UInt32 compressed_size_32 = compressed_size;
UInt32 uncompressed_size_32 = uncompressed_size;
memcpy(&compressed_buffer[1], reinterpret_cast<const char *>(&compressed_size_32), sizeof(compressed_size_32));
memcpy(&compressed_buffer[5], reinterpret_cast<const char *>(&uncompressed_size_32), sizeof(uncompressed_size_32));
compressed_buffer_ptr = &compressed_buffer[0];
break;
}
2012-01-05 18:35:22 +00:00
default:
throw Exception("Unknown compression method", ErrorCodes::UNKNOWN_COMPRESSION_METHOD);
}
2010-06-04 18:25:25 +00:00
2012-01-05 18:35:22 +00:00
uint128 checksum = CityHash128(compressed_buffer_ptr, compressed_size);
out.write(reinterpret_cast<const char *>(&checksum), sizeof(checksum));
2012-01-05 18:35:22 +00:00
out.write(compressed_buffer_ptr, compressed_size);
2011-05-05 19:10:17 +00:00
}
2011-06-26 21:30:59 +00:00
public:
CompressedWriteBuffer(
WriteBuffer & out_,
CompressionMethod method_ = CompressionMethod::LZ4,
size_t buf_size = DBMS_DEFAULT_BUFFER_SIZE)
: BufferWithOwnMemory<WriteBuffer>(buf_size), out(out_), method(method_), qlz_state(new qlz_state_compress) {}
2011-06-26 21:30:59 +00:00
/// Объём сжатых данных
2011-05-05 19:10:17 +00:00
size_t getCompressedBytes()
{
nextIfAtEnd();
return out.count();
2011-05-05 19:10:17 +00:00
}
/// Сколько несжатых байт было записано в буфер
2011-05-05 19:10:17 +00:00
size_t getUncompressedBytes()
{
return count();
}
/// Сколько байт находится в буфере (ещё не сжато)
size_t getRemainingBytes()
2011-05-05 19:10:17 +00:00
{
nextIfAtEnd();
return offset();
2010-06-04 18:25:25 +00:00
}
~CompressedWriteBuffer()
{
2013-11-18 17:17:45 +00:00
try
{
next();
2013-11-18 17:17:45 +00:00
}
catch (...)
{
2013-11-18 19:18:03 +00:00
tryLogCurrentException(__PRETTY_FUNCTION__);
2013-11-18 17:17:45 +00:00
}
delete qlz_state;
2010-06-04 18:25:25 +00:00
}
};
}