1 | #include <DataTypes/DataTypesNumber.h> |
2 | #include <Columns/ColumnsNumber.h> |
3 | #include "FunctionArrayMapped.h" |
4 | #include <Functions/FunctionFactory.h> |
5 | |
6 | |
7 | namespace DB |
8 | { |
9 | |
10 | template <bool reverse> |
11 | struct ArrayFillImpl |
12 | { |
13 | static bool needBoolean() { return true; } |
14 | static bool needExpression() { return true; } |
15 | static bool needOneArray() { return false; } |
16 | |
17 | static DataTypePtr getReturnType(const DataTypePtr & /*expression_return*/, const DataTypePtr & array_element) |
18 | { |
19 | return std::make_shared<DataTypeArray>(array_element); |
20 | } |
21 | |
22 | static ColumnPtr execute(const ColumnArray & array, ColumnPtr mapped) |
23 | { |
24 | const ColumnUInt8 * column_fill = typeid_cast<const ColumnUInt8 *>(&*mapped); |
25 | |
26 | const IColumn & in_data = array.getData(); |
27 | const IColumn::Offsets & in_offsets = array.getOffsets(); |
28 | auto column_data = in_data.cloneEmpty(); |
29 | IColumn & out_data = *column_data.get(); |
30 | |
31 | if (column_fill) |
32 | { |
33 | const IColumn::Filter & fill = column_fill->getData(); |
34 | |
35 | size_t array_begin = 0; |
36 | size_t array_end = 0; |
37 | size_t begin = 0; |
38 | size_t end = 0; |
39 | |
40 | out_data.reserve(in_data.size()); |
41 | |
42 | for (size_t i = 0; i < in_offsets.size(); ++i) |
43 | { |
44 | array_end = in_offsets[i] - 1; |
45 | |
46 | for (; end <= array_end; ++end) |
47 | { |
48 | if (end == array_end || fill[end + 1] != fill[begin]) |
49 | { |
50 | if (fill[begin]) |
51 | out_data.insertRangeFrom(in_data, begin, end + 1 - begin); |
52 | else |
53 | { |
54 | if constexpr (reverse) |
55 | { |
56 | if (end == array_end) |
57 | out_data.insertManyFrom(in_data, array_end, end + 1 - begin); |
58 | else |
59 | out_data.insertManyFrom(in_data, end + 1, end + 1 - begin); |
60 | } |
61 | else |
62 | { |
63 | if (begin == array_begin) |
64 | out_data.insertManyFrom(in_data, array_begin, end + 1 - begin); |
65 | else |
66 | out_data.insertManyFrom(in_data, begin - 1, end + 1 - begin); |
67 | } |
68 | } |
69 | |
70 | begin = end + 1; |
71 | } |
72 | } |
73 | |
74 | array_begin = array_end + 1; |
75 | } |
76 | } |
77 | else |
78 | { |
79 | auto column_fill_const = checkAndGetColumnConst<ColumnUInt8>(&*mapped); |
80 | |
81 | if (!column_fill_const) |
82 | throw Exception("Unexpected type of cut column" , ErrorCodes::ILLEGAL_COLUMN); |
83 | |
84 | if (column_fill_const->getValue<UInt8>()) |
85 | return ColumnArray::create( |
86 | array.getDataPtr(), |
87 | array.getOffsetsPtr() |
88 | ); |
89 | |
90 | size_t array_begin = 0; |
91 | size_t array_end = 0; |
92 | |
93 | out_data.reserve(in_data.size()); |
94 | |
95 | for (size_t i = 0; i < in_offsets.size(); ++i) |
96 | { |
97 | array_end = in_offsets[i] - 1; |
98 | |
99 | if constexpr (reverse) |
100 | out_data.insertManyFrom(in_data, array_end, array_end + 1 - array_begin); |
101 | else |
102 | out_data.insertManyFrom(in_data, array_begin, array_end + 1 - array_begin); |
103 | |
104 | array_begin = array_end + 1; |
105 | } |
106 | } |
107 | |
108 | return ColumnArray::create( |
109 | std::move(column_data), |
110 | array.getOffsetsPtr() |
111 | ); |
112 | } |
113 | }; |
114 | |
115 | struct NameArrayFill { static constexpr auto name = "arrayFill" ; }; |
116 | struct NameArrayReverseFill { static constexpr auto name = "arrayReverseFill" ; }; |
117 | using FunctionArrayFill = FunctionArrayMapped<ArrayFillImpl<false>, NameArrayFill>; |
118 | using FunctionArrayReverseFill = FunctionArrayMapped<ArrayFillImpl<true>, NameArrayReverseFill>; |
119 | |
120 | void registerFunctionsArrayFill(FunctionFactory & factory) |
121 | { |
122 | factory.registerFunction<FunctionArrayFill>(); |
123 | factory.registerFunction<FunctionArrayReverseFill>(); |
124 | } |
125 | |
126 | } |
127 | |