ClickHouse

1
#include <Processors/QueryPlan/Optimizations/actionsDAGUtils.h>
2

3
#include <Core/Field.h>
4
#include <Functions/IFunction.h>
5
#include <Columns/ColumnConst.h>
6

7
#include <stack>
8

9
namespace DB
10
{
11
MatchedTrees::Matches matchTrees(const ActionsDAG::NodeRawConstPtrs & inner_dag, const ActionsDAG & outer_dag, bool check_monotonicity)
12
{
13
    using Parents = std::set<const ActionsDAG::Node *>;
14
    std::unordered_map<const ActionsDAG::Node *, Parents> inner_parents;
15
    std::unordered_map<std::string_view, const ActionsDAG::Node *> inner_inputs;
16

17
    {
18
        std::stack<const ActionsDAG::Node *> stack;
19
        for (const auto * out : inner_dag)
20
        {
21
            if (inner_parents.contains(out))
22
                continue;
23

24
            stack.push(out);
25
            inner_parents.emplace(out, Parents());
26
            while (!stack.empty())
27
            {
28
                const auto * node = stack.top();
29
                stack.pop();
30

31
                if (node->type == ActionsDAG::ActionType::INPUT)
32
                    inner_inputs.emplace(node->result_name, node);
33

34
                for (const auto * child : node->children)
35
                {
36
                    auto [it, inserted] = inner_parents.emplace(child, Parents());
37
                    it->second.emplace(node);
38

39
                    if (inserted)
40
                        stack.push(child);
41
                }
42
            }
43
        }
44
    }
45

46
    struct Frame
47
    {
48
        const ActionsDAG::Node * node;
49
        ActionsDAG::NodeRawConstPtrs mapped_children;
50
    };
51

52
    MatchedTrees::Matches matches;
53
    std::stack<Frame> stack;
54

55
    for (const auto & node : outer_dag.getNodes())
56
    {
57
        if (matches.contains(&node))
58
            continue;
59

60
        stack.push(Frame{&node, {}});
61
        while (!stack.empty())
62
        {
63
            auto & frame = stack.top();
64
            frame.mapped_children.reserve(frame.node->children.size());
65

66
            while (frame.mapped_children.size() < frame.node->children.size())
67
            {
68
                const auto * child = frame.node->children[frame.mapped_children.size()];
69
                auto it = matches.find(child);
70
                if (it == matches.end())
71
                {
72
                    /// If match map does not contain a child, it was not visited.
73
                    stack.push(Frame{child, {}});
74
                    break;
75
                }
76
                /// A node from found match may be nullptr.
77
                /// It means that node is visited, but no match was found.
78
                if (it->second.monotonicity)
79
                    /// Ignore a match with monotonicity.
80
                    frame.mapped_children.push_back(nullptr);
81
                else
82
                    frame.mapped_children.push_back(it->second.node);
83

84
            }
85

86
            if (frame.mapped_children.size() < frame.node->children.size())
87
                continue;
88

89
            /// Create an empty match for current node.
90
            /// match.node will be set if match is found.
91
            auto & match = matches[frame.node];
92

93
            if (frame.node->type == ActionsDAG::ActionType::INPUT)
94
            {
95
                const ActionsDAG::Node * mapped = nullptr;
96
                if (auto it = inner_inputs.find(frame.node->result_name); it != inner_inputs.end())
97
                    mapped = it->second;
98

99
                match.node = mapped;
100
            }
101
            else if (frame.node->type == ActionsDAG::ActionType::ALIAS)
102
            {
103
                match = matches[frame.node->children.at(0)];
104
            }
105
            else if (frame.node->type == ActionsDAG::ActionType::FUNCTION)
106
            {
107
                //std::cerr << "... Processing " << frame.node->function_base->getName() << std::endl;
108

109
                bool found_all_children = true;
110
                const ActionsDAG::Node * any_child = nullptr;
111
                size_t num_children = frame.node->children.size();
112
                for (size_t i = 0; i < num_children; ++i)
113
                {
114
                    if (frame.mapped_children[i])
115
                        any_child = frame.mapped_children[i];
116
                    else if (!frame.node->children[i]->column || !isColumnConst(*frame.node->children[i]->column))
117
                        found_all_children = false;
118
                }
119

120
                if (found_all_children && any_child)
121
                {
122
                    Parents container;
123
                    Parents * intersection = &inner_parents[any_child];
124

125
                    if (frame.mapped_children.size() > 1)
126
                    {
127
                        std::vector<Parents *> other_parents;
128
                        size_t mapped_children_size = frame.mapped_children.size();
129
                        other_parents.reserve(mapped_children_size);
130
                        for (size_t i = 1; i < mapped_children_size; ++i)
131
                            if (frame.mapped_children[i])
132
                                other_parents.push_back(&inner_parents[frame.mapped_children[i]]);
133

134
                        for (const auto * parent : *intersection)
135
                        {
136
                            bool is_common = true;
137
                            for (const auto * set : other_parents)
138
                            {
139
                                if (!set->contains(parent))
140
                                {
141
                                    is_common = false;
142
                                    break;
143
                                }
144
                            }
145

146
                            if (is_common)
147
                                container.insert(parent);
148
                        }
149

150
                        intersection = &container;
151
                    }
152

153
                    //std::cerr << ".. Candidate parents " << intersection->size() << std::endl;
154

155
                    if (!intersection->empty())
156
                    {
157
                        auto func_name = frame.node->function_base->getName();
158
                        for (const auto * parent : *intersection)
159
                        {
160
                            //std::cerr << ".. candidate " << parent->result_name << std::endl;
161
                            if (parent->type == ActionsDAG::ActionType::FUNCTION && func_name == parent->function_base->getName())
162
                            {
163
                                const auto & children = parent->children;
164
                                if (children.size() == num_children)
165
                                {
166
                                    bool all_children_matched = true;
167
                                    for (size_t i = 0; all_children_matched && i < num_children; ++i)
168
                                    {
169
                                        if (frame.mapped_children[i] == nullptr)
170
                                        {
171
                                            all_children_matched = children[i]->column && isColumnConst(*children[i]->column)
172
                                                && children[i]->result_type->equals(*frame.node->children[i]->result_type)
173
                                                && assert_cast<const ColumnConst &>(*children[i]->column).getField() == assert_cast<const ColumnConst &>(*frame.node->children[i]->column).getField();
174
                                        }
175
                                        else
176
                                            all_children_matched = frame.mapped_children[i] == children[i];
177
                                    }
178

179
                                    if (all_children_matched)
180
                                    {
181
                                        match.node = parent;
182
                                        break;
183
                                    }
184
                                }
185
                            }
186
                        }
187
                    }
188
                }
189

190
                if (!match.node && check_monotonicity && frame.node->function_base->hasInformationAboutMonotonicity())
191
                {
192
                    size_t num_const_args = 0;
193
                    const ActionsDAG::Node * monotonic_child = nullptr;
194
                    for (const auto * child : frame.node->children)
195
                    {
196
                        if (child->column)
197
                            ++num_const_args;
198
                        else
199
                            monotonic_child = child;
200
                    }
201

202
                    if (monotonic_child && num_const_args + 1 == frame.node->children.size())
203
                    {
204
                        const auto & child_match = matches[monotonic_child];
205
                        if (child_match.node)
206
                        {
207
                            auto info = frame.node->function_base->getMonotonicityForRange(*monotonic_child->result_type, {}, {});
208
                            if (info.is_monotonic)
209
                            {
210
                                MatchedTrees::Monotonicity monotonicity;
211
                                monotonicity.direction *= info.is_positive ? 1 : -1;
212
                                monotonicity.strict = info.is_strict;
213

214
                                if (child_match.monotonicity)
215
                                {
216
                                    monotonicity.direction *= child_match.monotonicity->direction;
217
                                    if (!child_match.monotonicity->strict)
218
                                        monotonicity.strict = false;
219
                                }
220

221
                                match.node = child_match.node;
222
                                match.monotonicity = monotonicity;
223
                            }
224
                        }
225
                    }
226
                }
227
            }
228

229
            stack.pop();
230
        }
231
    }
232

233
    return matches;
234
}
235

236
}
237