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
17namespace DB
18{
19
20namespace 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
28namespace 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 */
36class FunctionArrayElement : public IFunction
37{
38public:
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
51private:
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
92namespace ArrayImpl
93{
94
95class NullMapBuilder
96{
97public:
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
127private:
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
136namespace
137{
138
139template <typename T>
140struct 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
225struct 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.
342struct 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
422FunctionPtr FunctionArrayElement::create(const Context &)
423{
424 return std::make_shared<FunctionArrayElement>();
425}
426
427
428template <typename DataType>
429bool 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
457template <typename IndexType, typename DataType>
458bool 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
480bool 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
520template <typename IndexType>
521bool 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
549bool 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
573template <typename IndexType>
574bool 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
592template <typename IndexType>
593bool 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
637template <typename IndexType>
638bool 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
670bool 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
726String FunctionArrayElement::getName() const
727{
728 return name;
729}
730
731DataTypePtr 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
743void 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
834void 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
882void registerFunctionArrayElement(FunctionFactory & factory)
883{
884 factory.registerFunction<FunctionArrayElement>();
885}
886
887}
888