1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/GatherUtils/GatherUtils.h> |
3 | #include <DataTypes/DataTypeArray.h> |
4 | #include <Columns/ColumnArray.h> |
5 | #include <Common/typeid_cast.h> |
6 | |
7 | |
8 | namespace DB |
9 | { |
10 | |
11 | namespace ErrorCodes |
12 | { |
13 | extern const int LOGICAL_ERROR; |
14 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
15 | } |
16 | |
17 | |
18 | class FunctionArrayPop : public IFunction |
19 | { |
20 | public: |
21 | FunctionArrayPop(bool pop_front_, const char * name_) : pop_front(pop_front_), name(name_) {} |
22 | |
23 | String getName() const override { return name; } |
24 | |
25 | bool isVariadic() const override { return false; } |
26 | size_t getNumberOfArguments() const override { return 1; } |
27 | |
28 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
29 | { |
30 | if (arguments[0]->onlyNull()) |
31 | return arguments[0]; |
32 | |
33 | auto array_type = typeid_cast<const DataTypeArray *>(arguments[0].get()); |
34 | if (!array_type) |
35 | throw Exception("First argument for function " + getName() + " must be an array but it has type " |
36 | + arguments[0]->getName() + "." , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
37 | |
38 | return arguments[0]; |
39 | } |
40 | |
41 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
42 | { |
43 | const auto & return_type = block.getByPosition(result).type; |
44 | |
45 | if (return_type->onlyNull()) |
46 | { |
47 | block.getByPosition(result).column = return_type->createColumnConstWithDefaultValue(input_rows_count); |
48 | return; |
49 | } |
50 | |
51 | auto result_column = return_type->createColumn(); |
52 | |
53 | const auto & array_column = block.getByPosition(arguments[0]).column; |
54 | |
55 | std::unique_ptr<GatherUtils::IArraySource> source; |
56 | |
57 | size_t size = array_column->size(); |
58 | |
59 | if (auto argument_column_array = typeid_cast<const ColumnArray *>(array_column.get())) |
60 | source = GatherUtils::createArraySource(*argument_column_array, false, size); |
61 | else |
62 | throw Exception{"First arguments for function " + getName() + " must be array." , ErrorCodes::LOGICAL_ERROR}; |
63 | |
64 | auto sink = GatherUtils::createArraySink(typeid_cast<ColumnArray &>(*result_column), size); |
65 | |
66 | if (pop_front) |
67 | GatherUtils::sliceFromLeftConstantOffsetUnbounded(*source, *sink, 1); |
68 | else |
69 | GatherUtils::sliceFromLeftConstantOffsetBounded(*source, *sink, 0, -1); |
70 | |
71 | block.getByPosition(result).column = std::move(result_column); |
72 | } |
73 | |
74 | bool useDefaultImplementationForConstants() const override { return true; } |
75 | bool useDefaultImplementationForNulls() const override { return false; } |
76 | |
77 | private: |
78 | bool pop_front; |
79 | const char * name; |
80 | }; |
81 | |
82 | } |
83 | |