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 | struct ArrayFirstImpl |
11 | { |
12 | static bool needBoolean() { return false; } |
13 | static bool needExpression() { return true; } |
14 | static bool needOneArray() { return false; } |
15 | |
16 | static DataTypePtr getReturnType(const DataTypePtr & /*expression_return*/, const DataTypePtr & array_element) |
17 | { |
18 | return array_element; |
19 | } |
20 | |
21 | static ColumnPtr execute(const ColumnArray & array, ColumnPtr mapped) |
22 | { |
23 | auto column_filter = typeid_cast<const ColumnUInt8 *>(&*mapped); |
24 | |
25 | if (!column_filter) |
26 | { |
27 | auto column_filter_const = checkAndGetColumnConst<ColumnUInt8>(&*mapped); |
28 | |
29 | if (!column_filter_const) |
30 | throw Exception("Unexpected type of filter column", ErrorCodes::ILLEGAL_COLUMN); |
31 | |
32 | if (column_filter_const->getValue<UInt8>()) |
33 | { |
34 | const auto & offsets = array.getOffsets(); |
35 | const auto & data = array.getData(); |
36 | auto out = data.cloneEmpty(); |
37 | out->reserve(data.size()); |
38 | |
39 | size_t pos{}; |
40 | for (size_t i = 0; i < offsets.size(); ++i) |
41 | { |
42 | if (offsets[i] - pos > 0) |
43 | out->insert(data[pos]); |
44 | else |
45 | out->insertDefault(); |
46 | |
47 | pos = offsets[i]; |
48 | } |
49 | |
50 | return out; |
51 | } |
52 | else |
53 | { |
54 | auto out = array.getData().cloneEmpty(); |
55 | out->insertDefault(); |
56 | return out->replicate(IColumn::Offsets(1, array.size())); |
57 | } |
58 | } |
59 | |
60 | const auto & filter = column_filter->getData(); |
61 | const auto & offsets = array.getOffsets(); |
62 | const auto & data = array.getData(); |
63 | auto out = data.cloneEmpty(); |
64 | out->reserve(data.size()); |
65 | |
66 | size_t pos{}; |
67 | for (size_t i = 0; i < offsets.size(); ++i) |
68 | { |
69 | auto exists = false; |
70 | for (; pos < offsets[i]; ++pos) |
71 | { |
72 | if (filter[pos]) |
73 | { |
74 | out->insert(data[pos]); |
75 | exists = true; |
76 | pos = offsets[i]; |
77 | break; |
78 | } |
79 | } |
80 | |
81 | if (!exists) |
82 | out->insertDefault(); |
83 | } |
84 | |
85 | return out; |
86 | } |
87 | }; |
88 | |
89 | struct NameArrayFirst { static constexpr auto name = "arrayFirst"; }; |
90 | using FunctionArrayFirst = FunctionArrayMapped<ArrayFirstImpl, NameArrayFirst>; |
91 | |
92 | void registerFunctionArrayFirst(FunctionFactory & factory) |
93 | { |
94 | factory.registerFunction<FunctionArrayFirst>(); |
95 | } |
96 | |
97 | } |
98 | |
99 |