-
Notifications
You must be signed in to change notification settings - Fork 9k
Expand file tree
/
Copy patharrayResize.cpp
More file actions
168 lines (137 loc) · 6.37 KB
/
Copy patharrayResize.cpp
File metadata and controls
168 lines (137 loc) · 6.37 KB
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
#include <Functions/array/arrayResize.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/FunctionFactory.h>
#include <Functions/GatherUtils/GatherUtils.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/DataTypeNullable.h>
#include <DataTypes/DataTypesNumber.h>
#include <DataTypes/getLeastSupertype.h>
#include <Columns/ColumnArray.h>
#include <Columns/ColumnConst.h>
#include <Interpreters/castColumn.h>
#include <IO/WriteHelpers.h>
#include <Common/typeid_cast.h>
namespace DB
{
namespace ErrorCodes
{
extern const int LOGICAL_ERROR;
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
}
namespace
{
ColumnPtr convertToStructure(const ColumnPtr & column, const IColumn & structure)
{
if (column->structureEquals(structure))
return column;
auto result = structure.cloneEmpty();
result->insertRangeFrom(*column, 0, column->size());
return result;
}
}
DataTypePtr FunctionArrayResize::getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const
{
FunctionArgumentDescriptors mandatory_args{
{"array", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isArray), nullptr, "Array"},
{"size", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isNumber), nullptr, "Number"}
};
FunctionArgumentDescriptors optional_args{
{"extender", nullptr, nullptr, "Any type"}
};
validateFunctionArguments(*this, arguments, mandatory_args, optional_args);
if (arguments[0].type->onlyNull())
return arguments[0].type;
/// Issue #48398
if (arguments[1].type->isNullable())
throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
"Second argument for function {} must not be Nullable.", getName());
if (arguments.size() == 2)
return arguments[0].type;
else
{
const auto * array_type = typeid_cast<const DataTypeArray *>(arguments[0].type.get());
auto data_types = {array_type->getNestedType(), arguments[2].type};
return std::make_shared<DataTypeArray>(getLeastSupertype(data_types));
}
}
ColumnPtr FunctionArrayResize::executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & return_type, size_t input_rows_count) const
{
if (return_type->onlyNull())
return return_type->createColumnConstWithDefaultValue(input_rows_count);
auto array_column = arguments[0].column;
auto size_column = arguments[1].column;
/// `getInt` on a Decimal column returns the raw unscaled value, so without this conversion
/// `arrayResize([1, 2, 3], 1.5::Decimal(2, 1))` would resize to 15 elements. Convert the size to a
/// plain integer to use its real value (rounded towards zero), consistent with Float and integer sizes.
if (isDecimal(arguments[1].type))
size_column = castColumn(arguments[1], std::make_shared<DataTypeInt64>());
if (!arguments[0].type->equals(*return_type))
array_column = castColumn(arguments[0], return_type);
const DataTypePtr & return_nested_type = typeid_cast<const DataTypeArray &>(*return_type).getNestedType();
size_t size = array_column->size();
std::unique_ptr<GatherUtils::IArraySource> array_source;
std::unique_ptr<GatherUtils::IValueSource> value_source;
bool is_const = false;
if (const auto * const_array_column = typeid_cast<const ColumnConst *>(array_column.get()))
{
is_const = true;
array_column = const_array_column->getDataColumnPtr();
}
if (const auto * argument_column_array = typeid_cast<const ColumnArray *>(array_column.get()))
array_source = GatherUtils::createArraySource(*argument_column_array, is_const, size);
else
throw Exception(ErrorCodes::LOGICAL_ERROR, "First arguments for function {} must be array.", getName());
auto result_column = array_column->cloneEmpty();
auto & result_array = typeid_cast<ColumnArray &>(*result_column);
ColumnPtr appended_column;
if (arguments.size() == 3)
{
appended_column = arguments[2].column;
if (!arguments[2].type->equals(*return_nested_type))
appended_column = castColumn(arguments[2], return_nested_type);
}
else
{
auto default_column = result_array.getData().cloneEmpty();
default_column->insertDefault();
appended_column = ColumnConst::create(std::move(default_column), size);
}
bool is_appended_const = false;
if (const auto * const_appended_column = typeid_cast<const ColumnConst *>(appended_column.get()))
{
is_appended_const = true;
appended_column = const_appended_column->getDataColumnPtr();
}
appended_column = convertToStructure(appended_column, result_array.getData());
value_source = GatherUtils::createValueSource(*appended_column, is_appended_const, size);
auto sink = GatherUtils::createArraySink(result_array, size);
if (isColumnConst(*size_column))
GatherUtils::resizeConstantSize(*array_source, *value_source, *sink, size_column->getInt(0));
else
GatherUtils::resizeDynamicSize(*array_source, *value_source, *sink, *size_column);
return result_column;
}
REGISTER_FUNCTION(ArrayResize)
{
FunctionDocumentation::Description description = "Changes the length of the array.";
FunctionDocumentation::Syntax syntax = "arrayResize(arr, size[, extender])";
FunctionDocumentation::Arguments arguments = {
{"arr", "Array to resize.", {"Array(T)"}},
{"size", R"(
-The new length of the array.
If `size` is less than the original size of the array, the array is truncated from the right.
If `size` is larger than the initial size of the array, the array is extended to the right with `extender` values or default values for the data type of the array items.
)"},
{"extender", "Value to use for extending the array. Can be `NULL`."}
};
FunctionDocumentation::ReturnedValue returned_value = {"An array of length `size`.", {"Array(T)"}};
FunctionDocumentation::Examples examples = {
{"Example 1", "SELECT arrayResize([1], 3);", "[1,0,0]"},
{"Example 2", "SELECT arrayResize([1], 3, NULL);", "[1,NULL,NULL]"},
};
FunctionDocumentation::IntroducedIn introduced_in = {1, 1};
FunctionDocumentation::Category category = FunctionDocumentation::Category::Array;
FunctionDocumentation documentation = {description, syntax, arguments, {}, returned_value, examples, introduced_in, category};
factory.registerFunction<FunctionArrayResize>(documentation);
}
}