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

163 lines
4.8 KiB
C
Raw Normal View History

2011-11-18 20:17:27 +00:00
#pragma once
2010-06-04 18:25:25 +00:00
2010-06-07 12:48:32 +00:00
#include <vector>
2010-06-04 18:25:25 +00:00
#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>
2010-06-04 18:25:25 +00:00
#include <DB/Core/Exception.h>
#include <DB/Core/ErrorCodes.h>
#include <DB/IO/ReadBuffer.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 CompressedReadBuffer : public BufferWithOwnMemory<ReadBuffer>
2010-06-04 18:25:25 +00:00
{
private:
ReadBuffer & in;
std::vector<char> compressed_buffer;
qlz_state_decompress * qlz_state;
2011-06-17 21:19:39 +00:00
/** Указатель на кусок памяти, куда будут разжиматься блоки.
* Это может быть либо свой кусок памяти из BufferWithOwnMemory (по-умолчанию),
* либо пользователь может попросить разжимать данные в свой кусок памяти (метод setMemory).
*/
Memory * maybe_own_memory;
/// Прочитать сжатые данные в compressed_buffer. Достать из их заголовка размер разжатых данных. Проверить чексумму.
bool readCompressedData(size_t & size_decompressed)
2010-06-04 18:25:25 +00:00
{
if (in.eof())
return false;
uint128 checksum;
in.readStrict(reinterpret_cast<char *>(&checksum), sizeof(checksum));
2011-06-26 21:30:59 +00:00
2011-06-17 21:19:39 +00:00
in.readStrict(&compressed_buffer[0], QUICKLZ_HEADER_SIZE);
2010-06-04 18:25:25 +00:00
2011-06-17 21:19:39 +00:00
size_t size_compressed = qlz_size_compressed(&compressed_buffer[0]);
if (size_compressed > DBMS_MAX_COMPRESSED_SIZE)
throw Exception("Too large size_compressed. Most likely corrupted data.", ErrorCodes::TOO_LARGE_SIZE_COMPRESSED);
2011-06-26 21:30:59 +00:00
size_decompressed = qlz_size_decompressed(&compressed_buffer[0]);
2010-06-04 18:25:25 +00:00
compressed_buffer.resize(size_compressed);
2011-06-17 21:19:39 +00:00
in.readStrict(&compressed_buffer[QUICKLZ_HEADER_SIZE], size_compressed - QUICKLZ_HEADER_SIZE);
2010-06-07 12:48:32 +00:00
if (checksum != CityHash128(&compressed_buffer[0], size_compressed))
throw Exception("Checksum doesnt match: corrupted data.", ErrorCodes::CHECKSUM_DOESNT_MATCH);
2011-06-26 21:30:59 +00:00
return true;
}
void decompress(char * to, size_t size_decompressed)
{
2012-01-05 18:35:22 +00:00
/// Старший бит первого байта определяет использованный метод сжатия.
if ((compressed_buffer[0] & 0x80) == 0)
2012-06-21 18:14:38 +00:00
{
if (!qlz_state)
qlz_state = new qlz_state_decompress;
qlz_decompress(&compressed_buffer[0], to, qlz_state);
2012-06-21 18:14:38 +00:00
}
2012-01-05 18:35:22 +00:00
else
LZ4_uncompress(&compressed_buffer[QUICKLZ_HEADER_SIZE], to, size_decompressed);
}
bool nextImpl()
{
size_t size_decompressed = 0;
if (!readCompressedData(size_decompressed))
return false;
maybe_own_memory->resize(size_decompressed);
internal_buffer = Buffer(&(*maybe_own_memory)[0], &(*maybe_own_memory)[size_decompressed]);
working_buffer = Buffer(&(*maybe_own_memory)[0], &(*maybe_own_memory)[size_decompressed]);
decompress(working_buffer.begin(), size_decompressed);
2010-06-04 18:25:25 +00:00
return true;
}
2011-06-26 21:30:59 +00:00
public:
CompressedReadBuffer(ReadBuffer & in_)
2012-06-21 18:43:29 +00:00
: BufferWithOwnMemory<ReadBuffer>(0),
in(in_),
compressed_buffer(QUICKLZ_HEADER_SIZE),
qlz_state(NULL),
maybe_own_memory(&memory)
2011-06-26 21:30:59 +00:00
{
}
~CompressedReadBuffer()
{
2012-06-21 18:14:38 +00:00
if (qlz_state)
delete qlz_state;
}
/// Использовать предоставленный пользователем кусок памяти для разжатия. (Для реализации кэша разжатых блоков.)
void setMemory(Memory & memory_)
{
maybe_own_memory = &memory_;
}
size_t readBig(char * to, size_t n)
{
size_t bytes_read = 0;
/// Если в буфере есть непрочитанные байты, то скопируем сколько надо в to.
if (pos < working_buffer.end())
bytes_read += read(to, std::min(static_cast<size_t>(working_buffer.end() - pos), n));
if (bytes_read < n)
bytes += offset();
/// Если надо ещё прочитать - будем, по возможности, разжимать сразу в to.
while (bytes_read < n)
{
size_t size_decompressed = 0;
if (!readCompressedData(size_decompressed))
{
working_buffer.resize(0);
pos = working_buffer.begin();
return bytes_read;
}
/// Если разжатый блок помещается целиком туда, куда его надо скопировать.
if (size_decompressed <= n - bytes_read)
{
decompress(to + bytes_read, size_decompressed);
bytes_read += size_decompressed;
bytes += size_decompressed;
}
else
{
maybe_own_memory->resize(size_decompressed);
internal_buffer = Buffer(&(*maybe_own_memory)[0], &(*maybe_own_memory)[size_decompressed]);
working_buffer = Buffer(&(*maybe_own_memory)[0], &(*maybe_own_memory)[size_decompressed]);
pos = working_buffer.begin();
decompress(working_buffer.begin(), size_decompressed);
bytes_read += read(to + bytes_read, n - bytes_read);
break;
}
}
return bytes_read;
}
2010-06-04 18:25:25 +00:00
};
}