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
10struct 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
89struct NameArrayFirst { static constexpr auto name = "arrayFirst"; };
90using FunctionArrayFirst = FunctionArrayMapped<ArrayFirstImpl, NameArrayFirst>;
91
92void registerFunctionArrayFirst(FunctionFactory & factory)
93{
94 factory.registerFunction<FunctionArrayFirst>();
95}
96
97}
98
99