1 | #include <Functions/IFunctionImpl.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/ColumnTuple.h> |
8 | #include <Columns/ColumnSet.h> |
9 | #include <Interpreters/Set.h> |
10 | |
11 | |
12 | namespace DB |
13 | { |
14 | |
15 | namespace ErrorCodes |
16 | { |
17 | extern const int ILLEGAL_COLUMN; |
18 | } |
19 | |
20 | /** in(x, set) - function for evaluating the IN |
21 | * notIn(x, set) - and NOT IN. |
22 | */ |
23 | |
24 | template <bool negative, bool global> |
25 | struct FunctionInName; |
26 | |
27 | template <> |
28 | struct FunctionInName<false, false> |
29 | { |
30 | static constexpr auto name = "in" ; |
31 | }; |
32 | |
33 | template <> |
34 | struct FunctionInName<false, true> |
35 | { |
36 | static constexpr auto name = "globalIn" ; |
37 | }; |
38 | |
39 | template <> |
40 | struct FunctionInName<true, false> |
41 | { |
42 | static constexpr auto name = "notIn" ; |
43 | }; |
44 | |
45 | template <> |
46 | struct FunctionInName<true, true> |
47 | { |
48 | static constexpr auto name = "globalNotIn" ; |
49 | }; |
50 | |
51 | template <bool negative, bool global> |
52 | class FunctionIn : public IFunction |
53 | { |
54 | public: |
55 | static constexpr auto name = FunctionInName<negative, global>::name; |
56 | static FunctionPtr create(const Context &) |
57 | { |
58 | return std::make_shared<FunctionIn>(); |
59 | } |
60 | |
61 | String getName() const override |
62 | { |
63 | return name; |
64 | } |
65 | |
66 | size_t getNumberOfArguments() const override |
67 | { |
68 | return 2; |
69 | } |
70 | |
71 | DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override |
72 | { |
73 | return std::make_shared<DataTypeUInt8>(); |
74 | } |
75 | |
76 | bool useDefaultImplementationForConstants() const override { return true; } |
77 | |
78 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t /*input_rows_count*/) override |
79 | { |
80 | /// NOTE: after updating this code, check that FunctionIgnoreExceptNull returns the same type of column. |
81 | |
82 | /// Second argument must be ColumnSet. |
83 | ColumnPtr column_set_ptr = block.getByPosition(arguments[1]).column; |
84 | const ColumnSet * column_set = checkAndGetColumnConstData<const ColumnSet>(column_set_ptr.get()); |
85 | if (!column_set) |
86 | column_set = checkAndGetColumn<const ColumnSet>(column_set_ptr.get()); |
87 | if (!column_set) |
88 | throw Exception("Second argument for function '" + getName() + "' must be Set; found " + column_set_ptr->getName(), |
89 | ErrorCodes::ILLEGAL_COLUMN); |
90 | |
91 | Block block_of_key_columns; |
92 | |
93 | /// First argument may be a tuple or a single column. |
94 | const ColumnWithTypeAndName & left_arg = block.getByPosition(arguments[0]); |
95 | const ColumnTuple * tuple = typeid_cast<const ColumnTuple *>(left_arg.column.get()); |
96 | const ColumnConst * const_tuple = checkAndGetColumnConst<ColumnTuple>(left_arg.column.get()); |
97 | const DataTypeTuple * type_tuple = typeid_cast<const DataTypeTuple *>(left_arg.type.get()); |
98 | |
99 | ColumnPtr materialized_tuple; |
100 | if (const_tuple) |
101 | { |
102 | materialized_tuple = const_tuple->convertToFullColumn(); |
103 | tuple = typeid_cast<const ColumnTuple *>(materialized_tuple.get()); |
104 | } |
105 | |
106 | auto set = column_set->getData(); |
107 | auto set_types = set->getDataTypes(); |
108 | if (tuple && (set_types.size() != 1 || !set_types[0]->equals(*type_tuple))) |
109 | { |
110 | const auto & tuple_columns = tuple->getColumns(); |
111 | const DataTypes & tuple_types = type_tuple->getElements(); |
112 | size_t tuple_size = tuple_columns.size(); |
113 | for (size_t i = 0; i < tuple_size; ++i) |
114 | block_of_key_columns.insert({ tuple_columns[i], tuple_types[i], "" }); |
115 | } |
116 | else |
117 | block_of_key_columns.insert(left_arg); |
118 | |
119 | block.getByPosition(result).column = set->execute(block_of_key_columns, negative); |
120 | } |
121 | }; |
122 | |
123 | |
124 | void registerFunctionsIn(FunctionFactory & factory) |
125 | { |
126 | factory.registerFunction<FunctionIn<false, false>>(); |
127 | factory.registerFunction<FunctionIn<false, true>>(); |
128 | factory.registerFunction<FunctionIn<true, false>>(); |
129 | factory.registerFunction<FunctionIn<true, true>>(); |
130 | } |
131 | |
132 | } |
133 | |