ClickHouse/src/IO/ZstdInflatingReadBuffer.cpp

71 lines
2.0 KiB
C++
Raw Normal View History

2020-11-17 16:02:10 +00:00
#include <IO/ZstdInflatingReadBuffer.h>
namespace DB
{
namespace ErrorCodes
{
extern const int ZSTD_DECODER_FAILED;
}
ZstdInflatingReadBuffer::ZstdInflatingReadBuffer(std::unique_ptr<ReadBuffer> in_, size_t buf_size, char * existing_memory, size_t alignment)
: BufferWithOwnMemory<ReadBuffer>(buf_size, existing_memory, alignment), in(std::move(in_))
{
dctx = ZSTD_createDCtx();
input = {nullptr, 0, 0};
output = {nullptr, 0, 0};
if (dctx == nullptr)
{
throw Exception(ErrorCodes::ZSTD_DECODER_FAILED, "zstd_stream_decoder init failed: zstd version: {}", ZSTD_VERSION_STRING);
}
}
ZstdInflatingReadBuffer::~ZstdInflatingReadBuffer()
{
ZSTD_freeDCtx(dctx);
}
bool ZstdInflatingReadBuffer::nextImpl()
{
if (eof)
return false;
if (input.pos >= input.size)
{
in->nextIfAtEnd();
input.src = reinterpret_cast<unsigned char *>(in->position());
input.pos = 0;
input.size = in->buffer().end() - in->position();
}
output.dst = reinterpret_cast<unsigned char *>(internal_buffer.begin());
output.size = internal_buffer.size();
output.pos = 0;
size_t ret = ZSTD_decompressStream(dctx, &output, &input);
if (ZSTD_isError(ret))
throw Exception(
2020-11-23 04:11:52 +00:00
ErrorCodes::ZSTD_DECODER_FAILED, "Zstd stream decoding failed: error code: {}; zstd version: {}", ret, ZSTD_VERSION_STRING);
2020-11-17 16:02:10 +00:00
in->position() = in->buffer().begin() + input.pos;
working_buffer.resize(output.pos);
if (in->eof())
{
eof = true;
2021-02-04 23:14:17 +00:00
return !working_buffer.empty();
2020-11-17 16:02:10 +00:00
}
2021-07-14 17:15:05 +00:00
else if (output.pos == 0)
{
2021-07-15 14:58:58 +00:00
/// It is possible, that input buffer is not at eof yet, but nothing was decompressed in current iteration.
/// But there are cases, when such behaviour is not allowed - i.e. if input buffer is not eof, then
/// it has to be guaranteed that working_buffer is not empty. So if it is empty, continue.
2021-07-14 17:15:05 +00:00
return nextImpl();
}
2020-11-17 16:02:10 +00:00
return true;
}
2020-11-17 18:36:18 +00:00
}