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 | /** arrayFilter(x -> predicate, array) - leave in the array only the elements for which the expression is true. |
11 | */ |
12 | struct ArrayFilterImpl |
13 | { |
14 | static bool needBoolean() { return true; } |
15 | static bool needExpression() { return true; } |
16 | static bool needOneArray() { return false; } |
17 | |
18 | static DataTypePtr getReturnType(const DataTypePtr & /*expression_return*/, const DataTypePtr & array_element) |
19 | { |
20 | return std::make_shared<DataTypeArray>(array_element); |
21 | } |
22 | |
23 | /// If there are several arrays, the first one is passed here. |
24 | static ColumnPtr execute(const ColumnArray & array, ColumnPtr mapped) |
25 | { |
26 | const ColumnUInt8 * column_filter = typeid_cast<const ColumnUInt8 *>(&*mapped); |
27 | |
28 | if (!column_filter) |
29 | { |
30 | auto column_filter_const = checkAndGetColumnConst<ColumnUInt8>(&*mapped); |
31 | |
32 | if (!column_filter_const) |
33 | throw Exception("Unexpected type of filter column" , ErrorCodes::ILLEGAL_COLUMN); |
34 | |
35 | if (column_filter_const->getValue<UInt8>()) |
36 | return array.clone(); |
37 | else |
38 | return ColumnArray::create( |
39 | array.getDataPtr()->cloneEmpty(), |
40 | ColumnArray::ColumnOffsets::create(array.size(), 0)); |
41 | } |
42 | |
43 | const IColumn::Filter & filter = column_filter->getData(); |
44 | ColumnPtr filtered = array.getData().filter(filter, -1); |
45 | |
46 | const IColumn::Offsets & in_offsets = array.getOffsets(); |
47 | auto column_offsets = ColumnArray::ColumnOffsets::create(in_offsets.size()); |
48 | IColumn::Offsets & out_offsets = column_offsets->getData(); |
49 | |
50 | size_t in_pos = 0; |
51 | size_t out_pos = 0; |
52 | for (size_t i = 0; i < in_offsets.size(); ++i) |
53 | { |
54 | for (; in_pos < in_offsets[i]; ++in_pos) |
55 | { |
56 | if (filter[in_pos]) |
57 | ++out_pos; |
58 | } |
59 | out_offsets[i] = out_pos; |
60 | } |
61 | |
62 | return ColumnArray::create(filtered, std::move(column_offsets)); |
63 | } |
64 | }; |
65 | |
66 | struct NameArrayFilter { static constexpr auto name = "arrayFilter" ; }; |
67 | using FunctionArrayFilter = FunctionArrayMapped<ArrayFilterImpl, NameArrayFilter>; |
68 | |
69 | void registerFunctionArrayFilter(FunctionFactory & factory) |
70 | { |
71 | factory.registerFunction<FunctionArrayFilter>(); |
72 | } |
73 | |
74 | } |
75 | |
76 | |
77 | |