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
23namespace 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