ClickHouse
133 строки · 4.7 Кб
1#include <Functions/IFunction.h>
2#include <Functions/castTypeToEither.h>
3#include <Functions/FunctionFactory.h>
4#include <DataTypes/DataTypesNumber.h>
5#include <DataTypes/DataTypeNullable.h>
6#include <Core/ColumnNumbers.h>
7#include <Columns/ColumnNullable.h>
8
9
10namespace DB
11{
12namespace ErrorCodes
13{
14extern const int BAD_ARGUMENTS;
15extern const int ILLEGAL_COLUMN;
16}
17
18namespace
19{
20
21/// Returns 1 if argument is zero or NULL.
22/// It can be used to negate filter in WHERE condition.
23/// "WHERE isZeroOrNull(expr)" will return exactly the same rows that "WHERE expr" will filter out.
24class FunctionIsZeroOrNull : public IFunction
25{
26public:
27static constexpr auto name = "isZeroOrNull";
28
29static FunctionPtr create(ContextPtr)
30{
31return std::make_shared<FunctionIsZeroOrNull>();
32}
33
34std::string getName() const override
35{
36return name;
37}
38
39size_t getNumberOfArguments() const override { return 1; }
40bool useDefaultImplementationForNulls() const override { return false; }
41bool useDefaultImplementationForConstants() const override { return true; }
42bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
43ColumnNumbers getArgumentsThatDontImplyNullableReturnType(size_t /*number_of_arguments*/) const override { return {0}; }
44
45DataTypePtr getReturnTypeImpl(const DataTypes & types) const override
46{
47if (!isNumber(removeNullable(types.at(0))) && !isNothing(removeNullable(types.at(0))))
48throw Exception(
49ErrorCodes::BAD_ARGUMENTS, "The argument of function {} must have simple numeric type, possibly Nullable or Null", name);
50
51return std::make_shared<DataTypeUInt8>();
52}
53
54ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
55{
56if (isNothing(removeNullable(arguments[0].type)))
57return DataTypeUInt8{}.createColumnConst(input_rows_count, 1);
58
59const ColumnPtr & input_column = arguments[0].column;
60
61ColumnPtr res;
62if (const ColumnNullable * input_column_nullable = checkAndGetColumn<ColumnNullable>(input_column.get()))
63{
64const NullMap & null_map = input_column_nullable->getNullMapData();
65const IColumn * nested_column = &input_column_nullable->getNestedColumn();
66
67if (!castTypeToEither<
68ColumnUInt8, ColumnUInt16, ColumnUInt32, ColumnUInt64,
69ColumnInt8, ColumnInt16, ColumnInt32, ColumnInt64,
70ColumnFloat32, ColumnFloat64>(
71nested_column, [&](const auto & column)
72{
73auto col = ColumnUInt8::create(input_rows_count);
74processNullable(column.getData(), null_map, col->getData(), input_rows_count);
75res = std::move(col);
76return true;
77}))
78{
79throw Exception(
80ErrorCodes::ILLEGAL_COLUMN,
81"The argument of function {} must have simple numeric type, possibly Nullable or Null",
82name);
83}
84}
85else
86{
87if (!castTypeToEither<
88ColumnUInt8, ColumnUInt16, ColumnUInt32, ColumnUInt64,
89ColumnInt8, ColumnInt16, ColumnInt32, ColumnInt64,
90ColumnFloat32, ColumnFloat64>(
91input_column.get(), [&](const auto & column)
92{
93auto col = ColumnUInt8::create(input_rows_count);
94processNotNullable(column.getData(), col->getData(), input_rows_count);
95res = std::move(col);
96return true;
97}))
98{
99throw Exception(
100ErrorCodes::ILLEGAL_COLUMN,
101"The argument of function {} must have simple numeric type, possibly Nullable or Null",
102name);
103}
104}
105
106return res;
107}
108
109private:
110template <typename InputData>
111void processNotNullable(const InputData & input_data, ColumnUInt8::Container & result_data, size_t input_rows_count) const
112{
113for (size_t i = 0; i < input_rows_count; ++i)
114result_data[i] = input_data[i] == 0;
115}
116
117template <typename InputData>
118void processNullable(const InputData & input_data, const NullMap & input_null_map,
119ColumnUInt8::Container & result_data, size_t input_rows_count) const
120{
121for (size_t i = 0; i < input_rows_count; ++i)
122result_data[i] = input_null_map[i] || input_data[i] == 0;
123}
124};
125
126}
127
128REGISTER_FUNCTION(IsZeroOrNull)
129{
130factory.registerFunction<FunctionIsZeroOrNull>();
131}
132
133}
134