This commit is contained in:
Nikolay Degterinsky 2022-07-07 01:47:33 +00:00
parent 653e7cbeaf
commit 0c783800c8
5 changed files with 44 additions and 15 deletions

View File

@ -524,26 +524,34 @@ try
const auto & out_file_node = query_with_output->out_file->as<ASTLiteral &>();
out_file = out_file_node.value.safeGet<std::string>();
std::string compression_method;
UInt64 compression_level = 3;
std::string compression_method_string;
if (query_with_output->compression)
{
const auto & compression_method_node = query_with_output->compression->as<ASTLiteral &>();
compression_method = compression_method_node.value.safeGet<std::string>();
compression_method_string = compression_method_node.value.safeGet<std::string>();
}
CompressionMethod compression_method = chooseCompressionMethod(out_file, compression_method_string);
UInt64 compression_level = 3;
if (query_with_output->compression_level)
{
const auto & compression_level_node = query_with_output->compression_level->as<ASTLiteral &>();
bool res = compression_level_node.value.tryGet<UInt64>(compression_level);
if (!res || compression_level < 1 || compression_level > 9)
throw Exception("Invalid compression level, must be positive integer in range 1-9", ErrorCodes::BAD_ARGUMENTS);
}
auto range = getCompressionLevelRange(compression_method);
if (!res || compression_level < range.first || compression_level > range.second)
throw Exception(
ErrorCodes::BAD_ARGUMENTS,
"Invalid compression level, must be positive integer in range {}-{}",
range.first,
range.second);
}
out_file_buf = wrapWriteBufferWithCompressionMethod(
std::make_unique<WriteBufferFromFile>(out_file, DBMS_DEFAULT_BUFFER_SIZE, O_WRONLY | O_EXCL | O_CREAT),
chooseCompressionMethod(out_file, compression_method),
compression_method,
compression_level
);

View File

@ -98,6 +98,19 @@ CompressionMethod chooseCompressionMethod(const std::string & path, const std::s
ErrorCodes::NOT_IMPLEMENTED);
}
std::pair<uint64_t, uint64_t> getCompressionLevelRange(const CompressionMethod & method)
{
switch (method)
{
case CompressionMethod::Zstd:
return {1, 22};
case CompressionMethod::Lz4:
return {1, 12};
default:
return {1, 9};
}
}
static std::unique_ptr<CompressedReadBufferWrapper> createCompressedWrapper(
std::unique_ptr<ReadBuffer> nested, CompressionMethod method, size_t buf_size, char * existing_memory, size_t alignment, int zstd_window_log_max)
{

View File

@ -46,6 +46,9 @@ std::string toContentEncodingName(CompressionMethod method);
*/
CompressionMethod chooseCompressionMethod(const std::string & path, const std::string & hint);
/// Get a range of the valid compression levels for the compression method.
std::pair<uint64_t, uint64_t> getCompressionLevelRange(const CompressionMethod & method);
std::unique_ptr<ReadBuffer> wrapReadBufferWithCompressionMethod(
std::unique_ptr<ReadBuffer> nested,
CompressionMethod method,

View File

@ -97,8 +97,13 @@ bool ParserQueryWithOutput::parseImpl(Pos & pos, ASTPtr & node, Expected & expec
if (!compression.parse(pos, query_with_output.compression, expected))
return false;
ParserKeyword s_compression_level("LEVEL");
if (s_compression_level.ignore(pos, expected))
{
ParserNumber compression_level;
compression_level.parse(pos, query_with_output.compression_level, expected);
if (!compression_level.parse(pos, query_with_output.compression_level, expected))
return false;
}
}
query_with_output.children.push_back(query_with_output.out_file);

View File

@ -14,8 +14,8 @@ mkdir "${WORKING_FOLDER_02353}"
for m in gz br xz zst lz4 bz2
do
${CLICKHOUSE_CLIENT} --query "SELECT number, 'Hello, world!' FROM numbers(6000) INTO OUTFILE '${WORKING_FOLDER_02353}/${m}_1.${m}' COMPRESSION '${m}' 1"
${CLICKHOUSE_CLIENT} --query "SELECT number, 'Hello, world!' FROM numbers(6000) INTO OUTFILE '${WORKING_FOLDER_02353}/${m}_9.${m}' COMPRESSION '${m}' 9"
${CLICKHOUSE_CLIENT} --query "SELECT number, 'Hello, world!' FROM numbers(6000) INTO OUTFILE '${WORKING_FOLDER_02353}/${m}_1.${m}' COMPRESSION '${m}' LEVEL 1"
${CLICKHOUSE_CLIENT} --query "SELECT number, 'Hello, world!' FROM numbers(6000) INTO OUTFILE '${WORKING_FOLDER_02353}/${m}_9.${m}' COMPRESSION '${m}' LEVEL 9"
${CLICKHOUSE_CLIENT} --query "SELECT count(), max(x), avg(length(s)) FROM file('${WORKING_FOLDER_02353}/${m}_1.${m}', 'TabSeparated', 'x UInt32, s String')"
${CLICKHOUSE_CLIENT} --query "SELECT count(), max(x), avg(length(s)) FROM file('${WORKING_FOLDER_02353}/${m}_9.${m}', 'TabSeparated', 'x UInt32, s String')"