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
|
#include "yql_co.h"
#include <yql/essentials/core/yql_join.h>
#include <yql/essentials/core/yql_opt_utils.h>
#include <yql/essentials/utils/log/log.h>
namespace NYql {
namespace {
using namespace NNodes;
TExprNode::TPtr EquiJoinPushAny(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
// Add any to EquiJoin tree
if (!IsEmitPruneKeysEnabled(optCtx.Types)) {
return node;
}
auto equiJoin = TCoEquiJoin(node);
auto equiJoinTree = equiJoin.Arg(equiJoin.ArgCount() - 2).Cast<TCoEquiJoinTuple>();
if (auto result = PushAnyInEquiJoin(ctx, equiJoinTree); result != equiJoinTree.Ptr()) {
return ctx.ChangeChild(*node, equiJoin.ArgCount() - 2, std::move(result));
}
return node;
}
TExprNode::TPtr EquiJoinEmitPruneKeys(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
// Add PruneKeys to EquiJoin inputs
if (!IsEmitPruneKeysEnabled(optCtx.Types)) {
return node;
}
auto equiJoin = TCoEquiJoin(node);
if (HasSetting(equiJoin.Arg(equiJoin.ArgCount() - 1).Ref(), "prune_keys_added")) {
return node;
}
THashMap<TStringBuf, THashSet<TStringBuf>> columnsForPruneKeysExtractor;
GetPruneKeysColumnsForJoinLeaves(equiJoin.Arg(equiJoin.ArgCount() - 2).Cast<TCoEquiJoinTuple>(), columnsForPruneKeysExtractor);
TExprNode::TListType children;
bool hasChanges = false;
for (size_t i = 0; i + 2 < equiJoin.ArgCount(); ++i) {
auto child = equiJoin.Arg(i).Cast<TCoEquiJoinInput>();
auto list = child.List();
auto scope = child.Scope();
if (!scope.Ref().IsAtom()) {
children.push_back(equiJoin.Arg(i).Ptr());
continue;
}
THashSet<TString> columns;
auto itemNames = columnsForPruneKeysExtractor.find(scope.Ref().Content());
if (itemNames == columnsForPruneKeysExtractor.end() || itemNames->second.empty()) {
children.push_back(equiJoin.Arg(i).Ptr());
continue;
}
for (const auto& elem : itemNames->second) {
columns.insert(TString(elem));
}
if (IsAlreadyDistinct(list.Ref(), columns)) {
children.push_back(equiJoin.Arg(i).Ptr());
continue;
}
auto pruneKeysCallable = IsOrdered(list.Ref(), columns) ? "PruneAdjacentKeys" : "PruneKeys";
YQL_CLOG(DEBUG, Core) << "Add " << pruneKeysCallable << " to EquiJoin input #" << i << ", label " << scope.Ref().Content();
children.push_back(ctx.Builder(child.Pos())
.List()
.Callable(0, pruneKeysCallable)
.Add(0, list.Ptr())
.Add(1, MakePruneKeysExtractorLambda(child.Ref(), columns, ctx))
.Seal()
.Add(1, scope.Ptr())
.Seal()
.Build());
hasChanges = true;
}
if (!hasChanges) {
return node;
}
children.push_back(equiJoin.Arg(equiJoin.ArgCount() - 2).Ptr());
children.push_back(AddSetting(
equiJoin.Arg(equiJoin.ArgCount() - 1).Ref(),
equiJoin.Arg(equiJoin.ArgCount() - 1).Pos(),
"prune_keys_added",
nullptr,
ctx));
return ctx.ChangeChildren(*node, std::move(children));
}
} // namespace
void RegisterCoFlowCallables3(TCallableOptimizerMap& map) {
using namespace std::placeholders;
map["EquiJoin"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
if (auto ret = EquiJoinPushAny(node, ctx, optCtx); ret != node) {
YQL_CLOG(DEBUG, Core) << "EquiJoinPushAny";
return ret;
}
if (auto ret = EquiJoinEmitPruneKeys(node, ctx, optCtx); ret != node) {
YQL_CLOG(DEBUG, Core) << "EquiJoinEmitPruneKeys";
return ret;
}
return node;
};
}
}
|
