1#include <DataTypes/DataTypesNumber.h>
2#include <Columns/ColumnsNumber.h>
3#include "FunctionArrayMapped.h"
4#include <Functions/FunctionFactory.h>
5
6
7namespace DB
8{
9
10template <bool reverse>
11struct 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
115struct NameArrayFill { static constexpr auto name = "arrayFill"; };
116struct NameArrayReverseFill { static constexpr auto name = "arrayReverseFill"; };
117using FunctionArrayFill = FunctionArrayMapped<ArrayFillImpl<false>, NameArrayFill>;
118using FunctionArrayReverseFill = FunctionArrayMapped<ArrayFillImpl<true>, NameArrayReverseFill>;
119
120void registerFunctionsArrayFill(FunctionFactory & factory)
121{
122 factory.registerFunction<FunctionArrayFill>();
123 factory.registerFunction<FunctionArrayReverseFill>();
124}
125
126}
127