-
Notifications
You must be signed in to change notification settings - Fork 8.9k
Expand file tree
/
Copy pathCompressionCodecMultiple.cpp
More file actions
179 lines (141 loc) · 6.42 KB
/
Copy pathCompressionCodecMultiple.cpp
File metadata and controls
179 lines (141 loc) · 6.42 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
#include <Compression/CompressionCodecMultiple.h>
#include <Compression/CompressionInfo.h>
#include <Compression/registerCompressionCodecs.h>
#include <Common/PODArray.h>
#include <Compression/CompressionFactory.h>
#include <IO/WriteHelpers.h>
#include <IO/WriteBufferFromString.h>
namespace DB
{
CompressionCodecMultiple::CompressionCodecMultiple(Codecs codecs_)
: codecs(codecs_)
{
ASTs arguments;
for (const auto & codec : codecs)
arguments.push_back(codec->getCodecDesc());
/// Special case, codec doesn't have name and contain list of codecs.
setCodecDescription("", arguments);
}
uint8_t CompressionCodecMultiple::getMethodByte() const
{
return static_cast<uint8_t>(CompressionMethodByte::Multiple);
}
void CompressionCodecMultiple::updateHash(SipHash & hash) const
{
for (const auto & codec : codecs)
codec->updateHash(hash);
}
UInt32 CompressionCodecMultiple::getMaxCompressedDataSize(UInt32 uncompressed_size) const
{
UInt32 compressed_size = uncompressed_size;
for (const auto & codec : codecs)
compressed_size = codec->getCompressedReserveSize(compressed_size);
/// TotalCodecs ByteForEachCodec data
return static_cast<UInt32>(sizeof(UInt8) + codecs.size() + compressed_size);
}
UInt32 CompressionCodecMultiple::doCompressData(const char * source, UInt32 source_size, char * dest) const
{
PODArray<char> compressed_buf;
PODArray<char> uncompressed_buf(source, source + source_size);
dest[0] = static_cast<UInt8>(codecs.size());
size_t codecs_byte_pos = 1;
for (size_t idx = 0; idx < codecs.size(); ++idx, ++codecs_byte_pos)
{
const auto codec = codecs[idx];
dest[codecs_byte_pos] = codec->getMethodByte();
compressed_buf.resize(codec->getCompressedReserveSize(source_size));
UInt32 size_compressed = codec->compress(uncompressed_buf.data(), source_size, compressed_buf.data());
uncompressed_buf.swap(compressed_buf);
source_size = size_compressed;
}
memcpy(&dest[1 + codecs.size()], uncompressed_buf.data(), source_size);
return static_cast<UInt32>(1 + codecs.size() + source_size);
}
UInt32 CompressionCodecMultiple::doDecompressData(const char * source, UInt32 source_size, char * dest, UInt32 decompressed_size) const
{
if (source_size < 1 || !source[0])
throw Exception(decompression_error_code, "Wrong compression methods list");
UInt8 compression_methods_size = source[0];
/// +1 for the compression_methods_size byte itself
if (static_cast<UInt32>(compression_methods_size) + 1 > source_size)
throw Exception(decompression_error_code, "Wrong compression methods list: header claims {} codecs"
" but compressed data is only {} bytes",
static_cast<UInt32>(compression_methods_size), source_size);
PODArray<char> compressed_buf(&source[compression_methods_size + 1], &source[source_size]);
PODArray<char> uncompressed_buf;
/// Insert all data into compressed buf
source_size -= (compression_methods_size + 1);
for (int idx = compression_methods_size - 1; idx >= 0; --idx)
{
UInt8 compression_method = source[idx + 1];
const auto codec = CompressionCodecFactory::instance().get(compression_method);
auto additional_size_at_the_end_of_buffer = codec->getAdditionalSizeAtTheEndOfBuffer();
if (compressed_buf.size() >= 1_GiB)
throw Exception(decompression_error_code, "Too large compressed size: {}", compressed_buf.size());
if (source_size < COMPRESSED_BLOCK_HEADER_SIZE)
throw Exception(decompression_error_code, "Compressed data is too short to contain a block header: {} bytes",
source_size);
{
UInt32 bytes_to_resize = 0;
if (common::addOverflow(static_cast<UInt32>(compressed_buf.size()), additional_size_at_the_end_of_buffer, bytes_to_resize))
throw Exception(decompression_error_code, "Too large compressed size: {}", compressed_buf.size());
compressed_buf.resize(compressed_buf.size() + additional_size_at_the_end_of_buffer);
}
UInt32 uncompressed_size = readDecompressedBlockSize(compressed_buf.data());
if (uncompressed_size >= 1_GiB)
throw Exception(decompression_error_code, "Too large uncompressed size: {}", uncompressed_size);
if (idx == 0 && uncompressed_size != decompressed_size)
throw Exception(decompression_error_code, "Wrong final decompressed size in codec Multiple, got {}, expected {}",
uncompressed_size, decompressed_size);
{
UInt32 bytes_to_resize = 0;
if (common::addOverflow(uncompressed_size, additional_size_at_the_end_of_buffer, bytes_to_resize))
throw Exception(decompression_error_code, "Too large uncompressed size: {}", uncompressed_size);
uncompressed_buf.resize(bytes_to_resize);
}
codec->decompress(compressed_buf.data(), source_size, uncompressed_buf.data());
uncompressed_buf.swap(compressed_buf);
/// The call to decompress will validate uncompressed_size (same readDecompressedBlockSize call as here)
source_size = uncompressed_size;
}
memcpy(dest, compressed_buf.data(), decompressed_size);
return decompressed_size;
}
VectorWithMemoryTracking<uint8_t> CompressionCodecMultiple::getCodecsBytesFromData(const char * source)
{
VectorWithMemoryTracking<uint8_t> result;
uint8_t compression_methods_size = source[0];
result.reserve(compression_methods_size);
for (size_t i = 0; i < compression_methods_size; ++i)
result.push_back(source[1 + i]);
return result;
}
bool CompressionCodecMultiple::isCompression() const
{
for (const auto & codec : codecs)
if (codec->isCompression())
return true;
return false;
}
bool CompressionCodecMultiple::isEncryption() const
{
for (const auto & codec : codecs)
if (codec->isEncryption())
return true;
return false;
}
bool CompressionCodecMultiple::isLossyCompression() const
{
for (const auto & codec : codecs)
if (codec->isLossyCompression())
return true;
return false;
}
void registerCodecMultiple(CompressionCodecFactory & factory)
{
factory.registerSimpleCompressionCodec("Multiple", static_cast<UInt8>(CompressionMethodByte::Multiple), [&] ()
{
return std::make_shared<CompressionCodecMultiple>();
});
}
}