1#include <Compression/CompressionFactory.h>
2#include <Parsers/parseQuery.h>
3#include <Parsers/ParserCreateQuery.h>
4#include <Parsers/ASTFunction.h>
5#include <Parsers/ASTIdentifier.h>
6#include <Parsers/ASTLiteral.h>
7#include <Common/typeid_cast.h>
8#include <Poco/String.h>
9#include <IO/ReadBuffer.h>
10#include <Parsers/queryToString.h>
11#include <Compression/CompressionCodecMultiple.h>
12#include <Compression/CompressionCodecLZ4.h>
13#include <IO/WriteHelpers.h>
14
15
16namespace DB
17{
18namespace ErrorCodes
19{
20 extern const int UNKNOWN_CODEC;
21 extern const int UNEXPECTED_AST_STRUCTURE;
22 extern const int DATA_TYPE_CANNOT_HAVE_ARGUMENTS;
23}
24
25CompressionCodecPtr CompressionCodecFactory::getDefaultCodec() const
26{
27 return default_codec;
28}
29
30
31CompressionCodecPtr CompressionCodecFactory::get(const String & family_name, std::optional<int> level) const
32{
33 if (level)
34 {
35 auto literal = std::make_shared<ASTLiteral>(static_cast<UInt64>(*level));
36 return get(makeASTFunction("CODEC", makeASTFunction(Poco::toUpper(family_name), literal)));
37 }
38 else
39 {
40 auto identifier = std::make_shared<ASTIdentifier>(Poco::toUpper(family_name));
41 return get(makeASTFunction("CODEC", identifier));
42 }
43}
44
45
46CompressionCodecPtr CompressionCodecFactory::get(const ASTPtr & ast, DataTypePtr column_type) const
47{
48 if (const auto * func = ast->as<ASTFunction>())
49 {
50 Codecs codecs;
51 codecs.reserve(func->arguments->children.size());
52 for (const auto & inner_codec_ast : func->arguments->children)
53 {
54 if (const auto * family_name = inner_codec_ast->as<ASTIdentifier>())
55 codecs.emplace_back(getImpl(family_name->name, {}, column_type));
56 else if (const auto * ast_func = inner_codec_ast->as<ASTFunction>())
57 codecs.emplace_back(getImpl(ast_func->name, ast_func->arguments, column_type));
58 else
59 throw Exception("Unexpected AST element for compression codec", ErrorCodes::UNEXPECTED_AST_STRUCTURE);
60 }
61
62 if (codecs.size() == 1)
63 return codecs.back();
64 else if (codecs.size() > 1)
65 return std::make_shared<CompressionCodecMultiple>(codecs);
66 }
67
68 throw Exception("Unknown codec family: " + queryToString(ast), ErrorCodes::UNKNOWN_CODEC);
69}
70
71CompressionCodecPtr CompressionCodecFactory::get(const UInt8 byte_code) const
72{
73 const auto family_code_and_creator = family_code_with_codec.find(byte_code);
74
75 if (family_code_and_creator == family_code_with_codec.end())
76 throw Exception("Unknown codec family code: " + toString(byte_code), ErrorCodes::UNKNOWN_CODEC);
77
78 return family_code_and_creator->second({}, nullptr);
79}
80
81
82CompressionCodecPtr CompressionCodecFactory::getImpl(const String & family_name, const ASTPtr & arguments, DataTypePtr column_type) const
83{
84 if (family_name == "Multiple")
85 throw Exception("Codec Multiple cannot be specified directly", ErrorCodes::UNKNOWN_CODEC);
86
87 const auto family_and_creator = family_name_with_codec.find(family_name);
88
89 if (family_and_creator == family_name_with_codec.end())
90 throw Exception("Unknown codec family: " + family_name, ErrorCodes::UNKNOWN_CODEC);
91
92 return family_and_creator->second(arguments, column_type);
93}
94
95void CompressionCodecFactory::registerCompressionCodecWithType(
96 const String & family_name,
97 std::optional<UInt8> byte_code,
98 CreatorWithType creator)
99{
100 if (creator == nullptr)
101 throw Exception("CompressionCodecFactory: the codec family " + family_name + " has been provided a null constructor",
102 ErrorCodes::LOGICAL_ERROR);
103
104 if (!family_name_with_codec.emplace(family_name, creator).second)
105 throw Exception("CompressionCodecFactory: the codec family name '" + family_name + "' is not unique", ErrorCodes::LOGICAL_ERROR);
106
107 if (byte_code)
108 if (!family_code_with_codec.emplace(*byte_code, creator).second)
109 throw Exception("CompressionCodecFactory: the codec family name '" + family_name + "' is not unique", ErrorCodes::LOGICAL_ERROR);
110}
111
112void CompressionCodecFactory::registerCompressionCodec(const String & family_name, std::optional<UInt8> byte_code, Creator creator)
113{
114 registerCompressionCodecWithType(family_name, byte_code, [family_name, creator](const ASTPtr & ast, DataTypePtr /* data_type */)
115 {
116 return creator(ast);
117 });
118}
119
120void CompressionCodecFactory::registerSimpleCompressionCodec(
121 const String & family_name,
122 std::optional<UInt8> byte_code,
123 SimpleCreator creator)
124{
125 registerCompressionCodec(family_name, byte_code, [family_name, creator](const ASTPtr & ast)
126 {
127 if (ast)
128 throw Exception("Compression codec " + family_name + " cannot have arguments", ErrorCodes::DATA_TYPE_CANNOT_HAVE_ARGUMENTS);
129 return creator();
130 });
131}
132
133
134void registerCodecNone(CompressionCodecFactory & factory);
135void registerCodecZSTD(CompressionCodecFactory & factory);
136void registerCodecDelta(CompressionCodecFactory & factory);
137void registerCodecT64(CompressionCodecFactory & factory);
138void registerCodecDoubleDelta(CompressionCodecFactory & factory);
139void registerCodecGorilla(CompressionCodecFactory & factory);
140
141CompressionCodecFactory::CompressionCodecFactory()
142{
143 default_codec = std::make_shared<CompressionCodecLZ4>();
144 registerCodecLZ4(*this);
145 registerCodecNone(*this);
146 registerCodecZSTD(*this);
147 registerCodecMultiple(*this);
148 registerCodecLZ4HC(*this);
149 registerCodecDelta(*this);
150 registerCodecT64(*this);
151 registerCodecDoubleDelta(*this);
152 registerCodecGorilla(*this);
153}
154
155CompressionCodecFactory & CompressionCodecFactory::instance()
156{
157 static CompressionCodecFactory ret;
158 return ret;
159}
160
161}
162