ClickHouse

1
#include <Columns/ColumnString.h>
2
#include <Core/Block.h>
3
#include <Functions/FunctionFactory.h>
4
#include <Functions/FunctionHelpers.h>
5
#include <Functions/IFunction.h>
6
#include <Interpreters/Context.h>
7
#include <Interpreters/DatabaseCatalog.h>
8
#include <Interpreters/HashJoin.h>
9
#include <Storages/StorageJoin.h>
10
#include <Storages/TableLockHolder.h>
11

12

13
namespace DB
14
{
15

16
namespace ErrorCodes
17
{
18
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
19
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
20
}
21

22
class HashJoin;
23
using StorageJoinPtr = std::shared_ptr<StorageJoin>;
24

25
namespace
26
{
27

28
template <bool or_null>
29
class ExecutableFunctionJoinGet final : public IExecutableFunction, WithContext
30
{
31
public:
32
    ExecutableFunctionJoinGet(ContextPtr context_,
33
                              TableLockHolder table_lock_,
34
                              StorageJoinPtr storage_join_,
35
                              const DB::Block & result_columns_)
36
        : WithContext(context_)
37
        , table_lock(std::move(table_lock_))
38
        , storage_join(std::move(storage_join_))
39
        , result_columns(result_columns_)
40
    {}
41

42
    static constexpr auto name = or_null ? "joinGetOrNull" : "joinGet";
43

44
    bool useDefaultImplementationForNulls() const override { return false; }
45
    bool useDefaultImplementationForLowCardinalityColumns() const override { return false; }
46
    bool useDefaultImplementationForConstants() const override { return true; }
47

48
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override;
49

50
    String getName() const override { return name; }
51

52
private:
53
    TableLockHolder table_lock;
54
    StorageJoinPtr storage_join;
55
    DB::Block result_columns;
56
};
57

58
template <bool or_null>
59
class FunctionJoinGet final : public IFunctionBase, WithContext
60
{
61
public:
62
    static constexpr auto name = or_null ? "joinGetOrNull" : "joinGet";
63

64
    FunctionJoinGet(ContextPtr context_,
65
                    TableLockHolder table_lock_,
66
                    StorageJoinPtr storage_join_, String attr_name_,
67
                    DataTypes argument_types_, DataTypePtr return_type_)
68
        : WithContext(context_)
69
        , table_lock(std::move(table_lock_))
70
        , storage_join(storage_join_)
71
        , attr_name(std::move(attr_name_))
72
        , argument_types(std::move(argument_types_))
73
        , return_type(std::move(return_type_))
74
    {
75
    }
76

77
    String getName() const override { return name; }
78

79
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
80

81
    const DataTypes & getArgumentTypes() const override { return argument_types; }
82
    const DataTypePtr & getResultType() const override { return return_type; }
83

84
    ExecutableFunctionPtr prepare(const ColumnsWithTypeAndName &) const override;
85

86
private:
87
    TableLockHolder table_lock;
88
    StorageJoinPtr storage_join;
89
    const String attr_name;
90
    DataTypes argument_types;
91
    DataTypePtr return_type;
92
};
93

94
template <bool or_null>
95
class JoinGetOverloadResolver final : public IFunctionOverloadResolver, WithContext
96
{
97
public:
98
    static constexpr auto name = or_null ? "joinGetOrNull" : "joinGet";
99
    static FunctionOverloadResolverPtr create(ContextPtr context_) { return std::make_unique<JoinGetOverloadResolver>(context_); }
100

101
    explicit JoinGetOverloadResolver(ContextPtr context_) : WithContext(context_) {}
102

103
    bool isDeterministic() const override { return false; }
104
    String getName() const override { return name; }
105

106
    FunctionBasePtr buildImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &) const override;
107
    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName &) const override { return {}; } // Not used
108

109
    bool useDefaultImplementationForNulls() const override { return false; }
110
    bool useDefaultImplementationForLowCardinalityColumns() const override { return false; }
111

112
    bool isVariadic() const override { return true; }
113
    size_t getNumberOfArguments() const override { return 0; }
114
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {0, 1}; }
115
};
116

117

118
template <bool or_null>
119
ColumnPtr ExecutableFunctionJoinGet<or_null>::executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const
120
{
121
    ColumnsWithTypeAndName keys;
122
    for (size_t i = 2; i < arguments.size(); ++i)
123
    {
124
        auto key = arguments[i];
125
        keys.emplace_back(std::move(key));
126
    }
127
    return storage_join->joinGet(keys, result_columns, getContext()).column;
128
}
129

130
template <bool or_null>
131
ExecutableFunctionPtr FunctionJoinGet<or_null>::prepare(const ColumnsWithTypeAndName &) const
132
{
133
    Block result_columns {{return_type->createColumn(), return_type, attr_name}};
134
    return std::make_unique<ExecutableFunctionJoinGet<or_null>>(getContext(), table_lock, storage_join, result_columns);
135
}
136

137
std::pair<std::shared_ptr<StorageJoin>, String>
138
getJoin(const ColumnsWithTypeAndName & arguments, ContextPtr context)
139
{
140
    String join_name;
141
    if (const auto * name_col = checkAndGetColumnConst<ColumnString>(arguments[0].column.get()))
142
    {
143
        join_name = name_col->getValue<String>();
144
    }
145
    else
146
        throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
147
                        "Illegal type {} of first argument of function joinGet, expected a const string.",
148
                        arguments[0].type->getName());
149

150
    auto qualified_name = QualifiedTableName::parseFromString(join_name);
151
    if (qualified_name.database.empty())
152
        qualified_name.database = context->getCurrentDatabase();
153

154
    auto table = DatabaseCatalog::instance().getTable({qualified_name.database, qualified_name.table}, std::const_pointer_cast<Context>(context));
155
    auto storage_join = std::dynamic_pointer_cast<StorageJoin>(table);
156
    if (!storage_join)
157
        throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Table {} should have engine StorageJoin", join_name);
158

159
    String attr_name;
160
    if (const auto * name_col = checkAndGetColumnConst<ColumnString>(arguments[1].column.get()))
161
    {
162
        attr_name = name_col->getValue<String>();
163
    }
164
    else
165
        throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
166
                        "Illegal type {} of second argument of function joinGet, expected a const string.",
167
                        arguments[1].type->getName());
168
    return std::make_pair(storage_join, attr_name);
169
}
170

171
template <bool or_null>
172
FunctionBasePtr JoinGetOverloadResolver<or_null>::buildImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &) const
173
{
174
    if (arguments.size() < 3)
175
        throw Exception(
176
            ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
177
            "Number of arguments for function '{}' doesn't match: passed {}, should be greater or equal to 3",
178
            getName() , arguments.size());
179
    auto [storage_join, attr_name] = getJoin(arguments, getContext());
180
    DataTypes data_types(arguments.size() - 2);
181
    DataTypes argument_types(arguments.size());
182
    for (size_t i = 0; i < arguments.size(); ++i)
183
    {
184
        if (i >= 2)
185
            data_types[i - 2] = arguments[i].type;
186
        argument_types[i] = arguments[i].type;
187
    }
188

189
    auto return_type = storage_join->joinGetCheckAndGetReturnType(data_types, attr_name, or_null || storage_join->useNulls());
190
    auto table_lock = storage_join->lockForShare(getContext()->getInitialQueryId(), getContext()->getSettingsRef().lock_acquire_timeout);
191

192
    if (storage_join->useNulls())
193
        return std::make_unique<FunctionJoinGet<true>>(getContext(), table_lock, storage_join, attr_name, argument_types, return_type);
194

195
    return std::make_unique<FunctionJoinGet<or_null>>(getContext(), table_lock, storage_join, attr_name, argument_types, return_type);
196
}
197

198
}
199

200
REGISTER_FUNCTION(JoinGet)
201
{
202
    // joinGet
203
    factory.registerFunction<JoinGetOverloadResolver<false>>();
204
    // joinGetOrNull
205
    factory.registerFunction<JoinGetOverloadResolver<true>>();
206
}
207

208
}
209