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 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
77struct NameArrayFirstIndex { static constexpr auto name = "arrayFirstIndex"; };
78using FunctionArrayFirstIndex = FunctionArrayMapped<ArrayFirstIndexImpl, NameArrayFirstIndex>;
79
80void registerFunctionArrayFirstIndex(FunctionFactory & factory)
81{
82 factory.registerFunction<FunctionArrayFirstIndex>();
83}
84
85}
86
87