1 | #include "RedisBlockInputStream.h" |
2 | |
3 | #if USE_POCO_REDIS |
4 | |
5 | # include <string> |
6 | # include <vector> |
7 | |
8 | # include <Poco/Redis/Array.h> |
9 | # include <Poco/Redis/Client.h> |
10 | # include <Poco/Redis/Command.h> |
11 | # include <Poco/Redis/Type.h> |
12 | |
13 | # include <Columns/ColumnNullable.h> |
14 | # include <Columns/ColumnString.h> |
15 | # include <Columns/ColumnsNumber.h> |
16 | # include <IO/ReadHelpers.h> |
17 | # include <IO/WriteHelpers.h> |
18 | # include <ext/range.h> |
19 | |
20 | # include "DictionaryStructure.h" |
21 | |
22 | |
23 | namespace DB |
24 | { |
25 | namespace ErrorCodes |
26 | { |
27 | extern const int TYPE_MISMATCH; |
28 | extern const int LOGICAL_ERROR; |
29 | extern const int NUMBER_OF_COLUMNS_DOESNT_MATCH; |
30 | extern const int INTERNAL_REDIS_ERROR; |
31 | } |
32 | |
33 | |
34 | RedisBlockInputStream::RedisBlockInputStream( |
35 | const std::shared_ptr<Poco::Redis::Client> & client_, |
36 | const RedisArray & keys_, |
37 | const RedisStorageType & storage_type_, |
38 | const DB::Block & sample_block, |
39 | const size_t max_block_size_) |
40 | : client(client_), keys(keys_), storage_type(storage_type_), max_block_size{max_block_size_} |
41 | { |
42 | description.init(sample_block); |
43 | } |
44 | |
45 | RedisBlockInputStream::~RedisBlockInputStream() = default; |
46 | |
47 | |
48 | namespace |
49 | { |
50 | using ValueType = ExternalResultDescription::ValueType; |
51 | |
52 | template <typename T> |
53 | inline void insert(IColumn & column, const String & stringValue) |
54 | { |
55 | assert_cast<ColumnVector<T> &>(column).insertValue(parse<T>(stringValue)); |
56 | } |
57 | |
58 | void insertValue(IColumn & column, const ValueType type, const Poco::Redis::BulkString & bulk_string) |
59 | { |
60 | if (bulk_string.isNull()) |
61 | throw Exception{"Type mismatch, expected not Null String" , ErrorCodes::TYPE_MISMATCH}; |
62 | |
63 | String stringValue = bulk_string.value(); |
64 | switch (type) |
65 | { |
66 | case ValueType::vtUInt8: |
67 | insert<UInt8>(column, stringValue); |
68 | break; |
69 | case ValueType::vtUInt16: |
70 | insert<UInt16>(column, stringValue); |
71 | break; |
72 | case ValueType::vtUInt32: |
73 | insert<UInt32>(column, stringValue); |
74 | break; |
75 | case ValueType::vtUInt64: |
76 | insert<UInt64>(column, stringValue); |
77 | break; |
78 | case ValueType::vtInt8: |
79 | insert<Int8>(column, stringValue); |
80 | break; |
81 | case ValueType::vtInt16: |
82 | insert<Int16>(column, stringValue); |
83 | break; |
84 | case ValueType::vtInt32: |
85 | insert<Int32>(column, stringValue); |
86 | break; |
87 | case ValueType::vtInt64: |
88 | insert<Int64>(column, stringValue); |
89 | break; |
90 | case ValueType::vtFloat32: |
91 | insert<Float32>(column, stringValue); |
92 | break; |
93 | case ValueType::vtFloat64: |
94 | insert<Float64>(column, stringValue); |
95 | break; |
96 | case ValueType::vtString: |
97 | assert_cast<ColumnString &>(column).insert(parse<String>(stringValue)); |
98 | break; |
99 | case ValueType::vtDate: |
100 | assert_cast<ColumnUInt16 &>(column).insertValue(parse<LocalDate>(stringValue).getDayNum()); |
101 | break; |
102 | case ValueType::vtDateTime: |
103 | assert_cast<ColumnUInt32 &>(column).insertValue(static_cast<UInt32>(parse<LocalDateTime>(stringValue))); |
104 | break; |
105 | case ValueType::vtUUID: |
106 | assert_cast<ColumnUInt128 &>(column).insertValue(parse<UUID>(stringValue)); |
107 | break; |
108 | } |
109 | } |
110 | } |
111 | |
112 | |
113 | Block RedisBlockInputStream::readImpl() |
114 | { |
115 | if (keys.isNull() || description.sample_block.rows() == 0 || cursor >= keys.size()) |
116 | all_read = true; |
117 | |
118 | if (all_read) |
119 | return {}; |
120 | |
121 | const size_t size = description.sample_block.columns(); |
122 | MutableColumns columns(size); |
123 | |
124 | for (const auto i : ext::range(0, size)) |
125 | columns[i] = description.sample_block.getByPosition(i).column->cloneEmpty(); |
126 | |
127 | const auto insertValueByIdx = [this, &columns](size_t idx, const auto & value) |
128 | { |
129 | if (description.types[idx].second) |
130 | { |
131 | ColumnNullable & column_nullable = static_cast<ColumnNullable &>(*columns[idx]); |
132 | insertValue(column_nullable.getNestedColumn(), description.types[idx].first, value); |
133 | column_nullable.getNullMapData().emplace_back(0); |
134 | } |
135 | else |
136 | insertValue(*columns[idx], description.types[idx].first, value); |
137 | }; |
138 | |
139 | if (storage_type == RedisStorageType::HASH_MAP) |
140 | { |
141 | size_t num_rows = 0; |
142 | for (; cursor < keys.size(); ++cursor) |
143 | { |
144 | const auto & keys_array = keys.get<RedisArray>(cursor); |
145 | if (keys_array.size() < 2) |
146 | { |
147 | throw Exception{"Too low keys in request to source: " + DB::toString(keys_array.size()) |
148 | + ", expected 2 or more" , ErrorCodes::LOGICAL_ERROR}; |
149 | } |
150 | |
151 | if (num_rows + keys_array.size() - 1 > max_block_size) |
152 | break; |
153 | |
154 | Poco::Redis::Command command_for_values("HMGET" ); |
155 | for (auto it = keys_array.begin(); it != keys_array.end(); ++it) |
156 | command_for_values.addRedisType(*it); |
157 | |
158 | auto values = client->execute<RedisArray>(command_for_values); |
159 | |
160 | if (keys_array.size() != values.size() + 1) // 'HMGET' primary_key secondary_keys |
161 | throw Exception{"Inconsistent sizes of keys and values in Redis request" , |
162 | ErrorCodes::NUMBER_OF_COLUMNS_DOESNT_MATCH}; |
163 | |
164 | const auto & primary_key = keys_array.get<RedisBulkString>(0); |
165 | for (size_t i = 0; i < values.size(); ++i) |
166 | { |
167 | const auto & secondary_key = keys_array.get<RedisBulkString>(i + 1); |
168 | const auto & value = values.get<RedisBulkString>(i); |
169 | |
170 | /// null string means 'no value for requested key' |
171 | if (!value.isNull()) |
172 | { |
173 | insertValueByIdx(0, primary_key); |
174 | insertValueByIdx(1, secondary_key); |
175 | insertValueByIdx(2, value); |
176 | ++num_rows; |
177 | } |
178 | } |
179 | } |
180 | } |
181 | else |
182 | { |
183 | Poco::Redis::Command command_for_values("MGET" ); |
184 | |
185 | size_t need_values = std::min(max_block_size, keys.size() - cursor); |
186 | for (size_t i = 0; i < need_values; ++i) |
187 | command_for_values.add(keys.get<RedisBulkString>(cursor + i)); |
188 | |
189 | auto values = client->execute<RedisArray>(command_for_values); |
190 | if (values.size() != need_values) |
191 | throw Exception{"Inconsistent sizes of keys and values in Redis request" , ErrorCodes::INTERNAL_REDIS_ERROR}; |
192 | |
193 | for (size_t i = 0; i < values.size(); ++i) |
194 | { |
195 | const auto & key = keys.get<RedisBulkString>(cursor + i); |
196 | const auto & value = values.get<RedisBulkString>(i); |
197 | |
198 | /// Null string means 'no value for requested key' |
199 | if (!value.isNull()) |
200 | { |
201 | insertValueByIdx(0, key); |
202 | insertValueByIdx(1, value); |
203 | } |
204 | } |
205 | cursor += need_values; |
206 | } |
207 | |
208 | return description.sample_block.cloneWithColumns(std::move(columns)); |
209 | } |
210 | } |
211 | |
212 | #endif |
213 | |