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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
|
#pragma once
#include <yql/essentials/ast/yql_expr.h>
#include <yql/essentials/minikql/comp_nodes/ut/mkql_computation_node_ut.h>
#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
namespace NKikimr::NMiniKQL {
// Template to count TMaybe nesting levels
template <typename T>
struct TMaybeTraits {
static constexpr ui32 value = 0;
using ResultType = T;
};
template <typename T>
struct TMaybeTraits<TMaybe<T>> {
static constexpr ui32 value = 1 + TMaybeTraits<T>::value;
using ResultType = TMaybeTraits<T>::ResultType;
};
// Template to get the value of the deepest set TMaybe or return 0
template <typename T>
TMaybeTraits<T>::ResultType GetInnerValue(const T& value) {
if constexpr (TMaybeTraits<T>::value == 0) {
// Base case: not a TMaybe type, return the value itself
return value;
} else {
// TMaybe type
if (value.Defined()) {
return GetInnerValue(value.GetRef());
} else {
return {}; // Return 0 if not set
}
}
}
// Template to get the level of the set TMaybe (returns 0 if not set)
template <typename T>
ui32 GetSettedLevel(const T& value) {
if constexpr (TMaybeTraits<T>::value == 0) {
// Base case: not a TMaybe type, always set at level 0
return 0;
} else {
// TMaybe type
if (value.Defined()) {
return 1 + GetSettedLevel(value.GetRef());
} else {
return 0; // Return 0 if not set
}
}
}
template <typename T>
struct TUnpackedMaybe {
ui32 SettedLevel;
ui32 MaybeLevel;
T Value;
};
class TBlockHelper {
public:
explicit TBlockHelper()
: Setup_(GetNodeTestFactory())
, Pb_(*Setup_.PgmBuilder)
{
}
template <typename T>
TRuntimeNode ConvertValueToLiteralNode(T node) = delete;
template <typename T>
TRuntimeNode ConvertValueToLiteralNode(T simpleNode)
requires(NYql::NUdf::TPrimitiveDataType<T>::Result)
{
return Pb_.NewDataLiteral<T>(simpleNode);
}
template <typename T>
TRuntimeNode ConvertValueToLiteralNode(TMaybe<T> maybeNode) {
TUnpackedMaybe unpacked{.SettedLevel = GetSettedLevel(maybeNode), .MaybeLevel = TMaybeTraits<TMaybe<T>>::value, .Value = GetInnerValue(maybeNode)};
auto data = ConvertValueToLiteralNode(unpacked.Value);
for (ui32 i = unpacked.SettedLevel; i < unpacked.MaybeLevel; i++) {
data = Pb_.NewEmptyOptional(Pb_.NewOptionalType(data.GetStaticType()));
}
for (ui32 i = 0; i < unpacked.SettedLevel; i++) {
data = Pb_.NewOptional(data);
}
return data;
}
template <typename... TArgs, std::size_t... Is>
TRuntimeNode ConvertValueToLiteralNodeTuple(std::tuple<TArgs...> maybeNode, std::index_sequence<Is...>) {
auto data = std::vector<TRuntimeNode>{ConvertValueToLiteralNode(std::get<Is>(maybeNode))...};
return Pb_.NewTuple(data);
}
template <typename... TArgs>
TRuntimeNode ConvertValueToLiteralNode(std::tuple<TArgs...> node) {
return ConvertValueToLiteralNodeTuple(node, std::index_sequence_for<TArgs...>{});
}
template <typename T>
TRuntimeNode ConvertNode(T node) {
return Pb_.AsScalar(ConvertValueToLiteralNode(node));
}
template <typename T>
TRuntimeNode ConvertNode(std::vector<T> nodes) {
TRuntimeNode::TList convertedNodes;
for (auto& node : nodes) {
convertedNodes.push_back(ConvertValueToLiteralNode(node));
}
return ConvertLiteralListToDatum(convertedNodes);
}
TRuntimeNode ConvertLiteralListToDatum(TRuntimeNode::TList nodes);
template <typename T, typename U, typename V>
void TestKernel(T left, U right, V expected, std::function<TRuntimeNode(TSetup<false>&, TRuntimeNode, TRuntimeNode)> binaryOp) {
NYql::TExprContext exprCtx;
auto leftNode = ConvertNode(left);
auto rightNode = ConvertNode(right);
auto expectedNode = ConvertNode(expected);
auto resultValue = Setup_.BuildGraph(binaryOp(Setup_, leftNode, rightNode))->GetValue();
auto expectedValue = Setup_.BuildGraph(expectedNode)->GetValue();
auto outDatum = TArrowBlock::From(resultValue).GetDatum();
auto expectedDatum = TArrowBlock::From(expectedValue).GetDatum();
UNIT_ASSERT_EQUAL_C(outDatum, expectedDatum, "Expected : " << DatumToString(expectedDatum) << "\n but got : " << DatumToString(outDatum));
}
template <typename T, typename V>
void TestKernel(T operand, V expected, std::function<TRuntimeNode(TSetup<false>&, TRuntimeNode)> unaryOp) {
NYql::TExprContext exprCtx;
auto node = ConvertNode(operand);
auto expectedNode = ConvertNode(expected);
auto resultValue = Setup_.BuildGraph(unaryOp(Setup_, node))->GetValue();
auto expectedValue = Setup_.BuildGraph(expectedNode)->GetValue();
auto outDatum = TArrowBlock::From(resultValue).GetDatum();
auto expectedDatum = TArrowBlock::From(expectedValue).GetDatum();
UNIT_ASSERT_EQUAL_C(outDatum, expectedDatum, "Expected : " << DatumToString(expectedDatum) << "\n but got : " << DatumToString(outDatum));
}
private:
TComputationNodeFactory GetNodeTestFactory();
TString DatumToString(arrow::Datum datum);
TRuntimeNode MaterializeBlockStream(TProgramBuilder& pgmBuilder, TRuntimeNode stream);
TSetup<false> Setup_;
TProgramBuilder& Pb_;
};
} // namespace NKikimr::NMiniKQL
|
