1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
|
#include <Functions/IFunction.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <DataTypes/DataTypeString.h>
#include <DataTypes/DataTypesNumber.h>
#include <Columns/ColumnString.h>
#include <Interpreters/Context.h>
#include <Common/Macros.h>
#include <Core/Field.h>
namespace DB
{
namespace ErrorCodes
{
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
}
namespace
{
/** Get scalar value of sub queries from query context via IAST::Hash.
*/
class FunctionGetScalar : public IFunction, WithContext
{
public:
static constexpr auto name = "__getScalar";
static FunctionPtr create(ContextPtr context_)
{
return std::make_shared<FunctionGetScalar>(context_);
}
explicit FunctionGetScalar(ContextPtr context_) : WithContext(context_) {}
String getName() const override
{
return name;
}
size_t getNumberOfArguments() const override
{
return 1;
}
bool useDefaultImplementationForLowCardinalityColumns() const override
{
return false;
}
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
{
if (arguments.size() != 1 || !isString(arguments[0].type) || !arguments[0].column || !isColumnConst(*arguments[0].column))
throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Function {} accepts one const string argument", getName());
auto scalar_name = assert_cast<const ColumnConst &>(*arguments[0].column).getValue<String>();
ContextPtr query_context = getContext()->hasQueryContext() ? getContext()->getQueryContext() : getContext();
scalar = query_context->getScalar(scalar_name).getByPosition(0);
return scalar.type;
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName &, const DataTypePtr &, size_t input_rows_count) const override
{
return ColumnConst::create(scalar.column, input_rows_count);
}
private:
mutable ColumnWithTypeAndName scalar;
};
/** Get special scalar values
*/
template <typename Scalar>
class FunctionGetSpecialScalar : public IFunction
{
public:
static constexpr auto name = Scalar::name;
static FunctionPtr create(ContextPtr context_)
{
return std::make_shared<FunctionGetSpecialScalar<Scalar>>(context_);
}
static ColumnWithTypeAndName createScalar(ContextPtr context_)
{
if (const auto * block = context_->tryGetSpecialScalar(Scalar::scalar_name))
return block->getByPosition(0);
else if (context_->hasQueryContext())
{
if (context_->getQueryContext()->hasScalar(Scalar::scalar_name))
return context_->getQueryContext()->getScalar(Scalar::scalar_name).getByPosition(0);
}
return {DataTypeUInt32().createColumnConst(1, 0), std::make_shared<DataTypeUInt32>(), Scalar::scalar_name};
}
explicit FunctionGetSpecialScalar(ContextPtr context_)
: scalar(createScalar(context_)), is_distributed(context_->isDistributed())
{
}
String getName() const override
{
return name;
}
bool isDeterministic() const override { return false; }
bool isSuitableForConstantFolding() const override { return !is_distributed; }
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
size_t getNumberOfArguments() const override
{
return 0;
}
DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName &) const override
{
return scalar.type;
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName &, const DataTypePtr &, size_t input_rows_count) const override
{
auto result = ColumnConst::create(scalar.column, input_rows_count);
if (!isSuitableForConstantFolding())
return result->convertToFullColumnIfConst();
return result;
}
private:
ColumnWithTypeAndName scalar;
bool is_distributed;
};
struct GetShardNum
{
static constexpr auto name = "shardNum";
static constexpr auto scalar_name = "_shard_num";
};
struct GetShardCount
{
static constexpr auto name = "shardCount";
static constexpr auto scalar_name = "_shard_count";
};
}
REGISTER_FUNCTION(GetScalar)
{
factory.registerFunction<FunctionGetScalar>();
factory.registerFunction<FunctionGetSpecialScalar<GetShardNum>>();
factory.registerFunction<FunctionGetSpecialScalar<GetShardCount>>();
}
}
|
