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
|
#include "mkql_computation_node_graph_saveload.h"
#include "mkql_computation_node_holders.h"
#include <yql/essentials/minikql/pack_num.h>
#include <yql/essentials/minikql/comp_nodes/mkql_saveload.h>
namespace NKikimr {
namespace NMiniKQL {
namespace {
void TraverseGraph(const NUdf::TUnboxedValue* roots, ui32 rootCount, TVector<NUdf::TUnboxedValue>& values) {
THashSet<NUdf::IBoxedValue*> dedup;
for (ui32 i = 0; i < rootCount; ++i) {
const auto& value = roots[i];
if (!value.IsBoxed()) {
continue;
}
auto* ptr = value.AsBoxed().Get();
if (dedup.contains(ptr)) {
continue;
}
dedup.insert(ptr);
values.push_back(value);
}
for (ui32 from = 0, to = values.size(); from != to; ++from) {
auto current = values[from];
auto count = current.GetTraverseCount();
for (ui32 i = 0; i < count; ++i) {
auto value = current.GetTraverseItem(i);
if (!value.IsBoxed()) {
continue;
}
auto* ptr = value.AsBoxed().Get();
if (dedup.contains(ptr)) {
continue;
}
dedup.insert(ptr);
values.push_back(value);
++to;
}
}
}
}
void SaveGraphState(const NUdf::TUnboxedValue* roots, ui32 rootCount, ui64 hash, TString& out) {
out.clear();
out.AppendNoAlias((const char*)&hash, sizeof(hash));
TVector<NUdf::TUnboxedValue> values;
TraverseGraph(roots, rootCount, values);
for (ui32 i = 0; i < values.size(); ++i) {
auto state = values[i].Save();
if (state.IsString() || state.IsEmbedded()) {
auto strRef = state.AsStringRef();
auto size = strRef.Size();
WriteUi64(out, size);
if (size) {
out.AppendNoAlias(strRef.Data(), size);
}
}
else if (state.IsBoxed()) {
TString taskState;
auto listIt = state.GetListIterator();
NUdf::TUnboxedValue str;
while (listIt.Next(str)) {
const TStringBuf strRef = str.AsStringRef();
taskState.AppendNoAlias(strRef.Data(), strRef.Size());
}
WriteUi64(out, taskState.size());
if (!taskState.empty()) {
out.AppendNoAlias(taskState.data(), taskState.size());
}
}
}
}
void LoadGraphState(const NUdf::TUnboxedValue* roots, ui32 rootCount, ui64 hash, const TStringBuf& in) {
TStringBuf state(in);
MKQL_ENSURE(state.size() >= sizeof(ui64), "Serialized state is corrupted - no hash");
ui64 storedHash = *(ui64*)state.data();
state.Skip(sizeof(storedHash));
MKQL_ENSURE(hash == storedHash, "Unable to load graph state, different hashes");
TVector<NUdf::TUnboxedValue> values;
TraverseGraph(roots, rootCount, values);
for (ui32 i = 0; i < values.size(); ++i) {
auto size = ReadUi64(state);
if (size) {
MKQL_ENSURE(size <= state.size(), "Serialized state is corrupted");
values[i].Load(NUdf::TStringRef(state.data(), size));
state.Skip(size);
}
}
MKQL_ENSURE(state.size() == 0, "State was not loaded correctly");
}
} // namespace NMiniKQL
} // namespace NKikimr
|
