blob: 045558609a458a6323267ea89dab2b71b2485a1a (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
|
#include <DataTypes/DataTypesNumber.h>
#include <Columns/ColumnsNumber.h>
#include <Functions/FunctionFactory.h>
#include "FunctionArrayMapped.h"
namespace DB
{
namespace ErrorCodes
{
extern const int ILLEGAL_COLUMN;
}
enum class ArrayFirstLastIndexStrategy
{
First,
Last
};
template <ArrayFirstLastIndexStrategy strategy>
struct ArrayFirstLastIndexImpl
{
static bool needBoolean() { return false; }
static bool needExpression() { return true; }
static bool needOneArray() { return false; }
static DataTypePtr getReturnType(const DataTypePtr & /*expression_return*/, const DataTypePtr & /*array_element*/)
{
return std::make_shared<DataTypeUInt32>();
}
static ColumnPtr execute(const ColumnArray & array, ColumnPtr mapped)
{
const auto * column_filter = typeid_cast<const ColumnUInt8 *>(&*mapped);
if (!column_filter)
{
const auto * column_filter_const = checkAndGetColumnConst<ColumnUInt8>(&*mapped);
if (!column_filter_const)
throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Unexpected type of filter column");
if (column_filter_const->getValue<UInt8>())
{
const auto & offsets = array.getOffsets();
auto out_column = ColumnUInt32::create(offsets.size());
auto & out_index = out_column->getData();
size_t offsets_size = offsets.size();
for (size_t offset_index = 0; offset_index < offsets_size; ++offset_index)
{
size_t start_offset = offsets[offset_index - 1];
size_t end_offset = offsets[offset_index];
if (end_offset > start_offset)
{
if constexpr (strategy == ArrayFirstLastIndexStrategy::First)
out_index[offset_index] = 1;
else
out_index[offset_index] = static_cast<UInt32>(end_offset - start_offset);
}
else
{
out_index[offset_index] = 0;
}
}
return out_column;
}
else
{
return DataTypeUInt32().createColumnConst(array.size(), 0u);
}
}
const auto & filter = column_filter->getData();
const auto & offsets = array.getOffsets();
size_t offsets_size = offsets.size();
auto out_column = ColumnUInt32::create(offsets_size);
auto & out_index = out_column->getData();
for (size_t offset_index = 0; offset_index < offsets_size; ++offset_index)
{
size_t start_offset = offsets[offset_index - 1];
size_t end_offset = offsets[offset_index];
size_t result_index = 0;
if constexpr (strategy == ArrayFirstLastIndexStrategy::First)
{
for (size_t index = 1; start_offset != end_offset; ++start_offset, ++index)
{
if (filter[start_offset])
{
result_index = index;
break;
}
}
}
else
{
for (size_t index = end_offset - start_offset; end_offset != start_offset; --end_offset, --index)
{
if (filter[end_offset - 1])
{
result_index = index;
break;
}
}
}
out_index[offset_index] = static_cast<UInt32>(result_index);
}
return out_column;
}
};
struct NameArrayFirstIndex { static constexpr auto name = "arrayFirstIndex"; };
using ArrayFirstIndexImpl = ArrayFirstLastIndexImpl<ArrayFirstLastIndexStrategy::First>;
using FunctionArrayFirstIndex = FunctionArrayMapped<ArrayFirstIndexImpl, NameArrayFirstIndex>;
struct NameArrayLastIndex { static constexpr auto name = "arrayLastIndex"; };
using ArrayLastIndexImpl = ArrayFirstLastIndexImpl<ArrayFirstLastIndexStrategy::Last>;
using FunctionArrayLastIndex = FunctionArrayMapped<ArrayLastIndexImpl, NameArrayLastIndex>;
REGISTER_FUNCTION(ArrayFirstIndex)
{
factory.registerFunction<FunctionArrayFirstIndex>();
factory.registerFunction<FunctionArrayLastIndex>();
}
}
|
