ClickHouse

1
#include <Functions/replicate.h>
2
#include <Functions/IFunction.h>
3
#include <Functions/FunctionFactory.h>
4
#include <Functions/FunctionHelpers.h>
5
#include <DataTypes/DataTypeArray.h>
6
#include <Columns/ColumnArray.h>
7
#include <Interpreters/Context.h>
8

9

10
namespace DB
11
{
12
namespace ErrorCodes
13
{
14
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
15
    extern const int ILLEGAL_COLUMN;
16
    extern const int TOO_FEW_ARGUMENTS_FOR_FUNCTION;
17
}
18

19
DataTypePtr FunctionReplicate::getReturnTypeImpl(const DataTypes & arguments) const
20
{
21
    if (arguments.size() < 2)
22
        throw Exception(ErrorCodes::TOO_FEW_ARGUMENTS_FOR_FUNCTION,
23
                        "Function {} expect at least two arguments, got {}", getName(), arguments.size());
24

25
    for (size_t i = 1; i < arguments.size(); ++i)
26
    {
27
        const DataTypeArray * array_type = checkAndGetDataType<DataTypeArray>(arguments[i].get());
28
        if (!array_type)
29
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
30
                            "Argument {} for function {} must be array.",
31
                            i + 1, getName());
32
    }
33
    return std::make_shared<DataTypeArray>(arguments[0]);
34
}
35

36
ColumnPtr FunctionReplicate::executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const
37
{
38
    ColumnPtr first_column = arguments[0].column;
39
    ColumnPtr offsets;
40

41
    for (size_t i = 1; i < arguments.size(); ++i)
42
    {
43
        const ColumnArray * array_column = checkAndGetColumn<ColumnArray>(arguments[i].column.get());
44
        ColumnPtr temp_column;
45
        if (!array_column)
46
        {
47
            const auto * const_array_column = checkAndGetColumnConst<ColumnArray>(arguments[i].column.get());
48
            if (!const_array_column)
49
                throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Unexpected column for replicate");
50
            temp_column = const_array_column->convertToFullColumn();
51
            array_column = checkAndGetColumn<ColumnArray>(temp_column.get());
52
        }
53

54
        if (!offsets || offsets->empty())
55
            offsets = array_column->getOffsetsPtr();
56
    }
57

58
    const auto & offsets_data = assert_cast<const ColumnArray::ColumnOffsets &>(*offsets).getData(); /// NOLINT
59
    return ColumnArray::create(first_column->replicate(offsets_data)->convertToFullColumnIfConst(), offsets);
60
}
61

62
REGISTER_FUNCTION(Replicate)
63
{
64
    factory.registerFunction<FunctionReplicate>();
65
}
66

67
}
68