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 | |
16 | namespace DB |
17 | { |
18 | namespace 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 | |
25 | CompressionCodecPtr CompressionCodecFactory::getDefaultCodec() const |
26 | { |
27 | return default_codec; |
28 | } |
29 | |
30 | |
31 | CompressionCodecPtr 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 | |
46 | CompressionCodecPtr 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 | |
71 | CompressionCodecPtr 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 | |
82 | CompressionCodecPtr 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 | |
95 | void 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 | |
112 | void 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 | |
120 | void 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 | |
134 | void registerCodecNone(CompressionCodecFactory & factory); |
135 | void registerCodecZSTD(CompressionCodecFactory & factory); |
136 | void registerCodecDelta(CompressionCodecFactory & factory); |
137 | void registerCodecT64(CompressionCodecFactory & factory); |
138 | void registerCodecDoubleDelta(CompressionCodecFactory & factory); |
139 | void registerCodecGorilla(CompressionCodecFactory & factory); |
140 | |
141 | CompressionCodecFactory::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 | |
155 | CompressionCodecFactory & CompressionCodecFactory::instance() |
156 | { |
157 | static CompressionCodecFactory ret; |
158 | return ret; |
159 | } |
160 | |
161 | } |
162 | |