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
8namespace DB
9{
10
11namespace ErrorCodes
12{
13 extern const int LOGICAL_ERROR;
14 extern const int ILLEGAL_TYPE_OF_ARGUMENT;
15}
16
17
18class FunctionArrayPop : public IFunction
19{
20public:
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
77private:
78 bool pop_front;
79 const char * name;
80};
81
82}
83