1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Functions/FunctionHelpers.h> |
4 | #include <DataTypes/DataTypeArray.h> |
5 | #include <DataTypes/DataTypeNullable.h> |
6 | #include <DataTypes/DataTypeTuple.h> |
7 | #include <Core/ColumnNumbers.h> |
8 | #include <Columns/ColumnArray.h> |
9 | #include <Columns/ColumnNullable.h> |
10 | #include <Columns/ColumnsNumber.h> |
11 | #include <Columns/ColumnString.h> |
12 | #include <Columns/ColumnTuple.h> |
13 | #include <Common/typeid_cast.h> |
14 | #include <Common/assert_cast.h> |
15 | |
16 | |
17 | namespace DB |
18 | { |
19 | |
20 | namespace ErrorCodes |
21 | { |
22 | extern const int LOGICAL_ERROR; |
23 | extern const int ILLEGAL_COLUMN; |
24 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
25 | extern const int ZERO_ARRAY_OR_TUPLE_INDEX; |
26 | } |
27 | |
28 | namespace ArrayImpl |
29 | { |
30 | class NullMapBuilder; |
31 | } |
32 | |
33 | /** arrayElement(arr, i) - get the array element by index. If index is not constant and out of range - return default value of data type. |
34 | * The index begins with 1. Also, the index can be negative - then it is counted from the end of the array. |
35 | */ |
36 | class FunctionArrayElement : public IFunction |
37 | { |
38 | public: |
39 | static constexpr auto name = "arrayElement" ; |
40 | static FunctionPtr create(const Context & context); |
41 | |
42 | String getName() const override; |
43 | |
44 | bool useDefaultImplementationForConstants() const override { return true; } |
45 | size_t getNumberOfArguments() const override { return 2; } |
46 | |
47 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override; |
48 | |
49 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override; |
50 | |
51 | private: |
52 | void perform(Block & block, const ColumnNumbers & arguments, size_t result, |
53 | ArrayImpl::NullMapBuilder & builder, size_t input_rows_count); |
54 | |
55 | template <typename DataType> |
56 | bool executeNumberConst(Block & block, const ColumnNumbers & arguments, size_t result, const Field & index, |
57 | ArrayImpl::NullMapBuilder & builder); |
58 | |
59 | template <typename IndexType, typename DataType> |
60 | bool executeNumber(Block & block, const ColumnNumbers & arguments, size_t result, const PaddedPODArray<IndexType> & indices, |
61 | ArrayImpl::NullMapBuilder & builder); |
62 | |
63 | bool executeStringConst(Block & block, const ColumnNumbers & arguments, size_t result, const Field & index, |
64 | ArrayImpl::NullMapBuilder & builder); |
65 | |
66 | template <typename IndexType> |
67 | bool executeString(Block & block, const ColumnNumbers & arguments, size_t result, const PaddedPODArray<IndexType> & indices, |
68 | ArrayImpl::NullMapBuilder & builder); |
69 | |
70 | bool executeGenericConst(Block & block, const ColumnNumbers & arguments, size_t result, const Field & index, |
71 | ArrayImpl::NullMapBuilder & builder); |
72 | |
73 | template <typename IndexType> |
74 | bool executeGeneric(Block & block, const ColumnNumbers & arguments, size_t result, const PaddedPODArray<IndexType> & indices, |
75 | ArrayImpl::NullMapBuilder & builder); |
76 | |
77 | template <typename IndexType> |
78 | bool executeConst(Block & block, const ColumnNumbers & arguments, size_t result, |
79 | const PaddedPODArray <IndexType> & indices, ArrayImpl::NullMapBuilder & builder, |
80 | size_t input_rows_count); |
81 | |
82 | template <typename IndexType> |
83 | bool executeArgument(Block & block, const ColumnNumbers & arguments, size_t result, |
84 | ArrayImpl::NullMapBuilder & builder, size_t input_rows_count); |
85 | |
86 | /** For a tuple array, the function is evaluated component-wise for each element of the tuple. |
87 | */ |
88 | bool executeTuple(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count); |
89 | }; |
90 | |
91 | |
92 | namespace ArrayImpl |
93 | { |
94 | |
95 | class NullMapBuilder |
96 | { |
97 | public: |
98 | operator bool() const { return src_null_map; } |
99 | bool operator!() const { return !src_null_map; } |
100 | |
101 | void initSource(const UInt8 * src_null_map_) |
102 | { |
103 | src_null_map = src_null_map_; |
104 | } |
105 | |
106 | void initSink(size_t size) |
107 | { |
108 | auto sink = ColumnUInt8::create(size); |
109 | sink_null_map = sink->getData().data(); |
110 | sink_null_map_holder = std::move(sink); |
111 | } |
112 | |
113 | void update(size_t from) |
114 | { |
115 | sink_null_map[index] = bool(src_null_map && src_null_map[from]); |
116 | ++index; |
117 | } |
118 | |
119 | void update() |
120 | { |
121 | sink_null_map[index] = bool(src_null_map); |
122 | ++index; |
123 | } |
124 | |
125 | ColumnPtr getNullMapColumnPtr() && { return std::move(sink_null_map_holder); } |
126 | |
127 | private: |
128 | const UInt8 * src_null_map = nullptr; |
129 | UInt8 * sink_null_map = nullptr; |
130 | MutableColumnPtr sink_null_map_holder; |
131 | size_t index = 0; |
132 | }; |
133 | |
134 | } |
135 | |
136 | namespace |
137 | { |
138 | |
139 | template <typename T> |
140 | struct ArrayElementNumImpl |
141 | { |
142 | /** Implementation for constant index. |
143 | * If negative = false - index is from beginning of array, started from 0. |
144 | * If negative = true - index is from end of array, started from 0. |
145 | */ |
146 | template <bool negative> |
147 | static void vectorConst( |
148 | const PaddedPODArray<T> & data, const ColumnArray::Offsets & offsets, |
149 | const ColumnArray::Offset index, |
150 | PaddedPODArray<T> & result, ArrayImpl::NullMapBuilder & builder) |
151 | { |
152 | size_t size = offsets.size(); |
153 | result.resize(size); |
154 | |
155 | ColumnArray::Offset current_offset = 0; |
156 | for (size_t i = 0; i < size; ++i) |
157 | { |
158 | size_t array_size = offsets[i] - current_offset; |
159 | |
160 | if (index < array_size) |
161 | { |
162 | size_t j = !negative ? (current_offset + index) : (offsets[i] - index - 1); |
163 | result[i] = data[j]; |
164 | if (builder) |
165 | builder.update(j); |
166 | } |
167 | else |
168 | { |
169 | result[i] = T(); |
170 | |
171 | if (builder) |
172 | builder.update(); |
173 | } |
174 | |
175 | current_offset = offsets[i]; |
176 | } |
177 | } |
178 | |
179 | /** Implementation for non-constant index. |
180 | */ |
181 | template <typename TIndex> |
182 | static void vector( |
183 | const PaddedPODArray<T> & data, const ColumnArray::Offsets & offsets, |
184 | const PaddedPODArray<TIndex> & indices, |
185 | PaddedPODArray<T> & result, ArrayImpl::NullMapBuilder & builder) |
186 | { |
187 | size_t size = offsets.size(); |
188 | result.resize(size); |
189 | |
190 | ColumnArray::Offset current_offset = 0; |
191 | for (size_t i = 0; i < size; ++i) |
192 | { |
193 | size_t array_size = offsets[i] - current_offset; |
194 | |
195 | TIndex index = indices[i]; |
196 | if (index > 0 && static_cast<size_t>(index) <= array_size) |
197 | { |
198 | size_t j = current_offset + index - 1; |
199 | result[i] = data[j]; |
200 | |
201 | if (builder) |
202 | builder.update(j); |
203 | } |
204 | else if (index < 0 && static_cast<size_t>(-index) <= array_size) |
205 | { |
206 | size_t j = offsets[i] + index; |
207 | result[i] = data[j]; |
208 | |
209 | if (builder) |
210 | builder.update(j); |
211 | } |
212 | else |
213 | { |
214 | result[i] = T(); |
215 | |
216 | if (builder) |
217 | builder.update(); |
218 | } |
219 | |
220 | current_offset = offsets[i]; |
221 | } |
222 | } |
223 | }; |
224 | |
225 | struct ArrayElementStringImpl |
226 | { |
227 | template <bool negative> |
228 | static void vectorConst( |
229 | const ColumnString::Chars & data, const ColumnArray::Offsets & offsets, const ColumnString::Offsets & string_offsets, |
230 | const ColumnArray::Offset index, |
231 | ColumnString::Chars & result_data, ColumnArray::Offsets & result_offsets, |
232 | ArrayImpl::NullMapBuilder & builder) |
233 | { |
234 | size_t size = offsets.size(); |
235 | result_offsets.resize(size); |
236 | result_data.reserve(data.size()); |
237 | |
238 | ColumnArray::Offset current_offset = 0; |
239 | ColumnArray::Offset current_result_offset = 0; |
240 | for (size_t i = 0; i < size; ++i) |
241 | { |
242 | size_t array_size = offsets[i] - current_offset; |
243 | |
244 | if (index < array_size) |
245 | { |
246 | size_t adjusted_index = !negative ? index : (array_size - index - 1); |
247 | |
248 | size_t j = current_offset + adjusted_index; |
249 | if (builder) |
250 | builder.update(j); |
251 | |
252 | ColumnArray::Offset string_pos = current_offset == 0 && adjusted_index == 0 |
253 | ? 0 |
254 | : string_offsets[current_offset + adjusted_index - 1]; |
255 | |
256 | ColumnArray::Offset string_size = string_offsets[current_offset + adjusted_index] - string_pos; |
257 | |
258 | result_data.resize(current_result_offset + string_size); |
259 | memcpySmallAllowReadWriteOverflow15(&result_data[current_result_offset], &data[string_pos], string_size); |
260 | current_result_offset += string_size; |
261 | result_offsets[i] = current_result_offset; |
262 | } |
263 | else |
264 | { |
265 | /// Insert an empty row. |
266 | result_data.resize(current_result_offset + 1); |
267 | result_data[current_result_offset] = 0; |
268 | current_result_offset += 1; |
269 | result_offsets[i] = current_result_offset; |
270 | |
271 | if (builder) |
272 | builder.update(); |
273 | } |
274 | |
275 | current_offset = offsets[i]; |
276 | } |
277 | } |
278 | |
279 | /** Implementation for non-constant index. |
280 | */ |
281 | template <typename TIndex> |
282 | static void vector( |
283 | const ColumnString::Chars & data, const ColumnArray::Offsets & offsets, const ColumnString::Offsets & string_offsets, |
284 | const PaddedPODArray<TIndex> & indices, |
285 | ColumnString::Chars & result_data, ColumnArray::Offsets & result_offsets, |
286 | ArrayImpl::NullMapBuilder & builder) |
287 | { |
288 | size_t size = offsets.size(); |
289 | result_offsets.resize(size); |
290 | result_data.reserve(data.size()); |
291 | |
292 | ColumnArray::Offset current_offset = 0; |
293 | ColumnArray::Offset current_result_offset = 0; |
294 | for (size_t i = 0; i < size; ++i) |
295 | { |
296 | size_t array_size = offsets[i] - current_offset; |
297 | size_t adjusted_index; /// index in array from zero |
298 | |
299 | TIndex index = indices[i]; |
300 | if (index > 0 && static_cast<size_t>(index) <= array_size) |
301 | adjusted_index = index - 1; |
302 | else if (index < 0 && static_cast<size_t>(-index) <= array_size) |
303 | adjusted_index = array_size + index; |
304 | else |
305 | adjusted_index = array_size; /// means no element should be taken |
306 | |
307 | if (adjusted_index < array_size) |
308 | { |
309 | size_t j = current_offset + adjusted_index; |
310 | if (builder) |
311 | builder.update(j); |
312 | |
313 | ColumnArray::Offset string_pos = current_offset == 0 && adjusted_index == 0 |
314 | ? 0 |
315 | : string_offsets[current_offset + adjusted_index - 1]; |
316 | |
317 | ColumnArray::Offset string_size = string_offsets[current_offset + adjusted_index] - string_pos; |
318 | |
319 | result_data.resize(current_result_offset + string_size); |
320 | memcpySmallAllowReadWriteOverflow15(&result_data[current_result_offset], &data[string_pos], string_size); |
321 | current_result_offset += string_size; |
322 | result_offsets[i] = current_result_offset; |
323 | } |
324 | else |
325 | { |
326 | /// Insert empty string |
327 | result_data.resize(current_result_offset + 1); |
328 | result_data[current_result_offset] = 0; |
329 | current_result_offset += 1; |
330 | result_offsets[i] = current_result_offset; |
331 | |
332 | if (builder) |
333 | builder.update(); |
334 | } |
335 | |
336 | current_offset = offsets[i]; |
337 | } |
338 | } |
339 | }; |
340 | |
341 | /// Generic implementation for other nested types. |
342 | struct ArrayElementGenericImpl |
343 | { |
344 | template <bool negative> |
345 | static void vectorConst( |
346 | const IColumn & data, const ColumnArray::Offsets & offsets, |
347 | const ColumnArray::Offset index, |
348 | IColumn & result, ArrayImpl::NullMapBuilder & builder) |
349 | { |
350 | size_t size = offsets.size(); |
351 | result.reserve(size); |
352 | |
353 | ColumnArray::Offset current_offset = 0; |
354 | for (size_t i = 0; i < size; ++i) |
355 | { |
356 | size_t array_size = offsets[i] - current_offset; |
357 | |
358 | if (index < array_size) |
359 | { |
360 | size_t j = !negative ? current_offset + index : offsets[i] - index - 1; |
361 | result.insertFrom(data, j); |
362 | if (builder) |
363 | builder.update(j); |
364 | } |
365 | else |
366 | { |
367 | result.insertDefault(); |
368 | if (builder) |
369 | builder.update(); |
370 | } |
371 | |
372 | current_offset = offsets[i]; |
373 | } |
374 | } |
375 | |
376 | /** Implementation for non-constant index. |
377 | */ |
378 | template <typename TIndex> |
379 | static void vector( |
380 | const IColumn & data, const ColumnArray::Offsets & offsets, |
381 | const PaddedPODArray<TIndex> & indices, |
382 | IColumn & result, ArrayImpl::NullMapBuilder & builder) |
383 | { |
384 | size_t size = offsets.size(); |
385 | result.reserve(size); |
386 | |
387 | ColumnArray::Offset current_offset = 0; |
388 | for (size_t i = 0; i < size; ++i) |
389 | { |
390 | size_t array_size = offsets[i] - current_offset; |
391 | |
392 | TIndex index = indices[i]; |
393 | if (index > 0 && static_cast<size_t>(index) <= array_size) |
394 | { |
395 | size_t j = current_offset + index - 1; |
396 | result.insertFrom(data, j); |
397 | if (builder) |
398 | builder.update(j); |
399 | } |
400 | else if (index < 0 && static_cast<size_t>(-index) <= array_size) |
401 | { |
402 | size_t j = offsets[i] + index; |
403 | result.insertFrom(data, j); |
404 | if (builder) |
405 | builder.update(j); |
406 | } |
407 | else |
408 | { |
409 | result.insertDefault(); |
410 | if (builder) |
411 | builder.update(); |
412 | } |
413 | |
414 | current_offset = offsets[i]; |
415 | } |
416 | } |
417 | }; |
418 | |
419 | } |
420 | |
421 | |
422 | FunctionPtr FunctionArrayElement::create(const Context &) |
423 | { |
424 | return std::make_shared<FunctionArrayElement>(); |
425 | } |
426 | |
427 | |
428 | template <typename DataType> |
429 | bool FunctionArrayElement::executeNumberConst(Block & block, const ColumnNumbers & arguments, size_t result, const Field & index, |
430 | ArrayImpl::NullMapBuilder & builder) |
431 | { |
432 | const ColumnArray * col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
433 | |
434 | if (!col_array) |
435 | return false; |
436 | |
437 | const ColumnVector<DataType> * col_nested = checkAndGetColumn<ColumnVector<DataType>>(&col_array->getData()); |
438 | |
439 | if (!col_nested) |
440 | return false; |
441 | |
442 | auto col_res = ColumnVector<DataType>::create(); |
443 | |
444 | if (index.getType() == Field::Types::UInt64) |
445 | ArrayElementNumImpl<DataType>::template vectorConst<false>( |
446 | col_nested->getData(), col_array->getOffsets(), safeGet<UInt64>(index) - 1, col_res->getData(), builder); |
447 | else if (index.getType() == Field::Types::Int64) |
448 | ArrayElementNumImpl<DataType>::template vectorConst<true>( |
449 | col_nested->getData(), col_array->getOffsets(), -safeGet<Int64>(index) - 1, col_res->getData(), builder); |
450 | else |
451 | throw Exception("Illegal type of array index" , ErrorCodes::LOGICAL_ERROR); |
452 | |
453 | block.getByPosition(result).column = std::move(col_res); |
454 | return true; |
455 | } |
456 | |
457 | template <typename IndexType, typename DataType> |
458 | bool FunctionArrayElement::executeNumber(Block & block, const ColumnNumbers & arguments, size_t result, const PaddedPODArray<IndexType> & indices, |
459 | ArrayImpl::NullMapBuilder & builder) |
460 | { |
461 | const ColumnArray * col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
462 | |
463 | if (!col_array) |
464 | return false; |
465 | |
466 | const ColumnVector<DataType> * col_nested = checkAndGetColumn<ColumnVector<DataType>>(&col_array->getData()); |
467 | |
468 | if (!col_nested) |
469 | return false; |
470 | |
471 | auto col_res = ColumnVector<DataType>::create(); |
472 | |
473 | ArrayElementNumImpl<DataType>::template vector<IndexType>( |
474 | col_nested->getData(), col_array->getOffsets(), indices, col_res->getData(), builder); |
475 | |
476 | block.getByPosition(result).column = std::move(col_res); |
477 | return true; |
478 | } |
479 | |
480 | bool FunctionArrayElement::executeStringConst(Block & block, const ColumnNumbers & arguments, size_t result, const Field & index, |
481 | ArrayImpl::NullMapBuilder & builder) |
482 | { |
483 | const ColumnArray * col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
484 | |
485 | if (!col_array) |
486 | return false; |
487 | |
488 | const ColumnString * col_nested = checkAndGetColumn<ColumnString>(&col_array->getData()); |
489 | |
490 | if (!col_nested) |
491 | return false; |
492 | |
493 | auto col_res = ColumnString::create(); |
494 | |
495 | if (index.getType() == Field::Types::UInt64) |
496 | ArrayElementStringImpl::vectorConst<false>( |
497 | col_nested->getChars(), |
498 | col_array->getOffsets(), |
499 | col_nested->getOffsets(), |
500 | safeGet<UInt64>(index) - 1, |
501 | col_res->getChars(), |
502 | col_res->getOffsets(), |
503 | builder); |
504 | else if (index.getType() == Field::Types::Int64) |
505 | ArrayElementStringImpl::vectorConst<true>( |
506 | col_nested->getChars(), |
507 | col_array->getOffsets(), |
508 | col_nested->getOffsets(), |
509 | -safeGet<Int64>(index) - 1, |
510 | col_res->getChars(), |
511 | col_res->getOffsets(), |
512 | builder); |
513 | else |
514 | throw Exception("Illegal type of array index" , ErrorCodes::LOGICAL_ERROR); |
515 | |
516 | block.getByPosition(result).column = std::move(col_res); |
517 | return true; |
518 | } |
519 | |
520 | template <typename IndexType> |
521 | bool FunctionArrayElement::executeString(Block & block, const ColumnNumbers & arguments, size_t result, const PaddedPODArray<IndexType> & indices, |
522 | ArrayImpl::NullMapBuilder & builder) |
523 | { |
524 | const ColumnArray * col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
525 | |
526 | if (!col_array) |
527 | return false; |
528 | |
529 | const ColumnString * col_nested = checkAndGetColumn<ColumnString>(&col_array->getData()); |
530 | |
531 | if (!col_nested) |
532 | return false; |
533 | |
534 | auto col_res = ColumnString::create(); |
535 | |
536 | ArrayElementStringImpl::vector<IndexType>( |
537 | col_nested->getChars(), |
538 | col_array->getOffsets(), |
539 | col_nested->getOffsets(), |
540 | indices, |
541 | col_res->getChars(), |
542 | col_res->getOffsets(), |
543 | builder); |
544 | |
545 | block.getByPosition(result).column = std::move(col_res); |
546 | return true; |
547 | } |
548 | |
549 | bool FunctionArrayElement::executeGenericConst(Block & block, const ColumnNumbers & arguments, size_t result, const Field & index, |
550 | ArrayImpl::NullMapBuilder & builder) |
551 | { |
552 | const ColumnArray * col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
553 | |
554 | if (!col_array) |
555 | return false; |
556 | |
557 | const auto & col_nested = col_array->getData(); |
558 | auto col_res = col_nested.cloneEmpty(); |
559 | |
560 | if (index.getType() == Field::Types::UInt64) |
561 | ArrayElementGenericImpl::vectorConst<false>( |
562 | col_nested, col_array->getOffsets(), safeGet<UInt64>(index) - 1, *col_res, builder); |
563 | else if (index.getType() == Field::Types::Int64) |
564 | ArrayElementGenericImpl::vectorConst<true>( |
565 | col_nested, col_array->getOffsets(), -safeGet<Int64>(index) - 1, *col_res, builder); |
566 | else |
567 | throw Exception("Illegal type of array index" , ErrorCodes::LOGICAL_ERROR); |
568 | |
569 | block.getByPosition(result).column = std::move(col_res); |
570 | return true; |
571 | } |
572 | |
573 | template <typename IndexType> |
574 | bool FunctionArrayElement::executeGeneric(Block & block, const ColumnNumbers & arguments, size_t result, const PaddedPODArray<IndexType> & indices, |
575 | ArrayImpl::NullMapBuilder & builder) |
576 | { |
577 | const ColumnArray * col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
578 | |
579 | if (!col_array) |
580 | return false; |
581 | |
582 | const auto & col_nested = col_array->getData(); |
583 | auto col_res = col_nested.cloneEmpty(); |
584 | |
585 | ArrayElementGenericImpl::vector<IndexType>( |
586 | col_nested, col_array->getOffsets(), indices, *col_res, builder); |
587 | |
588 | block.getByPosition(result).column = std::move(col_res); |
589 | return true; |
590 | } |
591 | |
592 | template <typename IndexType> |
593 | bool FunctionArrayElement::executeConst(Block & block, const ColumnNumbers & arguments, size_t result, |
594 | const PaddedPODArray <IndexType> & indices, ArrayImpl::NullMapBuilder & builder, |
595 | size_t input_rows_count) |
596 | { |
597 | const ColumnArray * col_array = checkAndGetColumnConstData<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
598 | |
599 | if (!col_array) |
600 | return false; |
601 | |
602 | auto res = block.getByPosition(result).type->createColumn(); |
603 | |
604 | size_t rows = input_rows_count; |
605 | const IColumn & array_elements = col_array->getData(); |
606 | size_t array_size = array_elements.size(); |
607 | |
608 | for (size_t i = 0; i < rows; ++i) |
609 | { |
610 | IndexType index = indices[i]; |
611 | if (index > 0 && static_cast<size_t>(index) <= array_size) |
612 | { |
613 | size_t j = index - 1; |
614 | res->insertFrom(array_elements, j); |
615 | if (builder) |
616 | builder.update(j); |
617 | } |
618 | else if (index < 0 && static_cast<size_t>(-index) <= array_size) |
619 | { |
620 | size_t j = array_size + index; |
621 | res->insertFrom(array_elements, j); |
622 | if (builder) |
623 | builder.update(j); |
624 | } |
625 | else |
626 | { |
627 | res->insertDefault(); |
628 | if (builder) |
629 | builder.update(); |
630 | } |
631 | } |
632 | |
633 | block.getByPosition(result).column = std::move(res); |
634 | return true; |
635 | } |
636 | |
637 | template <typename IndexType> |
638 | bool FunctionArrayElement::executeArgument(Block & block, const ColumnNumbers & arguments, size_t result, |
639 | ArrayImpl::NullMapBuilder & builder, size_t input_rows_count) |
640 | { |
641 | auto index = checkAndGetColumn<ColumnVector<IndexType>>(block.getByPosition(arguments[1]).column.get()); |
642 | |
643 | if (!index) |
644 | return false; |
645 | |
646 | const auto & index_data = index->getData(); |
647 | |
648 | if (builder) |
649 | builder.initSink(index_data.size()); |
650 | |
651 | if (!(executeNumber<IndexType, UInt8>(block, arguments, result, index_data, builder) |
652 | || executeNumber<IndexType, UInt16>(block, arguments, result, index_data, builder) |
653 | || executeNumber<IndexType, UInt32>(block, arguments, result, index_data, builder) |
654 | || executeNumber<IndexType, UInt64>(block, arguments, result, index_data, builder) |
655 | || executeNumber<IndexType, Int8>(block, arguments, result, index_data, builder) |
656 | || executeNumber<IndexType, Int16>(block, arguments, result, index_data, builder) |
657 | || executeNumber<IndexType, Int32>(block, arguments, result, index_data, builder) |
658 | || executeNumber<IndexType, Int64>(block, arguments, result, index_data, builder) |
659 | || executeNumber<IndexType, Float32>(block, arguments, result, index_data, builder) |
660 | || executeNumber<IndexType, Float64>(block, arguments, result, index_data, builder) |
661 | || executeConst<IndexType>(block, arguments, result, index_data, builder, input_rows_count) |
662 | || executeString<IndexType>(block, arguments, result, index_data, builder) |
663 | || executeGeneric<IndexType>(block, arguments, result, index_data, builder))) |
664 | throw Exception("Illegal column " + block.getByPosition(arguments[0]).column->getName() |
665 | + " of first argument of function " + getName(), ErrorCodes::ILLEGAL_COLUMN); |
666 | |
667 | return true; |
668 | } |
669 | |
670 | bool FunctionArrayElement::executeTuple(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
671 | { |
672 | const ColumnArray * col_array = typeid_cast<const ColumnArray *>(block.getByPosition(arguments[0]).column.get()); |
673 | |
674 | if (!col_array) |
675 | return false; |
676 | |
677 | const ColumnTuple * col_nested = typeid_cast<const ColumnTuple *>(&col_array->getData()); |
678 | |
679 | if (!col_nested) |
680 | return false; |
681 | |
682 | const auto & tuple_columns = col_nested->getColumns(); |
683 | size_t tuple_size = tuple_columns.size(); |
684 | |
685 | const DataTypes & tuple_types = typeid_cast<const DataTypeTuple &>( |
686 | *typeid_cast<const DataTypeArray &>(*block.getByPosition(arguments[0]).type).getNestedType()).getElements(); |
687 | |
688 | /** We will calculate the function for the tuple of the internals of the array. |
689 | * To do this, create a temporary block. |
690 | * It will consist of the following columns |
691 | * - the index of the array to be taken; |
692 | * - an array of the first elements of the tuples; |
693 | * - the result of taking the elements by the index for an array of the first elements of the tuples; |
694 | * - array of the second elements of the tuples; |
695 | * - result of taking elements by index for an array of second elements of tuples; |
696 | * ... |
697 | */ |
698 | Block block_of_temporary_results; |
699 | block_of_temporary_results.insert(block.getByPosition(arguments[1])); |
700 | |
701 | /// results of taking elements by index for arrays from each element of the tuples; |
702 | Columns result_tuple_columns; |
703 | |
704 | for (size_t i = 0; i < tuple_size; ++i) |
705 | { |
706 | ColumnWithTypeAndName array_of_tuple_section; |
707 | array_of_tuple_section.column = ColumnArray::create(tuple_columns[i], col_array->getOffsetsPtr()); |
708 | array_of_tuple_section.type = std::make_shared<DataTypeArray>(tuple_types[i]); |
709 | block_of_temporary_results.insert(array_of_tuple_section); |
710 | |
711 | ColumnWithTypeAndName array_elements_of_tuple_section; |
712 | array_elements_of_tuple_section.type = getReturnTypeImpl( |
713 | {block_of_temporary_results.getByPosition(i * 2 + 1).type, block_of_temporary_results.getByPosition(0).type}); |
714 | block_of_temporary_results.insert(array_elements_of_tuple_section); |
715 | |
716 | executeImpl(block_of_temporary_results, ColumnNumbers{i * 2 + 1, 0}, i * 2 + 2, input_rows_count); |
717 | |
718 | result_tuple_columns.emplace_back(std::move(block_of_temporary_results.getByPosition(i * 2 + 2).column)); |
719 | } |
720 | |
721 | block.getByPosition(result).column = ColumnTuple::create(result_tuple_columns); |
722 | |
723 | return true; |
724 | } |
725 | |
726 | String FunctionArrayElement::getName() const |
727 | { |
728 | return name; |
729 | } |
730 | |
731 | DataTypePtr FunctionArrayElement::getReturnTypeImpl(const DataTypes & arguments) const |
732 | { |
733 | const DataTypeArray * array_type = checkAndGetDataType<DataTypeArray>(arguments[0].get()); |
734 | if (!array_type) |
735 | throw Exception("First argument for function " + getName() + " must be array." , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
736 | |
737 | if (!isInteger(arguments[1])) |
738 | throw Exception("Second argument for function " + getName() + " must be integer." , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
739 | |
740 | return array_type->getNestedType(); |
741 | } |
742 | |
743 | void FunctionArrayElement::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
744 | { |
745 | /// Check nullability. |
746 | bool is_array_of_nullable = false; |
747 | |
748 | const ColumnArray * col_array = nullptr; |
749 | const ColumnArray * col_const_array = nullptr; |
750 | |
751 | col_array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
752 | if (col_array) |
753 | is_array_of_nullable = isColumnNullable(col_array->getData()); |
754 | else |
755 | { |
756 | col_const_array = checkAndGetColumnConstData<ColumnArray>(block.getByPosition(arguments[0]).column.get()); |
757 | if (col_const_array) |
758 | is_array_of_nullable = isColumnNullable(col_const_array->getData()); |
759 | else |
760 | throw Exception("Illegal column " + block.getByPosition(arguments[0]).column->getName() |
761 | + " of first argument of function " + getName(), ErrorCodes::ILLEGAL_COLUMN); |
762 | } |
763 | |
764 | if (!is_array_of_nullable) |
765 | { |
766 | ArrayImpl::NullMapBuilder builder; |
767 | perform(block, arguments, result, builder, input_rows_count); |
768 | } |
769 | else |
770 | { |
771 | /// Perform initializations. |
772 | ArrayImpl::NullMapBuilder builder; |
773 | Block source_block; |
774 | |
775 | const auto & input_type = typeid_cast<const DataTypeNullable &>(*typeid_cast<const DataTypeArray &>(*block.getByPosition(arguments[0]).type).getNestedType()).getNestedType(); |
776 | const auto & tmp_ret_type = typeid_cast<const DataTypeNullable &>(*block.getByPosition(result).type).getNestedType(); |
777 | |
778 | if (col_array) |
779 | { |
780 | const auto & nullable_col = typeid_cast<const ColumnNullable &>(col_array->getData()); |
781 | const auto & nested_col = nullable_col.getNestedColumnPtr(); |
782 | |
783 | /// Put nested_col inside a ColumnArray. |
784 | source_block = |
785 | { |
786 | { |
787 | ColumnArray::create(nested_col, col_array->getOffsetsPtr()), |
788 | std::make_shared<DataTypeArray>(input_type), |
789 | "" |
790 | }, |
791 | block.getByPosition(arguments[1]), |
792 | { |
793 | nullptr, |
794 | tmp_ret_type, |
795 | "" |
796 | } |
797 | }; |
798 | |
799 | builder.initSource(nullable_col.getNullMapData().data()); |
800 | } |
801 | else |
802 | { |
803 | /// ColumnConst(ColumnArray(ColumnNullable(...))) |
804 | const auto & nullable_col = assert_cast<const ColumnNullable &>(col_const_array->getData()); |
805 | const auto & nested_col = nullable_col.getNestedColumnPtr(); |
806 | |
807 | source_block = |
808 | { |
809 | { |
810 | ColumnConst::create(ColumnArray::create(nested_col, col_const_array->getOffsetsPtr()), input_rows_count), |
811 | std::make_shared<DataTypeArray>(input_type), |
812 | "" |
813 | }, |
814 | block.getByPosition(arguments[1]), |
815 | { |
816 | nullptr, |
817 | tmp_ret_type, |
818 | "" |
819 | } |
820 | }; |
821 | |
822 | builder.initSource(nullable_col.getNullMapData().data()); |
823 | } |
824 | |
825 | perform(source_block, {0, 1}, 2, builder, input_rows_count); |
826 | |
827 | /// Store the result. |
828 | const ColumnWithTypeAndName & source_col = source_block.getByPosition(2); |
829 | ColumnWithTypeAndName & dest_col = block.getByPosition(result); |
830 | dest_col.column = ColumnNullable::create(source_col.column, builder ? std::move(builder).getNullMapColumnPtr() : ColumnUInt8::create()); |
831 | } |
832 | } |
833 | |
834 | void FunctionArrayElement::perform(Block & block, const ColumnNumbers & arguments, size_t result, |
835 | ArrayImpl::NullMapBuilder & builder, size_t input_rows_count) |
836 | { |
837 | if (executeTuple(block, arguments, result, input_rows_count)) |
838 | { |
839 | } |
840 | else if (!isColumnConst(*block.getByPosition(arguments[1]).column)) |
841 | { |
842 | if (!(executeArgument<UInt8>(block, arguments, result, builder, input_rows_count) |
843 | || executeArgument<UInt16>(block, arguments, result, builder, input_rows_count) |
844 | || executeArgument<UInt32>(block, arguments, result, builder, input_rows_count) |
845 | || executeArgument<UInt64>(block, arguments, result, builder, input_rows_count) |
846 | || executeArgument<Int8>(block, arguments, result, builder, input_rows_count) |
847 | || executeArgument<Int16>(block, arguments, result, builder, input_rows_count) |
848 | || executeArgument<Int32>(block, arguments, result, builder, input_rows_count) |
849 | || executeArgument<Int64>(block, arguments, result, builder, input_rows_count))) |
850 | throw Exception("Second argument for function " + getName() + " must must have UInt or Int type." , |
851 | ErrorCodes::ILLEGAL_COLUMN); |
852 | } |
853 | else |
854 | { |
855 | Field index = (*block.getByPosition(arguments[1]).column)[0]; |
856 | |
857 | if (builder) |
858 | builder.initSink(input_rows_count); |
859 | |
860 | if (index == 0u) |
861 | throw Exception("Array indices are 1-based" , ErrorCodes::ZERO_ARRAY_OR_TUPLE_INDEX); |
862 | |
863 | if (!(executeNumberConst<UInt8>(block, arguments, result, index, builder) |
864 | || executeNumberConst<UInt16>(block, arguments, result, index, builder) |
865 | || executeNumberConst<UInt32>(block, arguments, result, index, builder) |
866 | || executeNumberConst<UInt64>(block, arguments, result, index, builder) |
867 | || executeNumberConst<Int8>(block, arguments, result, index, builder) |
868 | || executeNumberConst<Int16>(block, arguments, result, index, builder) |
869 | || executeNumberConst<Int32>(block, arguments, result, index, builder) |
870 | || executeNumberConst<Int64>(block, arguments, result, index, builder) |
871 | || executeNumberConst<Float32>(block, arguments, result, index, builder) |
872 | || executeNumberConst<Float64>(block, arguments, result, index, builder) |
873 | || executeStringConst (block, arguments, result, index, builder) |
874 | || executeGenericConst (block, arguments, result, index, builder))) |
875 | throw Exception("Illegal column " + block.getByPosition(arguments[0]).column->getName() |
876 | + " of first argument of function " + getName(), |
877 | ErrorCodes::ILLEGAL_COLUMN); |
878 | } |
879 | } |
880 | |
881 | |
882 | void registerFunctionArrayElement(FunctionFactory & factory) |
883 | { |
884 | factory.registerFunction<FunctionArrayElement>(); |
885 | } |
886 | |
887 | } |
888 | |