1 | #include <DataTypes/DataTypesNumber.h> |
2 | #include <DataTypes/DataTypesDecimal.h> |
3 | #include <DataTypes/DataTypeArray.h> |
4 | #include <DataTypes/DataTypeString.h> |
5 | #include <DataTypes/DataTypeFixedString.h> |
6 | #include <DataTypes/DataTypeTuple.h> |
7 | #include <DataTypes/DataTypeNullable.h> |
8 | #include <DataTypes/NumberTraits.h> |
9 | #include <DataTypes/getLeastSupertype.h> |
10 | #include <Columns/ColumnVector.h> |
11 | #include <Columns/ColumnDecimal.h> |
12 | #include <Columns/ColumnString.h> |
13 | #include <Columns/ColumnConst.h> |
14 | #include <Columns/ColumnArray.h> |
15 | #include <Columns/ColumnFixedString.h> |
16 | #include <Columns/ColumnTuple.h> |
17 | #include <Columns/ColumnNullable.h> |
18 | #include <Common/typeid_cast.h> |
19 | #include <Common/assert_cast.h> |
20 | #include <Functions/IFunctionImpl.h> |
21 | #include <Functions/FunctionHelpers.h> |
22 | #include <Functions/GatherUtils/GatherUtils.h> |
23 | #include <Functions/GatherUtils/Algorithms.h> |
24 | #include <Functions/FunctionIfBase.h> |
25 | #include <Functions/FunctionFactory.h> |
26 | #include <Interpreters/castColumn.h> |
27 | |
28 | |
29 | namespace DB |
30 | { |
31 | |
32 | namespace ErrorCodes |
33 | { |
34 | extern const int NOT_IMPLEMENTED; |
35 | } |
36 | |
37 | |
38 | using namespace GatherUtils; |
39 | |
40 | /** Selection function by condition: if(cond, then, else). |
41 | * cond - UInt8 |
42 | * then, else - numeric types for which there is a general type, or dates, datetimes, or strings, or arrays of these types. |
43 | */ |
44 | |
45 | |
46 | template <typename A, typename B, typename ResultType> |
47 | struct NumIfImpl |
48 | { |
49 | using ArrayCond = PaddedPODArray<UInt8>; |
50 | using ArrayA = PaddedPODArray<A>; |
51 | using ArrayB = PaddedPODArray<B>; |
52 | using ColVecResult = ColumnVector<ResultType>; |
53 | |
54 | static void vector_vector(const ArrayCond & cond, const ArrayA & a, const ArrayB & b, Block & block, size_t result, UInt32) |
55 | { |
56 | size_t size = cond.size(); |
57 | auto col_res = ColVecResult::create(size); |
58 | typename ColVecResult::Container & res = col_res->getData(); |
59 | |
60 | for (size_t i = 0; i < size; ++i) |
61 | res[i] = cond[i] ? static_cast<ResultType>(a[i]) : static_cast<ResultType>(b[i]); |
62 | block.getByPosition(result).column = std::move(col_res); |
63 | } |
64 | |
65 | static void vector_constant(const ArrayCond & cond, const ArrayA & a, B b, Block & block, size_t result, UInt32) |
66 | { |
67 | size_t size = cond.size(); |
68 | auto col_res = ColVecResult::create(size); |
69 | typename ColVecResult::Container & res = col_res->getData(); |
70 | |
71 | for (size_t i = 0; i < size; ++i) |
72 | res[i] = cond[i] ? static_cast<ResultType>(a[i]) : static_cast<ResultType>(b); |
73 | block.getByPosition(result).column = std::move(col_res); |
74 | } |
75 | |
76 | static void constant_vector(const ArrayCond & cond, A a, const ArrayB & b, Block & block, size_t result, UInt32) |
77 | { |
78 | size_t size = cond.size(); |
79 | auto col_res = ColVecResult::create(size); |
80 | typename ColVecResult::Container & res = col_res->getData(); |
81 | |
82 | for (size_t i = 0; i < size; ++i) |
83 | res[i] = cond[i] ? static_cast<ResultType>(a) : static_cast<ResultType>(b[i]); |
84 | block.getByPosition(result).column = std::move(col_res); |
85 | } |
86 | |
87 | static void constant_constant(const ArrayCond & cond, A a, B b, Block & block, size_t result, UInt32) |
88 | { |
89 | size_t size = cond.size(); |
90 | auto col_res = ColVecResult::create(size); |
91 | typename ColVecResult::Container & res = col_res->getData(); |
92 | |
93 | for (size_t i = 0; i < size; ++i) |
94 | res[i] = cond[i] ? static_cast<ResultType>(a) : static_cast<ResultType>(b); |
95 | block.getByPosition(result).column = std::move(col_res); |
96 | } |
97 | }; |
98 | |
99 | template <typename A, typename B, typename R> |
100 | struct NumIfImpl<Decimal<A>, Decimal<B>, Decimal<R>> |
101 | { |
102 | using ResultType = Decimal<R>; |
103 | using ArrayCond = PaddedPODArray<UInt8>; |
104 | using ArrayA = DecimalPaddedPODArray<Decimal<A>>; |
105 | using ArrayB = DecimalPaddedPODArray<Decimal<B>>; |
106 | using ColVecResult = ColumnDecimal<ResultType>; |
107 | |
108 | static void vector_vector(const ArrayCond & cond, const ArrayA & a, const ArrayB & b, Block & block, size_t result, UInt32 scale) |
109 | { |
110 | size_t size = cond.size(); |
111 | auto col_res = ColVecResult::create(size, scale); |
112 | typename ColVecResult::Container & res = col_res->getData(); |
113 | |
114 | for (size_t i = 0; i < size; ++i) |
115 | res[i] = cond[i] ? static_cast<ResultType>(a[i]) : static_cast<ResultType>(b[i]); |
116 | block.getByPosition(result).column = std::move(col_res); |
117 | } |
118 | |
119 | static void vector_constant(const ArrayCond & cond, const ArrayA & a, B b, Block & block, size_t result, UInt32 scale) |
120 | { |
121 | size_t size = cond.size(); |
122 | auto col_res = ColVecResult::create(size, scale); |
123 | typename ColVecResult::Container & res = col_res->getData(); |
124 | |
125 | for (size_t i = 0; i < size; ++i) |
126 | res[i] = cond[i] ? static_cast<ResultType>(a[i]) : static_cast<ResultType>(b); |
127 | block.getByPosition(result).column = std::move(col_res); |
128 | } |
129 | |
130 | static void constant_vector(const ArrayCond & cond, A a, const ArrayB & b, Block & block, size_t result, UInt32 scale) |
131 | { |
132 | size_t size = cond.size(); |
133 | auto col_res = ColVecResult::create(size, scale); |
134 | typename ColVecResult::Container & res = col_res->getData(); |
135 | |
136 | for (size_t i = 0; i < size; ++i) |
137 | res[i] = cond[i] ? static_cast<ResultType>(a) : static_cast<ResultType>(b[i]); |
138 | block.getByPosition(result).column = std::move(col_res); |
139 | } |
140 | |
141 | static void constant_constant(const ArrayCond & cond, A a, B b, Block & block, size_t result, UInt32 scale) |
142 | { |
143 | size_t size = cond.size(); |
144 | auto col_res = ColVecResult::create(size, scale); |
145 | typename ColVecResult::Container & res = col_res->getData(); |
146 | |
147 | for (size_t i = 0; i < size; ++i) |
148 | res[i] = cond[i] ? static_cast<ResultType>(a) : static_cast<ResultType>(b); |
149 | block.getByPosition(result).column = std::move(col_res); |
150 | } |
151 | }; |
152 | |
153 | template <typename A, typename B> |
154 | struct NumIfImpl<A, B, NumberTraits::Error> |
155 | { |
156 | private: |
157 | [[noreturn]] static void throw_error() |
158 | { |
159 | throw Exception("Internal logic error: invalid types of arguments 2 and 3 of if" , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
160 | } |
161 | public: |
162 | template <typename... Args> static void vector_vector(Args &&...) { throw_error(); } |
163 | template <typename... Args> static void vector_constant(Args &&...) { throw_error(); } |
164 | template <typename... Args> static void constant_vector(Args &&...) { throw_error(); } |
165 | template <typename... Args> static void constant_constant(Args &&...) { throw_error(); } |
166 | }; |
167 | |
168 | |
169 | class FunctionIf : public FunctionIfBase</*null_is_false=*/false> |
170 | { |
171 | public: |
172 | static constexpr auto name = "if" ; |
173 | static FunctionPtr create(const Context & context) { return std::make_shared<FunctionIf>(context); } |
174 | FunctionIf(const Context & context_) : context(context_) {} |
175 | |
176 | private: |
177 | template <typename T0, typename T1> |
178 | static constexpr bool allow_arrays = !std::is_same_v<T0, UInt128> && !std::is_same_v<T1, UInt128>; |
179 | |
180 | template <typename T0, typename T1> |
181 | static UInt32 decimalScale(Block & block [[maybe_unused]], const ColumnNumbers & arguments [[maybe_unused]]) |
182 | { |
183 | if constexpr (IsDecimalNumber<T0> && IsDecimalNumber<T1>) |
184 | { |
185 | UInt32 left_scale = getDecimalScale(*block.getByPosition(arguments[1]).type); |
186 | UInt32 right_scale = getDecimalScale(*block.getByPosition(arguments[2]).type); |
187 | if (left_scale != right_scale) |
188 | throw Exception("Conditional functions with different Decimal scales" , ErrorCodes::NOT_IMPLEMENTED); |
189 | return left_scale; |
190 | } |
191 | else |
192 | return std::numeric_limits<UInt32>::max(); |
193 | } |
194 | |
195 | template <typename T0, typename T1, typename ColVecT0, typename ColVecT1> |
196 | bool executeRightType( |
197 | const ColumnUInt8 * cond_col, |
198 | Block & block, |
199 | const ColumnNumbers & arguments, |
200 | size_t result, |
201 | const ColVecT0 * col_left) |
202 | { |
203 | using ResultType = typename NumberTraits::ResultOfIf<T0, T1>::Type; |
204 | |
205 | const IColumn * col_right_untyped = block.getByPosition(arguments[2]).column.get(); |
206 | UInt32 scale = decimalScale<T0, T1>(block, arguments); |
207 | |
208 | if (auto col_right_vec = checkAndGetColumn<ColVecT1>(col_right_untyped)) |
209 | { |
210 | NumIfImpl<T0, T1, ResultType>::vector_vector( |
211 | cond_col->getData(), col_left->getData(), col_right_vec->getData(), block, result, scale); |
212 | return true; |
213 | } |
214 | else if (auto col_right_const = checkAndGetColumnConst<ColVecT1>(col_right_untyped)) |
215 | { |
216 | NumIfImpl<T0, T1, ResultType>::vector_constant( |
217 | cond_col->getData(), col_left->getData(), col_right_const->template getValue<T1>(), block, result, scale); |
218 | return true; |
219 | } |
220 | |
221 | return false; |
222 | } |
223 | |
224 | template <typename T0, typename T1, typename ColVecT0, typename ColVecT1> |
225 | bool executeConstRightType( |
226 | const ColumnUInt8 * cond_col, |
227 | Block & block, |
228 | const ColumnNumbers & arguments, |
229 | size_t result, |
230 | const ColumnConst * col_left) |
231 | { |
232 | using ResultType = typename NumberTraits::ResultOfIf<T0, T1>::Type; |
233 | |
234 | const IColumn * col_right_untyped = block.getByPosition(arguments[2]).column.get(); |
235 | UInt32 scale = decimalScale<T0, T1>(block, arguments); |
236 | |
237 | if (auto col_right_vec = checkAndGetColumn<ColVecT1>(col_right_untyped)) |
238 | { |
239 | NumIfImpl<T0, T1, ResultType>::constant_vector( |
240 | cond_col->getData(), col_left->template getValue<T0>(), col_right_vec->getData(), block, result, scale); |
241 | return true; |
242 | } |
243 | else if (auto col_right_const = checkAndGetColumnConst<ColVecT1>(col_right_untyped)) |
244 | { |
245 | NumIfImpl<T0, T1, ResultType>::constant_constant( |
246 | cond_col->getData(), col_left->template getValue<T0>(), col_right_const->template getValue<T1>(), block, result, scale); |
247 | return true; |
248 | } |
249 | |
250 | return false; |
251 | } |
252 | |
253 | template <typename T0, typename T1, typename ColVecT0, typename ColVecT1> |
254 | bool executeRightTypeArray( |
255 | [[maybe_unused]] const ColumnUInt8 * cond_col, |
256 | [[maybe_unused]] Block & block, |
257 | [[maybe_unused]] const ColumnNumbers & arguments, |
258 | [[maybe_unused]] size_t result, |
259 | [[maybe_unused]] const ColumnArray * col_left_array, |
260 | [[maybe_unused]] size_t input_rows_count) |
261 | { |
262 | if constexpr (std::is_same_v<NumberTraits::Error, typename NumberTraits::ResultOfIf<T0, T1>::Type>) |
263 | return false; |
264 | else if constexpr (allow_arrays<T0, T1>) |
265 | { |
266 | using ResultType = typename NumberTraits::ResultOfIf<T0, T1>::Type; |
267 | |
268 | const IColumn * col_right_untyped = block.getByPosition(arguments[2]).column.get(); |
269 | |
270 | if (auto col_right_array = checkAndGetColumn<ColumnArray>(col_right_untyped)) |
271 | { |
272 | const ColVecT1 * col_right_vec = checkAndGetColumn<ColVecT1>(&col_right_array->getData()); |
273 | if (!col_right_vec) |
274 | return false; |
275 | |
276 | auto res = block.getByPosition(result).type->createColumn(); |
277 | |
278 | conditional( |
279 | NumericArraySource<T0>(*col_left_array), |
280 | NumericArraySource<T1>(*col_right_array), |
281 | NumericArraySink<ResultType>(assert_cast<ColumnArray &>(*res), input_rows_count), |
282 | cond_col->getData()); |
283 | |
284 | block.getByPosition(result).column = std::move(res); |
285 | return true; |
286 | } |
287 | else if (auto col_right_const_array = checkAndGetColumnConst<ColumnArray>(col_right_untyped)) |
288 | { |
289 | const ColumnArray * col_right_const_array_data = checkAndGetColumn<ColumnArray>(&col_right_const_array->getDataColumn()); |
290 | if (!checkColumn<ColVecT1>(&col_right_const_array_data->getData())) |
291 | return false; |
292 | |
293 | auto res = block.getByPosition(result).type->createColumn(); |
294 | |
295 | conditional( |
296 | NumericArraySource<T0>(*col_left_array), |
297 | ConstSource<NumericArraySource<T1>>(*col_right_const_array), |
298 | NumericArraySink<ResultType>(assert_cast<ColumnArray &>(*res), input_rows_count), |
299 | cond_col->getData()); |
300 | |
301 | block.getByPosition(result).column = std::move(res); |
302 | return true; |
303 | } |
304 | } |
305 | |
306 | return false; |
307 | } |
308 | |
309 | template <typename T0, typename T1, typename ColVecT0, typename ColVecT1> |
310 | bool executeConstRightTypeArray( |
311 | [[maybe_unused]] const ColumnUInt8 * cond_col, |
312 | [[maybe_unused]] Block & block, |
313 | [[maybe_unused]] const ColumnNumbers & arguments, |
314 | [[maybe_unused]] size_t result, |
315 | [[maybe_unused]] const ColumnConst * col_left_const_array, |
316 | [[maybe_unused]] size_t input_rows_count) |
317 | { |
318 | if constexpr (std::is_same_v<NumberTraits::Error, typename NumberTraits::ResultOfIf<T0, T1>::Type>) |
319 | return false; |
320 | else if constexpr (allow_arrays<T0, T1>) |
321 | { |
322 | using ResultType = typename NumberTraits::ResultOfIf<T0, T1>::Type; |
323 | |
324 | const IColumn * col_right_untyped = block.getByPosition(arguments[2]).column.get(); |
325 | |
326 | if (auto col_right_array = checkAndGetColumn<ColumnArray>(col_right_untyped)) |
327 | { |
328 | const ColVecT1 * col_right_vec = checkAndGetColumn<ColVecT1>(&col_right_array->getData()); |
329 | |
330 | if (!col_right_vec) |
331 | return false; |
332 | |
333 | auto res = block.getByPosition(result).type->createColumn(); |
334 | |
335 | conditional( |
336 | ConstSource<NumericArraySource<T0>>(*col_left_const_array), |
337 | NumericArraySource<T1>(*col_right_array), |
338 | NumericArraySink<ResultType>(assert_cast<ColumnArray &>(*res), input_rows_count), |
339 | cond_col->getData()); |
340 | |
341 | block.getByPosition(result).column = std::move(res); |
342 | return true; |
343 | } |
344 | else if (auto col_right_const_array = checkAndGetColumnConst<ColumnArray>(col_right_untyped)) |
345 | { |
346 | const ColumnArray * col_right_const_array_data = checkAndGetColumn<ColumnArray>(&col_right_const_array->getDataColumn()); |
347 | if (!checkColumn<ColVecT1>(&col_right_const_array_data->getData())) |
348 | return false; |
349 | |
350 | auto res = block.getByPosition(result).type->createColumn(); |
351 | |
352 | conditional( |
353 | ConstSource<NumericArraySource<T0>>(*col_left_const_array), |
354 | ConstSource<NumericArraySource<T1>>(*col_right_const_array), |
355 | NumericArraySink<ResultType>(assert_cast<ColumnArray &>(*res), input_rows_count), |
356 | cond_col->getData()); |
357 | |
358 | block.getByPosition(result).column = std::move(res); |
359 | return true; |
360 | } |
361 | } |
362 | |
363 | return false; |
364 | } |
365 | |
366 | template <typename T0, typename T1> |
367 | bool executeTyped(const ColumnUInt8 * cond_col, Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
368 | { |
369 | using ColVecT0 = std::conditional_t<IsDecimalNumber<T0>, ColumnDecimal<T0>, ColumnVector<T0>>; |
370 | using ColVecT1 = std::conditional_t<IsDecimalNumber<T1>, ColumnDecimal<T1>, ColumnVector<T1>>; |
371 | |
372 | const IColumn * col_left_untyped = block.getByPosition(arguments[1]).column.get(); |
373 | |
374 | bool left_ok = false; |
375 | bool right_ok = false; |
376 | |
377 | if (auto col_left = checkAndGetColumn<ColVecT0>(col_left_untyped)) |
378 | { |
379 | left_ok = true; |
380 | right_ok = executeRightType<T0, T1, ColVecT0, ColVecT1>(cond_col, block, arguments, result, col_left); |
381 | } |
382 | else if (auto col_const_left = checkAndGetColumnConst<ColVecT0>(col_left_untyped)) |
383 | { |
384 | left_ok = true; |
385 | right_ok = executeConstRightType<T0, T1, ColVecT0, ColVecT1>(cond_col, block, arguments, result, col_const_left); |
386 | } |
387 | else if (auto col_arr_left = checkAndGetColumn<ColumnArray>(col_left_untyped)) |
388 | { |
389 | if (auto col_arr_left_elems = checkAndGetColumn<ColVecT0>(&col_arr_left->getData())) |
390 | { |
391 | left_ok = true; |
392 | right_ok = executeRightTypeArray<T0, T1, ColVecT0, ColVecT1>( |
393 | cond_col, block, arguments, result, col_arr_left, input_rows_count); |
394 | } |
395 | } |
396 | else if (auto col_const_arr_left = checkAndGetColumnConst<ColumnArray>(col_left_untyped)) |
397 | { |
398 | if (checkColumn<ColVecT0>(&assert_cast<const ColumnArray &>(col_const_arr_left->getDataColumn()).getData())) |
399 | { |
400 | left_ok = true; |
401 | right_ok = executeConstRightTypeArray<T0, T1, ColVecT0, ColVecT1>( |
402 | cond_col, block, arguments, result, col_const_arr_left, input_rows_count); |
403 | } |
404 | } |
405 | |
406 | if (!left_ok) |
407 | return false; |
408 | |
409 | ColumnWithTypeAndName & right_column_typed = block.getByPosition(arguments[2]); |
410 | if (!right_ok) |
411 | throw Exception("Illegal column " + right_column_typed.column->getName() + " of third argument of function " + getName(), |
412 | ErrorCodes::ILLEGAL_COLUMN); |
413 | |
414 | return true; |
415 | } |
416 | |
417 | bool executeString(const ColumnUInt8 * cond_col, Block & block, const ColumnNumbers & arguments, size_t result) |
418 | { |
419 | const IColumn * col_then_untyped = block.getByPosition(arguments[1]).column.get(); |
420 | const IColumn * col_else_untyped = block.getByPosition(arguments[2]).column.get(); |
421 | |
422 | const ColumnString * col_then = checkAndGetColumn<ColumnString>(col_then_untyped); |
423 | const ColumnString * col_else = checkAndGetColumn<ColumnString>(col_else_untyped); |
424 | const ColumnFixedString * col_then_fixed = checkAndGetColumn<ColumnFixedString>(col_then_untyped); |
425 | const ColumnFixedString * col_else_fixed = checkAndGetColumn<ColumnFixedString>(col_else_untyped); |
426 | const ColumnConst * col_then_const = checkAndGetColumnConst<ColumnString>(col_then_untyped); |
427 | const ColumnConst * col_else_const = checkAndGetColumnConst<ColumnString>(col_else_untyped); |
428 | const ColumnConst * col_then_const_fixed = checkAndGetColumnConst<ColumnFixedString>(col_then_untyped); |
429 | const ColumnConst * col_else_const_fixed = checkAndGetColumnConst<ColumnFixedString>(col_else_untyped); |
430 | |
431 | const PaddedPODArray<UInt8> & cond_data = cond_col->getData(); |
432 | size_t rows = cond_data.size(); |
433 | |
434 | if ((col_then_fixed || col_then_const_fixed) |
435 | && (col_else_fixed || col_else_const_fixed)) |
436 | { |
437 | /// The result is FixedString. |
438 | |
439 | auto col_res_untyped = block.getByPosition(result).type->createColumn(); |
440 | ColumnFixedString * col_res = assert_cast<ColumnFixedString *>(col_res_untyped.get()); |
441 | auto sink = FixedStringSink(*col_res, rows); |
442 | |
443 | if (col_then_fixed && col_else_fixed) |
444 | conditional(FixedStringSource(*col_then_fixed), FixedStringSource(*col_else_fixed), sink, cond_data); |
445 | else if (col_then_fixed && col_else_const_fixed) |
446 | conditional(FixedStringSource(*col_then_fixed), ConstSource<FixedStringSource>(*col_else_const_fixed), sink, cond_data); |
447 | else if (col_then_const_fixed && col_else_fixed) |
448 | conditional(ConstSource<FixedStringSource>(*col_then_const_fixed), FixedStringSource(*col_else_fixed), sink, cond_data); |
449 | else if (col_then_const_fixed && col_else_const_fixed) |
450 | conditional(ConstSource<FixedStringSource>(*col_then_const_fixed), ConstSource<FixedStringSource>(*col_else_const_fixed), sink, cond_data); |
451 | |
452 | block.getByPosition(result).column = std::move(col_res_untyped); |
453 | return true; |
454 | } |
455 | |
456 | if ((col_then || col_then_const || col_then_fixed || col_then_const_fixed) |
457 | && (col_else || col_else_const || col_else_fixed || col_else_const_fixed)) |
458 | { |
459 | /// The result is String. |
460 | auto col_res = ColumnString::create(); |
461 | auto sink = StringSink(*col_res, rows); |
462 | |
463 | if (col_then && col_else) |
464 | conditional(StringSource(*col_then), StringSource(*col_else), sink, cond_data); |
465 | else if (col_then && col_else_const) |
466 | conditional(StringSource(*col_then), ConstSource<StringSource>(*col_else_const), sink, cond_data); |
467 | else if (col_then_const && col_else) |
468 | conditional(ConstSource<StringSource>(*col_then_const), StringSource(*col_else), sink, cond_data); |
469 | else if (col_then_const && col_else_const) |
470 | conditional(ConstSource<StringSource>(*col_then_const), ConstSource<StringSource>(*col_else_const), sink, cond_data); |
471 | else if (col_then && col_else_fixed) |
472 | conditional(StringSource(*col_then), FixedStringSource(*col_else_fixed), sink, cond_data); |
473 | else if (col_then_fixed && col_else) |
474 | conditional(FixedStringSource(*col_then_fixed), StringSource(*col_else), sink, cond_data); |
475 | else if (col_then_const && col_else_fixed) |
476 | conditional(ConstSource<StringSource>(*col_then_const), FixedStringSource(*col_else_fixed), sink, cond_data); |
477 | else if (col_then_fixed && col_else_const) |
478 | conditional(FixedStringSource(*col_then_fixed), ConstSource<StringSource>(*col_else_const), sink, cond_data); |
479 | else if (col_then && col_else_const_fixed) |
480 | conditional(StringSource(*col_then), ConstSource<FixedStringSource>(*col_else_const_fixed), sink, cond_data); |
481 | else if (col_then_const_fixed && col_else) |
482 | conditional(ConstSource<FixedStringSource>(*col_then_const_fixed), StringSource(*col_else), sink, cond_data); |
483 | else if (col_then_const && col_else_const_fixed) |
484 | conditional(ConstSource<StringSource>(*col_then_const), ConstSource<FixedStringSource>(*col_else_const_fixed), sink, cond_data); |
485 | else if (col_then_const_fixed && col_else_const) |
486 | conditional(ConstSource<FixedStringSource>(*col_then_const_fixed), ConstSource<StringSource>(*col_else_const), sink, cond_data); |
487 | |
488 | block.getByPosition(result).column = std::move(col_res); |
489 | return true; |
490 | } |
491 | |
492 | return false; |
493 | } |
494 | |
495 | bool executeGenericArray(const ColumnUInt8 * cond_col, Block & block, const ColumnNumbers & arguments, size_t result) |
496 | { |
497 | /// For generic implementation, arrays must be of same type. |
498 | if (!block.getByPosition(arguments[1]).type->equals(*block.getByPosition(arguments[2]).type)) |
499 | return false; |
500 | |
501 | const IColumn * col_then_untyped = block.getByPosition(arguments[1]).column.get(); |
502 | const IColumn * col_else_untyped = block.getByPosition(arguments[2]).column.get(); |
503 | |
504 | const ColumnArray * col_arr_then = checkAndGetColumn<ColumnArray>(col_then_untyped); |
505 | const ColumnArray * col_arr_else = checkAndGetColumn<ColumnArray>(col_else_untyped); |
506 | const ColumnConst * col_arr_then_const = checkAndGetColumnConst<ColumnArray>(col_then_untyped); |
507 | const ColumnConst * col_arr_else_const = checkAndGetColumnConst<ColumnArray>(col_else_untyped); |
508 | |
509 | const PaddedPODArray<UInt8> & cond_data = cond_col->getData(); |
510 | size_t rows = cond_data.size(); |
511 | |
512 | if ((col_arr_then || col_arr_then_const) |
513 | && (col_arr_else || col_arr_else_const)) |
514 | { |
515 | auto res = block.getByPosition(result).type->createColumn(); |
516 | auto col_res = assert_cast<ColumnArray *>(res.get()); |
517 | |
518 | if (col_arr_then && col_arr_else) |
519 | conditional(GenericArraySource(*col_arr_then), GenericArraySource(*col_arr_else), GenericArraySink(*col_res, rows), cond_data); |
520 | else if (col_arr_then && col_arr_else_const) |
521 | conditional(GenericArraySource(*col_arr_then), ConstSource<GenericArraySource>(*col_arr_else_const), GenericArraySink(*col_res, rows), cond_data); |
522 | else if (col_arr_then_const && col_arr_else) |
523 | conditional(ConstSource<GenericArraySource>(*col_arr_then_const), GenericArraySource(*col_arr_else), GenericArraySink(*col_res, rows), cond_data); |
524 | else if (col_arr_then_const && col_arr_else_const) |
525 | conditional(ConstSource<GenericArraySource>(*col_arr_then_const), ConstSource<GenericArraySource>(*col_arr_else_const), GenericArraySink(*col_res, rows), cond_data); |
526 | else |
527 | return false; |
528 | |
529 | block.getByPosition(result).column = std::move(res); |
530 | return true; |
531 | } |
532 | |
533 | return false; |
534 | } |
535 | |
536 | bool executeTuple(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
537 | { |
538 | /// Calculate function for each corresponding elements of tuples. |
539 | |
540 | const ColumnWithTypeAndName & arg1 = block.getByPosition(arguments[1]); |
541 | const ColumnWithTypeAndName & arg2 = block.getByPosition(arguments[2]); |
542 | |
543 | Columns col1_contents; |
544 | Columns col2_contents; |
545 | |
546 | if (const ColumnTuple * tuple1 = typeid_cast<const ColumnTuple *>(arg1.column.get())) |
547 | col1_contents = tuple1->getColumnsCopy(); |
548 | else if (const ColumnConst * const_tuple = checkAndGetColumnConst<ColumnTuple>(arg1.column.get())) |
549 | col1_contents = convertConstTupleToConstantElements(*const_tuple); |
550 | else |
551 | return false; |
552 | |
553 | if (const ColumnTuple * tuple2 = typeid_cast<const ColumnTuple *>(arg2.column.get())) |
554 | col2_contents = tuple2->getColumnsCopy(); |
555 | else if (const ColumnConst * const_tuple = checkAndGetColumnConst<ColumnTuple>(arg2.column.get())) |
556 | col2_contents = convertConstTupleToConstantElements(*const_tuple); |
557 | else |
558 | return false; |
559 | |
560 | const DataTypeTuple & type1 = static_cast<const DataTypeTuple &>(*arg1.type); |
561 | const DataTypeTuple & type2 = static_cast<const DataTypeTuple &>(*arg2.type); |
562 | |
563 | Block temporary_block; |
564 | temporary_block.insert(block.getByPosition(arguments[0])); |
565 | |
566 | size_t tuple_size = type1.getElements().size(); |
567 | Columns tuple_columns(tuple_size); |
568 | |
569 | for (size_t i = 0; i < tuple_size; ++i) |
570 | { |
571 | temporary_block.insert({nullptr, |
572 | getReturnTypeImpl({std::make_shared<DataTypeUInt8>(), type1.getElements()[i], type2.getElements()[i]}), |
573 | {}}); |
574 | |
575 | temporary_block.insert({col1_contents[i], type1.getElements()[i], {}}); |
576 | temporary_block.insert({col2_contents[i], type2.getElements()[i], {}}); |
577 | |
578 | /// temporary_block will be: cond, res_0, ..., res_i, then_i, else_i |
579 | executeImpl(temporary_block, {0, i + 2, i + 3}, i + 1, input_rows_count); |
580 | temporary_block.erase(i + 3); |
581 | temporary_block.erase(i + 2); |
582 | |
583 | tuple_columns[i] = temporary_block.getByPosition(i + 1).column; |
584 | } |
585 | |
586 | /// temporary_block is: cond, res_0, res_1, res_2... |
587 | |
588 | block.getByPosition(result).column = ColumnTuple::create(tuple_columns); |
589 | return true; |
590 | } |
591 | |
592 | void executeGeneric(const ColumnUInt8 * cond_col, Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
593 | { |
594 | /// Convert both columns to the common type (if needed). |
595 | |
596 | const ColumnWithTypeAndName & arg1 = block.getByPosition(arguments[1]); |
597 | const ColumnWithTypeAndName & arg2 = block.getByPosition(arguments[2]); |
598 | |
599 | DataTypePtr common_type = getLeastSupertype({arg1.type, arg2.type}); |
600 | |
601 | ColumnPtr col_then = castColumn(arg1, common_type, context); |
602 | ColumnPtr col_else = castColumn(arg2, common_type, context); |
603 | |
604 | MutableColumnPtr result_column = common_type->createColumn(); |
605 | result_column->reserve(input_rows_count); |
606 | |
607 | bool then_is_const = isColumnConst(*col_then); |
608 | bool else_is_const = isColumnConst(*col_else); |
609 | |
610 | const auto & cond_array = cond_col->getData(); |
611 | |
612 | if (then_is_const && else_is_const) |
613 | { |
614 | const IColumn & then_nested_column = assert_cast<const ColumnConst &>(*col_then).getDataColumn(); |
615 | const IColumn & else_nested_column = assert_cast<const ColumnConst &>(*col_else).getDataColumn(); |
616 | |
617 | for (size_t i = 0; i < input_rows_count; ++i) |
618 | { |
619 | if (cond_array[i]) |
620 | result_column->insertFrom(then_nested_column, 0); |
621 | else |
622 | result_column->insertFrom(else_nested_column, 0); |
623 | } |
624 | } |
625 | else if (then_is_const) |
626 | { |
627 | const IColumn & then_nested_column = assert_cast<const ColumnConst &>(*col_then).getDataColumn(); |
628 | |
629 | for (size_t i = 0; i < input_rows_count; ++i) |
630 | { |
631 | if (cond_array[i]) |
632 | result_column->insertFrom(then_nested_column, 0); |
633 | else |
634 | result_column->insertFrom(*col_else, i); |
635 | } |
636 | } |
637 | else if (else_is_const) |
638 | { |
639 | const IColumn & else_nested_column = assert_cast<const ColumnConst &>(*col_else).getDataColumn(); |
640 | |
641 | for (size_t i = 0; i < input_rows_count; ++i) |
642 | { |
643 | if (cond_array[i]) |
644 | result_column->insertFrom(*col_then, i); |
645 | else |
646 | result_column->insertFrom(else_nested_column, 0); |
647 | } |
648 | } |
649 | else |
650 | { |
651 | for (size_t i = 0; i < input_rows_count; ++i) |
652 | result_column->insertFrom(cond_array[i] ? *col_then : *col_else, i); |
653 | } |
654 | |
655 | block.getByPosition(result).column = std::move(result_column); |
656 | } |
657 | |
658 | bool executeForNullableCondition(Block & block, const ColumnNumbers & arguments, size_t result, size_t /*input_rows_count*/) |
659 | { |
660 | const ColumnWithTypeAndName & arg_cond = block.getByPosition(arguments[0]); |
661 | bool cond_is_null = arg_cond.column->onlyNull(); |
662 | |
663 | if (cond_is_null) |
664 | { |
665 | block.getByPosition(result).column = std::move(block.getByPosition(arguments[2]).column); |
666 | return true; |
667 | } |
668 | |
669 | if (auto * nullable = checkAndGetColumn<ColumnNullable>(*arg_cond.column)) |
670 | { |
671 | Block temporary_block |
672 | { |
673 | { nullable->getNestedColumnPtr(), removeNullable(arg_cond.type), arg_cond.name }, |
674 | block.getByPosition(arguments[1]), |
675 | block.getByPosition(arguments[2]), |
676 | block.getByPosition(result) |
677 | }; |
678 | |
679 | executeImpl(temporary_block, {0, 1, 2}, 3, temporary_block.rows()); |
680 | |
681 | block.getByPosition(result).column = std::move(temporary_block.getByPosition(3).column); |
682 | return true; |
683 | } |
684 | |
685 | return false; |
686 | } |
687 | |
688 | static ColumnPtr materializeColumnIfConst(const ColumnPtr & column) |
689 | { |
690 | return column->convertToFullColumnIfConst(); |
691 | } |
692 | |
693 | static ColumnPtr makeNullableColumnIfNot(const ColumnPtr & column) |
694 | { |
695 | if (isColumnNullable(*column)) |
696 | return column; |
697 | |
698 | return ColumnNullable::create( |
699 | materializeColumnIfConst(column), ColumnUInt8::create(column->size(), 0)); |
700 | } |
701 | |
702 | static ColumnPtr getNestedColumn(const ColumnPtr & column) |
703 | { |
704 | if (auto * nullable = checkAndGetColumn<ColumnNullable>(*column)) |
705 | return nullable->getNestedColumnPtr(); |
706 | |
707 | return column; |
708 | } |
709 | |
710 | bool executeForNullableThenElse(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
711 | { |
712 | const ColumnWithTypeAndName & arg_cond = block.getByPosition(arguments[0]); |
713 | const ColumnWithTypeAndName & arg_then = block.getByPosition(arguments[1]); |
714 | const ColumnWithTypeAndName & arg_else = block.getByPosition(arguments[2]); |
715 | |
716 | auto * then_is_nullable = checkAndGetColumn<ColumnNullable>(*arg_then.column); |
717 | auto * else_is_nullable = checkAndGetColumn<ColumnNullable>(*arg_else.column); |
718 | |
719 | if (!then_is_nullable && !else_is_nullable) |
720 | return false; |
721 | |
722 | /** Calculate null mask of result and nested column separately. |
723 | */ |
724 | ColumnPtr result_null_mask; |
725 | |
726 | { |
727 | Block temporary_block( |
728 | { |
729 | arg_cond, |
730 | { |
731 | then_is_nullable |
732 | ? then_is_nullable->getNullMapColumnPtr() |
733 | : DataTypeUInt8().createColumnConstWithDefaultValue(input_rows_count), |
734 | std::make_shared<DataTypeUInt8>(), |
735 | "" |
736 | }, |
737 | { |
738 | else_is_nullable |
739 | ? else_is_nullable->getNullMapColumnPtr() |
740 | : DataTypeUInt8().createColumnConstWithDefaultValue(input_rows_count), |
741 | std::make_shared<DataTypeUInt8>(), |
742 | "" |
743 | }, |
744 | { |
745 | nullptr, |
746 | std::make_shared<DataTypeUInt8>(), |
747 | "" |
748 | } |
749 | }); |
750 | |
751 | executeImpl(temporary_block, {0, 1, 2}, 3, temporary_block.rows()); |
752 | |
753 | result_null_mask = temporary_block.getByPosition(3).column; |
754 | } |
755 | |
756 | ColumnPtr result_nested_column; |
757 | |
758 | { |
759 | Block temporary_block( |
760 | { |
761 | arg_cond, |
762 | { |
763 | getNestedColumn(arg_then.column), |
764 | removeNullable(arg_then.type), |
765 | "" |
766 | }, |
767 | { |
768 | getNestedColumn(arg_else.column), |
769 | removeNullable(arg_else.type), |
770 | "" |
771 | }, |
772 | { |
773 | nullptr, |
774 | removeNullable(block.getByPosition(result).type), |
775 | "" |
776 | } |
777 | }); |
778 | |
779 | executeImpl(temporary_block, {0, 1, 2}, 3, temporary_block.rows()); |
780 | |
781 | result_nested_column = temporary_block.getByPosition(3).column; |
782 | } |
783 | |
784 | block.getByPosition(result).column = ColumnNullable::create( |
785 | materializeColumnIfConst(result_nested_column), materializeColumnIfConst(result_null_mask)); |
786 | return true; |
787 | } |
788 | |
789 | bool executeForNullThenElse(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
790 | { |
791 | const ColumnWithTypeAndName & arg_cond = block.getByPosition(arguments[0]); |
792 | const ColumnWithTypeAndName & arg_then = block.getByPosition(arguments[1]); |
793 | const ColumnWithTypeAndName & arg_else = block.getByPosition(arguments[2]); |
794 | |
795 | bool then_is_null = arg_then.column->onlyNull(); |
796 | bool else_is_null = arg_else.column->onlyNull(); |
797 | |
798 | if (!then_is_null && !else_is_null) |
799 | return false; |
800 | |
801 | if (then_is_null && else_is_null) |
802 | { |
803 | block.getByPosition(result).column = block.getByPosition(result).type->createColumnConstWithDefaultValue(input_rows_count); |
804 | return true; |
805 | } |
806 | |
807 | const ColumnUInt8 * cond_col = typeid_cast<const ColumnUInt8 *>(arg_cond.column.get()); |
808 | const ColumnConst * cond_const_col = checkAndGetColumnConst<ColumnVector<UInt8>>(arg_cond.column.get()); |
809 | |
810 | /// If then is NULL, we create Nullable column with null mask OR-ed with condition. |
811 | if (then_is_null) |
812 | { |
813 | if (cond_col) |
814 | { |
815 | if (isColumnNullable(*arg_else.column)) |
816 | { |
817 | auto arg_else_column = arg_else.column; |
818 | auto result_column = (*std::move(arg_else_column)).mutate(); |
819 | assert_cast<ColumnNullable &>(*result_column).applyNullMap(assert_cast<const ColumnUInt8 &>(*arg_cond.column)); |
820 | block.getByPosition(result).column = std::move(result_column); |
821 | } |
822 | else |
823 | { |
824 | block.getByPosition(result).column = ColumnNullable::create( |
825 | materializeColumnIfConst(arg_else.column), arg_cond.column); |
826 | } |
827 | } |
828 | else if (cond_const_col) |
829 | { |
830 | if (cond_const_col->getValue<UInt8>()) |
831 | block.getByPosition(result).column = block.getByPosition(result).type->createColumn()->cloneResized(input_rows_count); |
832 | else |
833 | block.getByPosition(result).column = makeNullableColumnIfNot(arg_else.column); |
834 | } |
835 | else |
836 | throw Exception("Illegal column " + arg_cond.column->getName() + " of first argument of function " + getName() |
837 | + ". Must be ColumnUInt8 or ColumnConstUInt8." , |
838 | ErrorCodes::ILLEGAL_COLUMN); |
839 | return true; |
840 | } |
841 | |
842 | /// If else is NULL, we create Nullable column with null mask OR-ed with negated condition. |
843 | if (else_is_null) |
844 | { |
845 | if (cond_col) |
846 | { |
847 | size_t size = input_rows_count; |
848 | auto & null_map_data = cond_col->getData(); |
849 | |
850 | auto negated_null_map = ColumnUInt8::create(); |
851 | auto & negated_null_map_data = negated_null_map->getData(); |
852 | negated_null_map_data.resize(size); |
853 | |
854 | for (size_t i = 0; i < size; ++i) |
855 | negated_null_map_data[i] = !null_map_data[i]; |
856 | |
857 | if (isColumnNullable(*arg_then.column)) |
858 | { |
859 | auto arg_then_column = arg_then.column; |
860 | auto result_column = (*std::move(arg_then_column)).mutate(); |
861 | assert_cast<ColumnNullable &>(*result_column).applyNegatedNullMap(assert_cast<const ColumnUInt8 &>(*arg_cond.column)); |
862 | block.getByPosition(result).column = std::move(result_column); |
863 | } |
864 | else |
865 | { |
866 | block.getByPosition(result).column = ColumnNullable::create( |
867 | materializeColumnIfConst(arg_then.column), std::move(negated_null_map)); |
868 | } |
869 | } |
870 | else if (cond_const_col) |
871 | { |
872 | if (cond_const_col->getValue<UInt8>()) |
873 | block.getByPosition(result).column = makeNullableColumnIfNot(arg_then.column); |
874 | else |
875 | block.getByPosition(result).column = block.getByPosition(result).type->createColumn()->cloneResized(input_rows_count); |
876 | } |
877 | else |
878 | throw Exception("Illegal column " + arg_cond.column->getName() + " of first argument of function " + getName() |
879 | + ". Must be ColumnUInt8 or ColumnConstUInt8." , |
880 | ErrorCodes::ILLEGAL_COLUMN); |
881 | return true; |
882 | } |
883 | |
884 | return false; |
885 | } |
886 | |
887 | public: |
888 | String getName() const override |
889 | { |
890 | return name; |
891 | } |
892 | |
893 | size_t getNumberOfArguments() const override { return 3; } |
894 | |
895 | bool useDefaultImplementationForNulls() const override { return false; } |
896 | ColumnNumbers getArgumentsThatDontImplyNullableReturnType(size_t /*number_of_arguments*/) const override { return {0}; } |
897 | |
898 | /// Get result types by argument types. If the function does not apply to these arguments, throw an exception. |
899 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
900 | { |
901 | if (arguments[0]->onlyNull()) |
902 | return arguments[2]; |
903 | |
904 | if (arguments[0]->isNullable()) |
905 | return getReturnTypeImpl({ |
906 | removeNullable(arguments[0]), arguments[1], arguments[2]}); |
907 | |
908 | if (!WhichDataType(arguments[0]).isUInt8()) |
909 | throw Exception("Illegal type " + arguments[0]->getName() + " of first argument (condition) of function if. Must be UInt8." , |
910 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
911 | |
912 | return getLeastSupertype({arguments[1], arguments[2]}); |
913 | } |
914 | |
915 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
916 | { |
917 | if (executeForNullableCondition(block, arguments, result, input_rows_count) |
918 | || executeForNullThenElse(block, arguments, result, input_rows_count) |
919 | || executeForNullableThenElse(block, arguments, result, input_rows_count)) |
920 | return; |
921 | |
922 | const ColumnWithTypeAndName & arg_cond = block.getByPosition(arguments[0]); |
923 | const ColumnWithTypeAndName & arg_then = block.getByPosition(arguments[1]); |
924 | const ColumnWithTypeAndName & arg_else = block.getByPosition(arguments[2]); |
925 | |
926 | /// A case for identical then and else (pointers are the same). |
927 | if (arg_then.column.get() == arg_else.column.get()) |
928 | { |
929 | /// Just point result to them. |
930 | block.getByPosition(result).column = arg_then.column; |
931 | return; |
932 | } |
933 | |
934 | const ColumnUInt8 * cond_col = typeid_cast<const ColumnUInt8 *>(arg_cond.column.get()); |
935 | const ColumnConst * cond_const_col = checkAndGetColumnConst<ColumnVector<UInt8>>(arg_cond.column.get()); |
936 | ColumnPtr materialized_cond_col; |
937 | |
938 | if (cond_const_col) |
939 | { |
940 | if (arg_then.type->equals(*arg_else.type)) |
941 | { |
942 | block.getByPosition(result).column = cond_const_col->getValue<UInt8>() |
943 | ? arg_then.column |
944 | : arg_else.column; |
945 | return; |
946 | } |
947 | else |
948 | { |
949 | materialized_cond_col = cond_const_col->convertToFullColumn(); |
950 | cond_col = typeid_cast<const ColumnUInt8 *>(&*materialized_cond_col); |
951 | } |
952 | } |
953 | |
954 | if (!cond_col) |
955 | throw Exception("Illegal column " + arg_cond.column->getName() + " of first argument of function " + getName() |
956 | + ". Must be ColumnUInt8 or ColumnConstUInt8." , |
957 | ErrorCodes::ILLEGAL_COLUMN); |
958 | |
959 | auto call = [&](const auto & types) -> bool |
960 | { |
961 | using Types = std::decay_t<decltype(types)>; |
962 | using T0 = typename Types::LeftType; |
963 | using T1 = typename Types::RightType; |
964 | |
965 | if constexpr (IsDecimalNumber<T0> == IsDecimalNumber<T1>) |
966 | return executeTyped<T0, T1>(cond_col, block, arguments, result, input_rows_count); |
967 | else |
968 | throw Exception("Conditional function with Decimal and non Decimal" , ErrorCodes::NOT_IMPLEMENTED); |
969 | }; |
970 | |
971 | TypeIndex left_id = arg_then.type->getTypeId(); |
972 | TypeIndex right_id = arg_else.type->getTypeId(); |
973 | |
974 | if (auto left_array = checkAndGetDataType<DataTypeArray>(arg_then.type.get())) |
975 | left_id = left_array->getNestedType()->getTypeId(); |
976 | |
977 | if (auto rigth_array = checkAndGetDataType<DataTypeArray>(arg_else.type.get())) |
978 | right_id = rigth_array->getNestedType()->getTypeId(); |
979 | |
980 | if (!(callOnBasicTypes<true, true, true, false>(left_id, right_id, call) |
981 | || executeTyped<UInt128, UInt128>(cond_col, block, arguments, result, input_rows_count) |
982 | || executeString(cond_col, block, arguments, result) |
983 | || executeGenericArray(cond_col, block, arguments, result) |
984 | || executeTuple(block, arguments, result, input_rows_count))) |
985 | { |
986 | executeGeneric(cond_col, block, arguments, result, input_rows_count); |
987 | } |
988 | } |
989 | |
990 | const Context & context; |
991 | }; |
992 | |
993 | void registerFunctionIf(FunctionFactory & factory) |
994 | { |
995 | factory.registerFunction<FunctionIf>(FunctionFactory::CaseInsensitive); |
996 | } |
997 | |
998 | } |
999 | |