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
|
#include <Interpreters/ExpressionActions.h>
#include <Interpreters/convertFieldToType.h>
#include <Parsers/ASTFunction.h>
#include <Parsers/ASTLiteral.h>
#include <Parsers/ASTIdentifier.h>
#include <DataTypes/DataTypesNumber.h>
#include <Interpreters/OptimizeShardingKeyRewriteInVisitor.h>
namespace
{
using namespace DB;
Field executeFunctionOnField(
const Field & field,
const std::string & name,
const ExpressionActionsPtr & sharding_expr,
const DataTypePtr & type,
const std::string & sharding_key_column_name)
{
ColumnWithTypeAndName column;
column.column = type->createColumnConst(1, field);
column.name = name;
column.type = type;
Block block{column};
size_t num_rows = 1;
sharding_expr->execute(block, num_rows);
ColumnWithTypeAndName & ret = block.getByName(sharding_key_column_name);
return (*ret.column)[0];
}
/// @param sharding_column_value - one of values from IN
/// @param sharding_column_name - name of that column
/// @return true if shard may contain such value (or it is unknown), otherwise false.
bool shardContains(
Field sharding_column_value,
const std::string & sharding_column_name,
const OptimizeShardingKeyRewriteInMatcher::Data & data)
{
/// Implicit conversion.
sharding_column_value = convertFieldToType(sharding_column_value, *data.sharding_key_type);
/// NULL is not allowed in sharding key,
/// so it should be safe to assume that shard cannot contain it.
if (sharding_column_value.isNull())
return false;
Field sharding_value = executeFunctionOnField(
sharding_column_value, sharding_column_name,
data.sharding_key_expr, data.sharding_key_type,
data.sharding_key_column_name);
/// The value from IN can be non-numeric,
/// but in this case it should be convertible to numeric type, let's try.
///
/// NOTE: that conversion should not be done for signed types,
/// since it uses accurate cast, that will return Null,
/// but we need static_cast<> (as createBlockSelector()).
if (!isInt64OrUInt64FieldType(sharding_value.getType()))
sharding_value = convertFieldToType(sharding_value, DataTypeUInt64());
/// In case of conversion is not possible (NULL), shard cannot contain the value anyway.
if (sharding_value.isNull())
return false;
UInt64 value = sharding_value.get<UInt64>();
const auto shard_num = data.slots[value % data.slots.size()] + 1;
return data.shard_info.shard_num == shard_num;
}
}
namespace DB
{
bool OptimizeShardingKeyRewriteInMatcher::needChildVisit(ASTPtr & /*node*/, const ASTPtr & /*child*/)
{
return true;
}
void OptimizeShardingKeyRewriteInMatcher::visit(ASTPtr & node, Data & data)
{
if (auto * function = node->as<ASTFunction>())
visit(*function, data);
}
void OptimizeShardingKeyRewriteInMatcher::visit(ASTFunction & function, Data & data)
{
if (function.name != "in")
return;
auto * left = function.arguments->children.front().get();
auto * right = function.arguments->children.back().get();
auto * identifier = left->as<ASTIdentifier>();
if (!identifier)
return;
if (!data.sharding_key_expr->getRequiredColumnsWithTypes().contains(identifier->name()))
return;
/// NOTE: that we should not take care about empty tuple,
/// since after optimize_skip_unused_shards,
/// at least one element should match each shard.
if (auto * tuple_func = right->as<ASTFunction>(); tuple_func && tuple_func->name == "tuple")
{
auto * tuple_elements = tuple_func->children.front()->as<ASTExpressionList>();
std::erase_if(tuple_elements->children, [&](auto & child)
{
auto * literal = child->template as<ASTLiteral>();
return literal && !shardContains(literal->value, identifier->name(), data);
});
}
else if (auto * tuple_literal = right->as<ASTLiteral>();
tuple_literal && tuple_literal->value.getType() == Field::Types::Tuple)
{
auto & tuple = tuple_literal->value.get<Tuple &>();
std::erase_if(tuple, [&](auto & child)
{
return !shardContains(child, identifier->name(), data);
});
}
}
}
|
