blob: b0018d474d59d2b5d219085b9c6bbac4eb1df4a4 (
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
|
#include <Analyzer/Passes/IfChainToMultiIfPass.h>
#include <DataTypes/DataTypesNumber.h>
#include <Analyzer/InDepthQueryTreeVisitor.h>
#include <Analyzer/FunctionNode.h>
#include <Functions/FunctionFactory.h>
namespace DB
{
namespace
{
class IfChainToMultiIfPassVisitor : public InDepthQueryTreeVisitorWithContext<IfChainToMultiIfPassVisitor>
{
public:
using Base = InDepthQueryTreeVisitorWithContext<IfChainToMultiIfPassVisitor>;
using Base::Base;
explicit IfChainToMultiIfPassVisitor(FunctionOverloadResolverPtr multi_if_function_ptr_, ContextPtr context)
: Base(std::move(context))
, multi_if_function_ptr(std::move(multi_if_function_ptr_))
{}
void enterImpl(QueryTreeNodePtr & node)
{
if (!getSettings().optimize_if_chain_to_multiif)
return;
auto * function_node = node->as<FunctionNode>();
if (!function_node || function_node->getFunctionName() != "if" || function_node->getArguments().getNodes().size() != 3)
return;
std::vector<QueryTreeNodePtr> multi_if_arguments;
auto & function_node_arguments = function_node->getArguments().getNodes();
multi_if_arguments.insert(multi_if_arguments.end(), function_node_arguments.begin(), function_node_arguments.end());
QueryTreeNodePtr if_chain_node = multi_if_arguments.back();
while (true)
{
/// Check if last `multiIf` argument is `if` function
auto * if_chain_function_node = if_chain_node->as<FunctionNode>();
if (!if_chain_function_node || if_chain_function_node->getFunctionName() != "if" || if_chain_function_node->getArguments().getNodes().size() != 3)
break;
/// Replace last `multiIf` argument with `if` function arguments
multi_if_arguments.pop_back();
auto & if_chain_function_node_arguments = if_chain_function_node->getArguments().getNodes();
multi_if_arguments.insert(multi_if_arguments.end(), if_chain_function_node_arguments.begin(), if_chain_function_node_arguments.end());
/// Use last `multiIf` argument for next check
if_chain_node = multi_if_arguments.back();
}
/// Do not replace `if` with 3 arguments to `multiIf`
if (multi_if_arguments.size() <= 3)
return;
auto multi_if_function = std::make_shared<FunctionNode>("multiIf");
multi_if_function->getArguments().getNodes() = std::move(multi_if_arguments);
multi_if_function->resolveAsFunction(multi_if_function_ptr->build(multi_if_function->getArgumentColumns()));
node = std::move(multi_if_function);
}
private:
FunctionOverloadResolverPtr multi_if_function_ptr;
};
}
void IfChainToMultiIfPass::run(QueryTreeNodePtr query_tree_node, ContextPtr context)
{
auto multi_if_function_ptr = FunctionFactory::instance().get("multiIf", context);
IfChainToMultiIfPassVisitor visitor(std::move(multi_if_function_ptr), std::move(context));
visitor.visit(query_tree_node);
}
}
|
