ClickHouse
183 строки · 7.6 Кб
1#include <Functions/IFunction.h>
2#include <Functions/FunctionFactory.h>
3#include <Functions/FunctionHelpers.h>
4#include <DataTypes/DataTypeTuple.h>
5#include <DataTypes/DataTypesNumber.h>
6#include <Columns/ColumnConst.h>
7#include <Columns/ColumnsNumber.h>
8#include <Columns/ColumnTuple.h>
9#include <Columns/ColumnSet.h>
10#include <Interpreters/Set.h>
11
12
13namespace DB
14{
15namespace ErrorCodes
16{
17extern const int ILLEGAL_COLUMN;
18extern const int LOGICAL_ERROR;
19}
20
21namespace
22{
23
24/** in(x, set) - function for evaluating the IN
25* notIn(x, set) - and NOT IN.
26*/
27
28template <bool negative, bool global, bool null_is_skipped, bool ignore_set>
29struct FunctionInName;
30
31template <> struct FunctionInName<false, false, true, false> { static constexpr auto name = "in"; };
32template <> struct FunctionInName<false, true, true, false> { static constexpr auto name = "globalIn"; };
33template <> struct FunctionInName<true, false, true, false> { static constexpr auto name = "notIn"; };
34template <> struct FunctionInName<true, true, true, false> { static constexpr auto name = "globalNotIn"; };
35template <> struct FunctionInName<false, false, false, false> { static constexpr auto name = "nullIn"; };
36template <> struct FunctionInName<false, true, false, false> { static constexpr auto name = "globalNullIn"; };
37template <> struct FunctionInName<true, false, false, false> { static constexpr auto name = "notNullIn"; };
38template <> struct FunctionInName<true, true, false, false> { static constexpr auto name = "globalNotNullIn"; };
39template <> struct FunctionInName<false, false, true, true> { static constexpr auto name = "inIgnoreSet"; };
40template <> struct FunctionInName<false, true, true, true> { static constexpr auto name = "globalInIgnoreSet"; };
41template <> struct FunctionInName<true, false, true, true> { static constexpr auto name = "notInIgnoreSet"; };
42template <> struct FunctionInName<true, true, true, true> { static constexpr auto name = "globalNotInIgnoreSet"; };
43template <> struct FunctionInName<false, false, false, true> { static constexpr auto name = "nullInIgnoreSet"; };
44template <> struct FunctionInName<false, true, false, true> { static constexpr auto name = "globalNullInIgnoreSet"; };
45template <> struct FunctionInName<true, false, false, true> { static constexpr auto name = "notNullInIgnoreSet"; };
46template <> struct FunctionInName<true, true, false, true> { static constexpr auto name = "globalNotNullInIgnoreSet"; };
47
48template <bool negative, bool global, bool null_is_skipped, bool ignore_set>
49class FunctionIn : public IFunction
50{
51public:
52/// ignore_set flag means that we don't use set from the second argument, just return zero column.
53/// It is needed to perform type analysis without creation of set.
54static constexpr auto name = FunctionInName<negative, global, null_is_skipped, ignore_set>::name;
55
56static FunctionPtr create(ContextPtr)
57{
58return std::make_shared<FunctionIn>();
59}
60
61String getName() const override
62{
63return name;
64}
65
66size_t getNumberOfArguments() const override
67{
68return 2;
69}
70
71DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override
72{
73return std::make_shared<DataTypeUInt8>();
74}
75
76bool useDefaultImplementationForConstants() const override
77{
78/// Never return constant for -IgnoreSet functions to avoid constant folding.
79return !ignore_set;
80}
81
82bool useDefaultImplementationForNulls() const override { return null_is_skipped; }
83
84bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
85
86ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, [[maybe_unused]] size_t input_rows_count) const override
87{
88if constexpr (ignore_set)
89return ColumnUInt8::create(input_rows_count, 0u);
90if (input_rows_count == 0)
91return ColumnUInt8::create();
92
93/// Second argument must be ColumnSet.
94ColumnPtr column_set_ptr = arguments[1].column;
95const ColumnSet * column_set = checkAndGetColumnConstData<const ColumnSet>(column_set_ptr.get());
96if (!column_set)
97column_set = checkAndGetColumn<const ColumnSet>(column_set_ptr.get());
98if (!column_set)
99throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Second argument for function '{}' must be Set; found {}",
100getName(), column_set_ptr->getName());
101
102ColumnsWithTypeAndName columns_of_key_columns;
103
104/// First argument may be a tuple or a single column.
105const ColumnWithTypeAndName & left_arg = arguments[0];
106const ColumnTuple * tuple = typeid_cast<const ColumnTuple *>(left_arg.column.get());
107const ColumnConst * const_tuple = checkAndGetColumnConst<ColumnTuple>(left_arg.column.get());
108const DataTypeTuple * type_tuple = typeid_cast<const DataTypeTuple *>(left_arg.type.get());
109
110ColumnPtr materialized_tuple;
111if (const_tuple)
112{
113materialized_tuple = const_tuple->convertToFullColumn();
114tuple = typeid_cast<const ColumnTuple *>(materialized_tuple.get());
115}
116
117auto future_set = column_set->getData();
118if (!future_set)
119throw Exception(ErrorCodes::LOGICAL_ERROR, "No Set is passed as the second argument for function '{}'", getName());
120
121auto set = future_set->get();
122if (!set)
123throw Exception(ErrorCodes::LOGICAL_ERROR, "Not-ready Set is passed as the second argument for function '{}'", getName());
124
125auto set_types = set->getDataTypes();
126
127if (tuple && set_types.size() != 1 && set_types.size() == tuple->tupleSize())
128{
129const auto & tuple_columns = tuple->getColumns();
130const DataTypes & tuple_types = type_tuple->getElements();
131size_t tuple_size = tuple_columns.size();
132for (size_t i = 0; i < tuple_size; ++i)
133columns_of_key_columns.emplace_back(tuple_columns[i], tuple_types[i], "_" + toString(i));
134}
135else
136columns_of_key_columns.emplace_back(left_arg);
137
138bool is_const = false;
139if (columns_of_key_columns.size() == 1)
140{
141auto & arg = columns_of_key_columns.at(0);
142const auto * col = arg.column.get();
143if (const auto * const_col = typeid_cast<const ColumnConst *>(col))
144{
145col = &const_col->getDataColumn();
146is_const = true;
147}
148}
149
150auto res = set->execute(columns_of_key_columns, negative);
151
152if (is_const)
153res = ColumnUInt8::create(input_rows_count, res->getUInt(0));
154
155if (res->size() != input_rows_count)
156throw Exception(ErrorCodes::LOGICAL_ERROR, "Output size is different from input size, expect {}, get {}", input_rows_count, res->size());
157
158return res;
159}
160};
161
162template<bool ignore_set>
163void registerFunctionsInImpl(FunctionFactory & factory)
164{
165factory.registerFunction<FunctionIn<false, false, true, ignore_set>>();
166factory.registerFunction<FunctionIn<false, true, true, ignore_set>>();
167factory.registerFunction<FunctionIn<true, false, true, ignore_set>>();
168factory.registerFunction<FunctionIn<true, true, true, ignore_set>>();
169factory.registerFunction<FunctionIn<false, false, false, ignore_set>>();
170factory.registerFunction<FunctionIn<false, true, false, ignore_set>>();
171factory.registerFunction<FunctionIn<true, false, false, ignore_set>>();
172factory.registerFunction<FunctionIn<true, true, false, ignore_set>>();
173}
174
175}
176
177REGISTER_FUNCTION(In)
178{
179registerFunctionsInImpl<false>(factory);
180registerFunctionsInImpl<true>(factory);
181}
182
183}
184