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
10/** arrayFilter(x -> predicate, array) - leave in the array only the elements for which the expression is true.
11 */
12struct 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
66struct NameArrayFilter { static constexpr auto name = "arrayFilter"; };
67using FunctionArrayFilter = FunctionArrayMapped<ArrayFilterImpl, NameArrayFilter>;
68
69void registerFunctionArrayFilter(FunctionFactory & factory)
70{
71 factory.registerFunction<FunctionArrayFilter>();
72}
73
74}
75
76
77