blob: 13b6311a877d5bbd5f12270b2d1e41f20a21cbd4 (
plain) (
blame)
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
|
#include <Common/typeid_cast.h>
#include <Common/checkStackSize.h>
#include <Parsers/ASTLiteral.h>
#include <Parsers/ASTFunction.h>
#include <Parsers/ASTHelpers.h>
#include <Parsers/ASTExpressionList.h>
#include <Interpreters/OptimizeIfWithConstantConditionVisitor.h>
#include <IO/WriteHelpers.h>
namespace DB
{
namespace ErrorCodes
{
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
}
static bool tryExtractConstValueFromCondition(const ASTPtr & condition, bool & value)
{
/// numeric constant in condition
if (const auto * literal = condition->as<ASTLiteral>())
{
if (literal->value.getType() == Field::Types::Int64 ||
literal->value.getType() == Field::Types::UInt64)
{
value = literal->value.get<Int64>();
return true;
}
}
/// cast of numeric constant in condition to UInt8
/// Note: this solution is ad-hoc and only implemented for metrica use case (one of the best customers).
/// We should allow any constant condition (or maybe remove this optimization completely) later.
if (const auto * function = condition->as<ASTFunction>())
{
if (isFunctionCast(function))
{
if (const auto * expr_list = function->arguments->as<ASTExpressionList>())
{
if (expr_list->children.size() != 2)
throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function CAST must have exactly two arguments");
const ASTPtr & type_ast = expr_list->children.at(1);
if (const auto * type_literal = type_ast->as<ASTLiteral>())
{
if (type_literal->value.getType() == Field::Types::String)
{
const auto & type_str = type_literal->value.get<std::string>();
if (type_str == "UInt8" || type_str == "Nullable(UInt8)")
return tryExtractConstValueFromCondition(expr_list->children.at(0), value);
}
}
}
}
else if (function->name == "toUInt8" || function->name == "toInt8" || function->name == "identity")
{
if (const auto * expr_list = function->arguments->as<ASTExpressionList>())
{
if (expr_list->children.size() != 1)
throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} must have exactly two arguments", function->name);
return tryExtractConstValueFromCondition(expr_list->children.at(0), value);
}
}
}
return false;
}
void OptimizeIfWithConstantConditionVisitor::visit(ASTPtr & current_ast)
{
if (!current_ast)
return;
checkStackSize();
for (ASTPtr & child : current_ast->children)
{
auto * function_node = child->as<ASTFunction>();
if (!function_node || function_node->name != "if")
{
visit(child);
continue;
}
if (!function_node->arguments)
throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Wrong number of arguments for function 'if' (0 instead of 3)");
if (function_node->arguments->children.size() != 3)
throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
"Wrong number of arguments for function 'if' ({} instead of 3)",
function_node->arguments->children.size());
visit(function_node->arguments);
const auto * args = function_node->arguments->as<ASTExpressionList>();
ASTPtr condition_expr = args->children[0];
ASTPtr then_expr = args->children[1];
ASTPtr else_expr = args->children[2];
bool condition;
if (tryExtractConstValueFromCondition(condition_expr, condition))
{
ASTPtr replace_ast = condition ? then_expr : else_expr;
ASTPtr child_copy = child;
String replace_alias = replace_ast->tryGetAlias();
String if_alias = child->tryGetAlias();
if (replace_alias.empty())
{
replace_ast->setAlias(if_alias);
child = replace_ast;
}
else
{
/// Only copy of one node is required here.
/// But IAST has only method for deep copy of subtree.
/// This can be a reason of performance degradation in case of deep queries.
ASTPtr replace_ast_deep_copy = replace_ast->clone();
replace_ast_deep_copy->setAlias(if_alias);
child = replace_ast_deep_copy;
}
if (!if_alias.empty())
{
auto alias_it = aliases.find(if_alias);
if (alias_it != aliases.end() && alias_it->second.get() == child_copy.get())
alias_it->second = child;
}
}
}
}
}
|
