ClickHouse

1
#include <Functions/IFunction.h>
2
#include <Functions/FunctionFactory.h>
3
#include <DataTypes/DataTypeNullable.h>
4
#include <Core/ColumnNumbers.h>
5
#include <Columns/ColumnNullable.h>
6

7

8
namespace DB
9
{
10

11
namespace ErrorCodes
12
{
13
    extern const int ILLEGAL_COLUMN;
14
}
15

16
namespace
17
{
18

19
/// Implements the function assumeNotNull which takes 1 argument and works as follows:
20
/// - if the argument is a nullable column, return its embedded column;
21
/// - otherwise return the original argument.
22
/// NOTE: assumeNotNull may not be called with the NULL value.
23
class FunctionAssumeNotNull : public IFunction
24
{
25
public:
26
    static constexpr auto name = "assumeNotNull";
27

28
    static FunctionPtr create(ContextPtr)
29
    {
30
        return std::make_shared<FunctionAssumeNotNull>();
31
    }
32

33
    std::string getName() const override
34
    {
35
        return name;
36
    }
37

38
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
39

40
    size_t getNumberOfArguments() const override { return 1; }
41
    bool useDefaultImplementationForNulls() const override { return false; }
42
    bool useDefaultImplementationForConstants() const override { return true; }
43
    ColumnNumbers getArgumentsThatDontImplyNullableReturnType(size_t /*number_of_arguments*/) const override { return {0}; }
44

45
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
46
    {
47
        return removeNullable(arguments[0]);
48
    }
49

50
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const override
51
    {
52
        const ColumnPtr & col = arguments[0].column;
53

54
        if (arguments[0].type->onlyNull() && !col->empty())
55
            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Cannot create non-empty column with type Nothing");
56

57
        if (const auto * nullable_col = checkAndGetColumn<ColumnNullable>(*col))
58
            return nullable_col->getNestedColumnPtr();
59
        else
60
            return col;
61
    }
62
};
63

64
}
65

66
REGISTER_FUNCTION(AssumeNotNull)
67
{
68
    factory.registerFunction<FunctionAssumeNotNull>();
69
}
70

71
}
72