mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-10-01 22:20:50 +00:00
134 lines
3.8 KiB
C++
134 lines
3.8 KiB
C++
#include <memory>
|
|
#include <city.h>
|
|
#include <lz4.h>
|
|
#include <lz4hc.h>
|
|
#include <zstd.h>
|
|
|
|
#include <Common/unaligned.h>
|
|
#include <Core/Types.h>
|
|
|
|
#include <IO/CompressedWriteBuffer.h>
|
|
|
|
|
|
namespace DB
|
|
{
|
|
|
|
namespace ErrorCodes
|
|
{
|
|
extern const int CANNOT_COMPRESS;
|
|
extern const int UNKNOWN_COMPRESSION_METHOD;
|
|
}
|
|
|
|
|
|
void CompressedWriteBuffer::nextImpl()
|
|
{
|
|
if (!offset())
|
|
return;
|
|
|
|
size_t uncompressed_size = offset();
|
|
size_t compressed_size = 0;
|
|
char * compressed_buffer_ptr = nullptr;
|
|
|
|
/** The format of compressed block - see CompressedStream.h
|
|
*/
|
|
|
|
switch (method)
|
|
{
|
|
case CompressionMethod::LZ4:
|
|
case CompressionMethod::LZ4HC:
|
|
{
|
|
static constexpr size_t header_size = 1 + sizeof(UInt32) + sizeof(UInt32);
|
|
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wold-style-cast"
|
|
compressed_buffer.resize(header_size + LZ4_COMPRESSBOUND(uncompressed_size));
|
|
#pragma GCC diagnostic pop
|
|
|
|
compressed_buffer[0] = static_cast<UInt8>(CompressionMethodByte::LZ4);
|
|
|
|
if (method == CompressionMethod::LZ4)
|
|
compressed_size = header_size + LZ4_compress_default(
|
|
working_buffer.begin(),
|
|
&compressed_buffer[header_size],
|
|
uncompressed_size,
|
|
LZ4_COMPRESSBOUND(uncompressed_size));
|
|
else
|
|
compressed_size = header_size + LZ4_compress_HC(
|
|
working_buffer.begin(),
|
|
&compressed_buffer[header_size],
|
|
uncompressed_size,
|
|
LZ4_COMPRESSBOUND(uncompressed_size),
|
|
0);
|
|
|
|
UInt32 compressed_size_32 = compressed_size;
|
|
UInt32 uncompressed_size_32 = uncompressed_size;
|
|
|
|
unalignedStore(&compressed_buffer[1], compressed_size_32);
|
|
unalignedStore(&compressed_buffer[5], 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,
|
|
1);
|
|
|
|
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;
|
|
|
|
unalignedStore(&compressed_buffer[1], compressed_size_32);
|
|
unalignedStore(&compressed_buffer[5], uncompressed_size_32);
|
|
|
|
compressed_buffer_ptr = &compressed_buffer[0];
|
|
break;
|
|
}
|
|
default:
|
|
throw Exception("Unknown compression method", ErrorCodes::UNKNOWN_COMPRESSION_METHOD);
|
|
}
|
|
|
|
CityHash_v1_0_2::uint128 checksum = CityHash_v1_0_2::CityHash128(compressed_buffer_ptr, compressed_size);
|
|
out.write(reinterpret_cast<const char *>(&checksum), sizeof(checksum));
|
|
|
|
out.write(compressed_buffer_ptr, compressed_size);
|
|
}
|
|
|
|
|
|
CompressedWriteBuffer::CompressedWriteBuffer(
|
|
WriteBuffer & out_,
|
|
CompressionMethod method_,
|
|
size_t buf_size)
|
|
: BufferWithOwnMemory<WriteBuffer>(buf_size), out(out_), method(method_)
|
|
{
|
|
}
|
|
|
|
|
|
CompressedWriteBuffer::~CompressedWriteBuffer()
|
|
{
|
|
try
|
|
{
|
|
next();
|
|
}
|
|
catch (...)
|
|
{
|
|
tryLogCurrentException(__PRETTY_FUNCTION__);
|
|
}
|
|
}
|
|
|
|
}
|