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

131 lines
3.9 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 <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;
2012-01-05 18:35:22 +00:00
CompressionMethod::Enum method;
2010-06-04 18:25:25 +00:00
std::vector<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;
char * compressed_buffer_ptr = NULL;
/** Для того, чтобы различить между QuickLZ и LZ4 и сохранить обратную совместимость (со случаем, когда использовался только QuickLZ),
* используем старший бит первого байта в сжатых данных (который сейчас не используется в QuickLZ).
* PS. Если потребуется использовать другие библиотеки, то можно использовать ещё один бит первого байта, или старший бит размера.
*/
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_ptr = &compressed_buffer[0];
break;
}
case CompressionMethod::LZ4:
{
/** В случае LZ4, в начале запишем заголовок такого же размера и структуры, как в QuickLZ
* 1 байт, чтобы отличить LZ4 от QuickLZ.
* 4 байта - размер сжатых данных
* 4 байта - размер несжатых данных.
*/
2012-05-02 15:45:53 +00:00
compressed_buffer.resize(QUICKLZ_HEADER_SIZE + LZ4_compressBound(uncompressed_size));
2012-01-05 18:35:22 +00:00
compressed_buffer[0] = 0x82; /// Второй бит - для совместимости с QuickLZ - обозначает, что размеры записываются 4 байтами.
compressed_size = QUICKLZ_HEADER_SIZE + LZ4_compress(
working_buffer.begin(),
&compressed_buffer[QUICKLZ_HEADER_SIZE],
uncompressed_size);
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;
}
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_,
2012-04-14 00:18:07 +00:00
CompressionMethod::Enum 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()
{
if (!std::uncaught_exception())
next();
delete qlz_state;
2010-06-04 18:25:25 +00:00
}
};
}