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 ArrayFirstIndexImpl |
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 std::make_shared<DataTypeUInt32>(); |
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 | auto out_column = ColumnUInt32::create(offsets.size()); |
36 | auto & out_index = out_column->getData(); |
37 | |
38 | size_t pos{}; |
39 | for (size_t i = 0; i < offsets.size(); ++i) |
40 | { |
41 | out_index[i] = offsets[i] - pos > 0; |
42 | pos = offsets[i]; |
43 | } |
44 | |
45 | return out_column; |
46 | } |
47 | else |
48 | return DataTypeUInt32().createColumnConst(array.size(), 0u); |
49 | } |
50 | |
51 | const auto & filter = column_filter->getData(); |
52 | const auto & offsets = array.getOffsets(); |
53 | auto out_column = ColumnUInt32::create(offsets.size()); |
54 | auto & out_index = out_column->getData(); |
55 | |
56 | size_t pos{}; |
57 | for (size_t i = 0; i < offsets.size(); ++i) |
58 | { |
59 | UInt32 index{}; |
60 | for (size_t idx{1}; pos < offsets[i]; ++pos, ++idx) |
61 | { |
62 | if (filter[pos]) |
63 | { |
64 | index = idx; |
65 | pos = offsets[i]; |
66 | break; |
67 | } |
68 | } |
69 | |
70 | out_index[i] = index; |
71 | } |
72 | |
73 | return out_column; |
74 | } |
75 | }; |
76 | |
77 | struct NameArrayFirstIndex { static constexpr auto name = "arrayFirstIndex"; }; |
78 | using FunctionArrayFirstIndex = FunctionArrayMapped<ArrayFirstIndexImpl, NameArrayFirstIndex>; |
79 | |
80 | void registerFunctionArrayFirstIndex(FunctionFactory & factory) |
81 | { |
82 | factory.registerFunction<FunctionArrayFirstIndex>(); |
83 | } |
84 | |
85 | } |
86 | |
87 |