1 | #pragma once |
2 | |
3 | #include <Common/typeid_cast.h> |
4 | #include <Core/Block.h> |
5 | #include <DataTypes/DataTypesNumber.h> |
6 | #include <Interpreters/Context.h> |
7 | #include <Parsers/ASTExpressionList.h> |
8 | #include <Parsers/ASTSelectQuery.h> |
9 | #include <Parsers/ASTFunction.h> |
10 | #include <Storages/SelectQueryInfo.h> |
11 | #include <Storages/MergeTree/KeyCondition.h> |
12 | |
13 | |
14 | namespace DB |
15 | { |
16 | |
17 | /// Builds reverse polish notation |
18 | template <typename RPNElement> |
19 | class RPNBuilder |
20 | { |
21 | public: |
22 | using RPN = std::vector<RPNElement>; |
23 | using AtomFromASTFunc = std::function< |
24 | bool(const ASTPtr & node, const Context & context, Block & block_with_constants, RPNElement & out)>; |
25 | |
26 | RPNBuilder(const SelectQueryInfo & query_info, const Context & context_, const AtomFromASTFunc & atomFromAST_) |
27 | : context(context_), atomFromAST(atomFromAST_) |
28 | { |
29 | /** Evaluation of expressions that depend only on constants. |
30 | * For the index to be used, if it is written, for example `WHERE Date = toDate(now())`. |
31 | */ |
32 | block_with_constants = KeyCondition::getBlockWithConstants(query_info.query, query_info.syntax_analyzer_result, context); |
33 | |
34 | /// Trasform WHERE section to Reverse Polish notation |
35 | const ASTSelectQuery & select = typeid_cast<const ASTSelectQuery &>(*query_info.query); |
36 | if (select.where()) |
37 | { |
38 | traverseAST(select.where()); |
39 | |
40 | if (select.prewhere()) |
41 | { |
42 | traverseAST(select.prewhere()); |
43 | rpn.emplace_back(RPNElement::FUNCTION_AND); |
44 | } |
45 | } |
46 | else if (select.prewhere()) |
47 | { |
48 | traverseAST(select.prewhere()); |
49 | } |
50 | else |
51 | { |
52 | rpn.emplace_back(RPNElement::FUNCTION_UNKNOWN); |
53 | } |
54 | } |
55 | |
56 | RPN && () { return std::move(rpn); } |
57 | |
58 | private: |
59 | void traverseAST(const ASTPtr & node) |
60 | { |
61 | RPNElement element; |
62 | |
63 | if (ASTFunction * func = typeid_cast<ASTFunction *>(&*node)) |
64 | { |
65 | if (operatorFromAST(func, element)) |
66 | { |
67 | auto & args = typeid_cast<ASTExpressionList &>(*func->arguments).children; |
68 | for (size_t i = 0, size = args.size(); i < size; ++i) |
69 | { |
70 | traverseAST(args[i]); |
71 | |
72 | /** The first part of the condition is for the correct support of `and` and `or` functions of arbitrary arity |
73 | * - in this case `n - 1` elements are added (where `n` is the number of arguments). |
74 | */ |
75 | if (i != 0 || element.function == RPNElement::FUNCTION_NOT) |
76 | rpn.emplace_back(std::move(element)); |
77 | } |
78 | |
79 | return; |
80 | } |
81 | } |
82 | |
83 | if (!atomFromAST(node, context, block_with_constants, element)) |
84 | { |
85 | element.function = RPNElement::FUNCTION_UNKNOWN; |
86 | } |
87 | |
88 | rpn.emplace_back(std::move(element)); |
89 | } |
90 | |
91 | bool operatorFromAST(const ASTFunction * func, RPNElement & out) |
92 | { |
93 | /// Functions AND, OR, NOT. |
94 | /** Also a special function `indexHint` - works as if instead of calling a function there are just parentheses |
95 | * (or, the same thing - calling the function `and` from one argument). |
96 | */ |
97 | const ASTs & args = typeid_cast<const ASTExpressionList &>(*func->arguments).children; |
98 | |
99 | if (func->name == "not" ) |
100 | { |
101 | if (args.size() != 1) |
102 | return false; |
103 | |
104 | out.function = RPNElement::FUNCTION_NOT; |
105 | } |
106 | else |
107 | { |
108 | if (func->name == "and" || func->name == "indexHint" ) |
109 | out.function = RPNElement::FUNCTION_AND; |
110 | else if (func->name == "or" ) |
111 | out.function = RPNElement::FUNCTION_OR; |
112 | else |
113 | return false; |
114 | } |
115 | |
116 | return true; |
117 | } |
118 | |
119 | const Context & context; |
120 | const AtomFromASTFunc & atomFromAST; |
121 | Block block_with_constants; |
122 | RPN rpn; |
123 | }; |
124 | |
125 | |
126 | }; |
127 | |