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
12namespace DB
13{
14
15namespace 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
24template <bool negative, bool global>
25struct FunctionInName;
26
27template <>
28struct FunctionInName<false, false>
29{
30 static constexpr auto name = "in";
31};
32
33template <>
34struct FunctionInName<false, true>
35{
36 static constexpr auto name = "globalIn";
37};
38
39template <>
40struct FunctionInName<true, false>
41{
42 static constexpr auto name = "notIn";
43};
44
45template <>
46struct FunctionInName<true, true>
47{
48 static constexpr auto name = "globalNotIn";
49};
50
51template <bool negative, bool global>
52class FunctionIn : public IFunction
53{
54public:
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
124void 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