Skip to content

Commit d329420

Browse files
Tony-RomanovTony-Romanov
authored
Merge ad4c0c1 into dbf1e80 
2 parents dbf1e80 + ad4c0c1 commit d329420

14 files changed

+170
-109
lines changed

‎ydb/library/yql/core/arrow_kernels/registry/ut/registry_ut.cpp

+63-8
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,57 @@ Y_UNIT_TEST_SUITE(TKernelRegistryTest) {
112112
});
113113
}
114114

115+
Y_UNIT_TEST(TestAddSubMulOps) {
116+
for (const auto oper : {TKernelRequestBuilder::EBinaryOp::Add, TKernelRequestBuilder::EBinaryOp::Sub, TKernelRequestBuilder::EBinaryOp::Mul}) {
117+
for (const auto slot : {EDataSlot::Int8, EDataSlot::Int16, EDataSlot::Int32, EDataSlot::Int64, EDataSlot::Uint8, EDataSlot::Uint16, EDataSlot::Uint32, EDataSlot::Uint64, EDataSlot::Float, EDataSlot::Double}) {
118+
TestOne([slot, oper](auto& b,auto& ctx) {
119+
const auto blockUint8Type = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::Uint8));
120+
const auto blockType = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(slot));
121+
return b.AddBinaryOp(oper, blockUint8Type, blockType, blockType);
122+
});
123+
TestOne([slot, oper](auto& b,auto& ctx) {
124+
const auto blockUint8Type = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::Uint8));
125+
const auto blockType = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(slot));
126+
return b.AddBinaryOp(oper, blockType, blockUint8Type, blockType);
127+
});
128+
TestOne([slot, oper](auto& b,auto& ctx) {
129+
const auto blockType = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(slot));
130+
return b.AddBinaryOp(oper, blockType, blockType, blockType);
131+
});
132+
}
133+
}
134+
}
135+
136+
Y_UNIT_TEST(TestDivModOps) {
137+
for (const auto oper : {TKernelRequestBuilder::EBinaryOp::Div, TKernelRequestBuilder::EBinaryOp::Mod}) {
138+
for (const auto slot : {EDataSlot::Int8, EDataSlot::Int16, EDataSlot::Int32, EDataSlot::Int64, EDataSlot::Uint8, EDataSlot::Uint16, EDataSlot::Uint32, EDataSlot::Uint64, EDataSlot::Float, EDataSlot::Double}) {
139+
TestOne([slot, oper](auto& b,auto& ctx) {
140+
const auto blockUint8Type = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::Uint8));
141+
const auto rawType = ctx.template MakeType<TDataExprType>(slot);
142+
const auto blockType = ctx.template MakeType<TBlockExprType>(rawType);
143+
const auto returnType = EDataSlot::Float != slot && EDataSlot::Double != slot ?
144+
ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TOptionalExprType>(rawType)) : blockType;
145+
return b.AddBinaryOp(oper, blockUint8Type, blockType, returnType);
146+
});
147+
TestOne([slot, oper](auto& b,auto& ctx) {
148+
const auto blockUint8Type = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::Uint8));
149+
const auto rawType = ctx.template MakeType<TDataExprType>(slot);
150+
const auto blockType = ctx.template MakeType<TBlockExprType>(rawType);
151+
const auto returnType = EDataSlot::Float != slot && EDataSlot::Double != slot ?
152+
ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TOptionalExprType>(rawType)) : blockType;
153+
return b.AddBinaryOp(oper, blockType, blockUint8Type, returnType);
154+
});
155+
TestOne([slot, oper](auto& b,auto& ctx) {
156+
const auto rawType = ctx.template MakeType<TDataExprType>(slot);
157+
const auto blockType = ctx.template MakeType<TBlockExprType>(rawType);
158+
const auto returnType = EDataSlot::Float != slot && EDataSlot::Double != slot ?
159+
ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TOptionalExprType>(rawType)) : blockType;
160+
return b.AddBinaryOp(oper, blockType, blockType, returnType);
161+
});
162+
}
163+
}
164+
}
165+
115166
Y_UNIT_TEST(TestSize) {
116167
TestOne([](auto& b,auto& ctx) {
117168
auto blockStrType = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::String));
@@ -121,17 +172,21 @@ Y_UNIT_TEST_SUITE(TKernelRegistryTest) {
121172
}
122173

123174
Y_UNIT_TEST(TestMinus) {
124-
TestOne([](auto& b,auto& ctx) {
125-
auto blockInt32Type = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::Int32));
126-
return b.AddUnaryOp(TKernelRequestBuilder::EUnaryOp::Minus, blockInt32Type, blockInt32Type);
127-
});
175+
for (const auto slot : {EDataSlot::Int8, EDataSlot::Int16, EDataSlot::Int32, EDataSlot::Int64, EDataSlot::Float, EDataSlot::Double}) {
176+
TestOne([slot](auto& b,auto& ctx) {
177+
const auto blockType = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(slot));
178+
return b.AddUnaryOp(TKernelRequestBuilder::EUnaryOp::Minus, blockType, blockType);
179+
});
180+
}
128181
}
129182

130183
Y_UNIT_TEST(TestAbs) {
131-
TestOne([](auto& b,auto& ctx) {
132-
auto blockInt32Type = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(EDataSlot::Int32));
133-
return b.AddUnaryOp(TKernelRequestBuilder::EUnaryOp::Abs, blockInt32Type, blockInt32Type);
134-
});
184+
for (const auto slot : {EDataSlot::Int8, EDataSlot::Int16, EDataSlot::Int32, EDataSlot::Int64, EDataSlot::Float, EDataSlot::Double}) {
185+
TestOne([slot](auto& b,auto& ctx) {
186+
const auto blockType = ctx.template MakeType<TBlockExprType>(ctx.template MakeType<TDataExprType>(slot));
187+
return b.AddUnaryOp(TKernelRequestBuilder::EUnaryOp::Abs, blockType, blockType);
188+
});
189+
}
135190
}
136191

137192
Y_UNIT_TEST(TestCoalesece) {

‎ydb/library/yql/minikql/arrow/mkql_functions.cpp

+4-4
Original file line numberDiff line numberDiff line change
@@ -160,14 +160,14 @@ bool FindArrowFunction(TStringBuf name, const TArrayRef<TType*>& inputTypes, TTy
160160
}
161161

162162
bool match = false;
163-
switch (kernel->Family.NullMode) {
164-
case TKernelFamily::ENullMode::Default:
163+
switch (kernel->NullMode) {
164+
case TKernel::ENullMode::Default:
165165
match = returnIsOptional == hasOptionals;
166166
break;
167-
case TKernelFamily::ENullMode::AlwaysNull:
167+
case TKernel::ENullMode::AlwaysNull:
168168
match = returnIsOptional;
169169
break;
170-
case TKernelFamily::ENullMode::AlwaysNotNull:
170+
case TKernel::ENullMode::AlwaysNotNull:
171171
match = !returnIsOptional;
172172
break;
173173
}

‎ydb/library/yql/minikql/invoke_builtins/mkql_builtins.cpp

-59
Original file line numberDiff line numberDiff line change
@@ -19,65 +19,6 @@ namespace NMiniKQL {
1919

2020
namespace {
2121

22-
class TForeignKernel : public TKernel {
23-
public:
24-
TForeignKernel(const TKernelFamily& family, const std::vector<NUdf::TDataTypeId>& argTypes, NUdf::TDataTypeId returnType,
25-
const std::shared_ptr<arrow::compute::Function>& function)
26-
: TKernel(family, argTypes, returnType)
27-
, Function(function)
28-
, ArrowKernel(ResolveKernel(Function, argTypes))
29-
{}
30-
31-
const arrow::compute::ScalarKernel& GetArrowKernel() const final {
32-
return ArrowKernel;
33-
}
34-
35-
private:
36-
static const arrow::compute::ScalarKernel& ResolveKernel(const std::shared_ptr<arrow::compute::Function>& function,
37-
const std::vector<NUdf::TDataTypeId>& argTypes) {
38-
std::vector<arrow::ValueDescr> args;
39-
for (const auto& t : argTypes) {
40-
args.emplace_back();
41-
auto slot = NUdf::FindDataSlot(t);
42-
MKQL_ENSURE(slot, "Unexpected data type");
43-
MKQL_ENSURE(ConvertArrowType(*slot, args.back().type), "Can't get arrow type");
44-
}
45-
46-
const auto kernel = ARROW_RESULT(function->DispatchExact(args));
47-
return *static_cast<const arrow::compute::ScalarKernel*>(kernel);
48-
}
49-
50-
private:
51-
const std::shared_ptr<arrow::compute::Function> Function;
52-
const arrow::compute::ScalarKernel& ArrowKernel;
53-
};
54-
55-
template <typename TInput1, typename TOutput>
56-
void RegisterUnary(const arrow::compute::FunctionRegistry& registry, std::string_view name, TKernelFamilyMap& kernelFamilyMap) {
57-
auto func = ARROW_RESULT(registry.GetFunction(std::string(name)));
58-
59-
std::vector<NUdf::TDataTypeId> argTypes({ NUdf::TDataType<TInput1>::Id });
60-
NUdf::TDataTypeId returnType = NUdf::TDataType<TOutput>::Id;
61-
62-
auto family = std::make_unique<TKernelFamilyBase>();
63-
family->Adopt(argTypes, returnType, std::make_unique<TForeignKernel>(*family, argTypes, returnType, func));
64-
65-
Y_ENSURE(kernelFamilyMap.emplace(TString(name), std::move(family)).second);
66-
}
67-
68-
template <typename TInput1, typename TInput2, typename TOutput>
69-
void RegisterBinary(const arrow::compute::FunctionRegistry& registry, std::string_view name, TKernelFamilyMap& kernelFamilyMap) {
70-
auto func = ARROW_RESULT(registry.GetFunction(std::string(name)));
71-
72-
std::vector<NUdf::TDataTypeId> argTypes({ NUdf::TDataType<TInput1>::Id, NUdf::TDataType<TInput2>::Id });
73-
NUdf::TDataTypeId returnType = NUdf::TDataType<TOutput>::Id;
74-
75-
auto family = std::make_unique<TKernelFamilyBase>();
76-
family->Adopt(argTypes, returnType, std::make_unique<TForeignKernel>(*family, argTypes, returnType, func));
77-
78-
Y_ENSURE(kernelFamilyMap.emplace(TString(name), std::move(family)).second);
79-
}
80-
8122
void RegisterDefaultOperations(IBuiltinFunctionRegistry& registry, TKernelFamilyMap& kernelFamilyMap) {
8223
RegisterAdd(registry);
8324
RegisterAdd(kernelFamilyMap);

‎ydb/library/yql/minikql/invoke_builtins/mkql_builtins_abs.cpp

+2
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ inline T Abs(T v) {
2323

2424
template<typename TInput, typename TOutput>
2525
struct TAbs : public TSimpleArithmeticUnary<TInput, TOutput, TAbs<TInput, TOutput>> {
26+
static constexpr bool DefaultNulls = true;
27+
2628
static TOutput Do(TInput val)
2729
{
2830
return Abs<TInput>(val);

‎ydb/library/yql/minikql/invoke_builtins/mkql_builtins_add.cpp

+1-1
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ void RegisterAdd(IBuiltinFunctionRegistry& registry) {
193193
}
194194

195195
void RegisterAdd(TKernelFamilyMap& kernelFamilyMap) {
196-
kernelFamilyMap["Add"] = std::make_unique<TBinaryNumericKernelFamily<TAdd>>();
196+
kernelFamilyMap["Add"] = std::make_unique<TBinaryNumericKernelFamily<TAdd, TAdd>>();
197197
}
198198

199199
void RegisterAggrAdd(IBuiltinFunctionRegistry& registry) {

‎ydb/library/yql/minikql/invoke_builtins/mkql_builtins_div.cpp

+4-2
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ template<typename TLeft, typename TRight, typename TOutput>
1212
struct TDiv : public TSimpleArithmeticBinary<TLeft, TRight, TOutput, TDiv<TLeft, TRight, TOutput>> {
1313
static_assert(std::is_floating_point<TOutput>::value, "expected floating point");
1414

15+
static constexpr bool DefaultNulls = true;
16+
1517
static TOutput Do(TOutput left, TOutput right)
1618
{
1719
return left / right;
@@ -60,7 +62,7 @@ struct TIntegralDiv {
6062
const auto result = PHINode::Create(type, 2, "result", done);
6163
result->addIncoming(zero, block);
6264

63-
if (std::is_signed<TOutput>() && sizeof(TOutput) <= sizeof(TLeft)) {
65+
if constexpr (std::is_signed<TOutput>() && sizeof(TOutput) <= sizeof(TLeft)) {
6466
const auto min = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, lv, ConstantInt::get(lv->getType(), Min<TOutput>()), "min", block);
6567
const auto one = CmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, rv, ConstantInt::get(rv->getType(), -1), "one", block);
6668
const auto two = BinaryOperator::CreateAnd(min, one, "two", block);
@@ -167,7 +169,7 @@ void RegisterDiv(IBuiltinFunctionRegistry& registry) {
167169
}
168170

169171
void RegisterDiv(TKernelFamilyMap& kernelFamilyMap) {
170-
kernelFamilyMap["Div"] = std::make_unique<TBinaryNumericKernelFamily<TIntegralDiv>>(TKernelFamily::ENullMode::AlwaysNull);
172+
kernelFamilyMap["Div"] = std::make_unique<TBinaryNumericKernelFamily<TIntegralDiv, TDiv>>();
171173
}
172174

173175
} // namespace NMiniKQL

0 commit comments

Comments
 (0)