Skip to content

Commit dd45bbe

Browse files
authored
Add decimal unary kernels (#8254)
1 parent 68c746e commit dd45bbe

File tree

10 files changed

+176
-5
lines changed

10 files changed

+176
-5
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ struct TAbs : public TSimpleArithmeticUnary<TInput, TOutput, TAbs<TInput, TOutpu
5454
#endif
5555
};
5656

57-
struct TDecimalAbs {
57+
struct TDecimalAbs : public TDecimalUnary<TDecimalAbs> {
5858
static NUdf::TUnboxedValuePod Execute(const NUdf::TUnboxedValuePod& arg) {
5959
const auto a = arg.GetInt128();
6060
return a < 0 ? NUdf::TUnboxedValuePod(-a) : arg;
@@ -83,7 +83,9 @@ void RegisterAbs(IBuiltinFunctionRegistry& registry) {
8383
}
8484

8585
void RegisterAbs(TKernelFamilyMap& kernelFamilyMap) {
86-
kernelFamilyMap["Abs"] = std::make_unique<TUnaryNumericKernelFamily<TAbs>>();
86+
auto family = std::make_unique<TUnaryNumericKernelFamily<TAbs>>();
87+
AddUnaryDecimalKernels<TDecimalAbs>(*family);
88+
kernelFamilyMap["Abs"] = std::move(family);
8789
}
8890

8991
} // namespace NMiniKQL

ydb/library/yql/minikql/invoke_builtins/mkql_builtins_impl.cpp

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -760,5 +760,48 @@ arrow::Status ExecDecimalBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
760760
}
761761
}
762762

763+
arrow::Status ExecDecimalScalarImpl(arrow::compute::KernelContext* kernelCtx,
764+
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
765+
TPrimitiveDataTypeGetter typeGetter, TUntypedUnaryScalarFuncPtr func) {
766+
if (const auto& arg = batch.values.front(); !arg.scalar()->is_valid) {
767+
*res = arrow::MakeNullScalar(typeGetter());
768+
} else {
769+
const auto valPtr = GetPrimitiveScalarValuePtr(*arg.scalar());
770+
std::shared_ptr<arrow::Buffer> buffer(ARROW_RESULT(arrow::AllocateBuffer(16, kernelCtx->memory_pool())));
771+
auto resDatum = arrow::Datum(std::make_shared<TPrimitiveDataType<NYql::NDecimal::TInt128>::TScalarResult>(buffer));
772+
const auto resPtr = GetPrimitiveScalarValueMutablePtr(*resDatum.scalar());
773+
func(valPtr, resPtr);
774+
*res = resDatum.scalar();
775+
}
776+
777+
return arrow::Status::OK();
778+
}
779+
780+
arrow::Status ExecDecimalArrayImpl(const arrow::compute::ExecBatch& batch, arrow::Datum* res,
781+
TUntypedUnaryArrayFuncPtr func) {
782+
const auto& arg = batch.values.front();
783+
auto& resArr = *res->array();
784+
785+
const auto& arr = *arg.array();
786+
auto length = arr.length;
787+
const auto valPtr = arr.buffers[1]->data();
788+
auto resPtr = resArr.buffers[1]->mutable_data();
789+
func(valPtr, resPtr, length, arr.offset);
790+
return arrow::Status::OK();
791+
}
792+
793+
arrow::Status ExecDecimalUnaryImpl(arrow::compute::KernelContext* kernelCtx,
794+
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
795+
TPrimitiveDataTypeGetter typeGetter,
796+
TUntypedUnaryScalarFuncPtr scalarFunc, TUntypedUnaryArrayFuncPtr arrayFunc) {
797+
MKQL_ENSURE(batch.values.size() == 1, "Expected single argument");
798+
const auto& arg = batch.values[0];
799+
if (arg.is_scalar()) {
800+
return ExecDecimalScalarImpl(kernelCtx, batch, res, typeGetter, scalarFunc);
801+
} else {
802+
return ExecDecimalArrayImpl(batch, res, arrayFunc);
803+
}
804+
}
805+
763806
} // namespace NMiniKQL
764807
} // namespace NKikimr

ydb/library/yql/minikql/invoke_builtins/mkql_builtins_impl.h.txt

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,15 @@ struct TSimpleArithmeticUnary : public TArithmeticConstraintsSame<TInput, TOutpu
150150
#endif
151151
};
152152

153+
template <class TImpl>
154+
struct TDecimalUnary {
155+
static void DoPtr(
156+
const NYql::NDecimal::TInt128* arg,
157+
NYql::NDecimal::TInt128* res) {
158+
*res = TImpl::Execute(NUdf::TUnboxedValuePod(*arg)).GetInt128();
159+
}
160+
};
161+
153162
template <typename TLeft, typename TRight, typename TOutput, class TImpl, bool CustomCast = false>
154163
struct TSimpleArithmeticBinary : public TArithmeticConstraintsBinary<TLeft, TRight, TOutput> {
155164
static NUdf::TUnboxedValuePod Execute(const NUdf::TUnboxedValuePod& left, const NUdf::TUnboxedValuePod& right) {
@@ -1021,6 +1030,11 @@ arrow::Status ExecDecimalBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
10211030
TUntypedBinaryArrayOptFuncPtr arrayScalarFunc,
10221031
TUntypedBinaryArrayOptFuncPtr arrayArrayFunc);
10231032

1033+
arrow::Status ExecDecimalUnaryImpl(arrow::compute::KernelContext* kernelCtx,
1034+
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
1035+
TPrimitiveDataTypeGetter typeGetter,
1036+
TUntypedUnaryScalarFuncPtr scalarFunc, TUntypedUnaryArrayFuncPtr arrayFunc);
1037+
10241038
template<typename TInput1, bool Tz1, typename TInput2, bool Tz2, typename TOutput, EPropagateTz PropagateTz,
10251039
typename TFuncInstance, TKernel::ENullMode NullMode>
10261040
struct TBinaryKernelExecs;
@@ -1259,7 +1273,7 @@ private:
12591273
template<typename TInput, bool Tz, typename TOutput, bool PropagateTz, class TFuncInstance>
12601274
struct TUnaryKernelExecs
12611275
{
1262-
using TTypedUnaryScalarFuncPtr = void(*)(
1276+
using TTypedUnaryScalarFuncPtr = void(*)(
12631277
const typename TPrimitiveDataType<TInput>::TLayout*,
12641278
typename TPrimitiveDataType<TOutput>::TLayout*
12651279
);
@@ -1296,6 +1310,23 @@ struct TUnaryKernelExecs
12961310
}
12971311
};
12981312

1313+
1314+
template<class TFuncInstance>
1315+
struct TUnaryDecimalKernelExecs
1316+
{
1317+
using TInput = NYql::NDecimal::TInt128;
1318+
using TOutput = NYql::NDecimal::TInt128;
1319+
1320+
static arrow::Status Exec(arrow::compute::KernelContext* kernelCtx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
1321+
auto func = &TFuncInstance::DoPtr;
1322+
auto arrayFunc = &TUnaryKernelExecs<TInput, false, TOutput, false, TFuncInstance>::ArrayCore;
1323+
return ExecDecimalUnaryImpl(kernelCtx, batch, res,
1324+
&GetPrimitiveDataType<TOutput>,
1325+
(TUntypedUnaryScalarFuncPtr)func,
1326+
(TUntypedUnaryArrayFuncPtr)arrayFunc);
1327+
}
1328+
};
1329+
12991330
using TStatelessArrayKernelExec = arrow::Status(*)(arrow::compute::KernelContext*, const arrow::compute::ExecBatch&, arrow::Datum*);
13001331

13011332
void AddUnaryKernelImpl(TKernelFamilyBase& owner, NUdf::EDataSlot arg1, NUdf::EDataSlot res,
@@ -1316,6 +1347,13 @@ void AddUnaryKernel(TKernelFamilyBase& owner) {
13161347
AddUnaryKernelImpl(owner, TInput::Slot, TOutput::Slot, &TExecs::Exec, TFuncInstance::NullMode);
13171348
}
13181349

1350+
template<class TFunc>
1351+
void AddUnaryDecimalKernels(TKernelFamilyBase& owner) {
1352+
using TExecs = TUnaryDecimalKernelExecs<TFunc>;
1353+
1354+
AddUnaryKernelImpl(owner, NUdf::EDataSlot::Decimal, NUdf::EDataSlot::Decimal, &TExecs::Exec, TKernel::ENullMode::Default);
1355+
}
1356+
13191357
void AddBinaryKernelImpl(TKernelFamilyBase& owner, NUdf::EDataSlot arg1, NUdf::EDataSlot arg2, NUdf::EDataSlot res,
13201358
TStatelessArrayKernelExec exec, TKernel::ENullMode nullMode);
13211359

ydb/library/yql/minikql/invoke_builtins/mkql_builtins_minus.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ struct TMinus : public TSimpleArithmeticUnary<TInput, TOutput, TMinus<TInput, TO
2525
#endif
2626
};
2727

28-
struct TDecimalMinus {
28+
struct TDecimalMinus: TDecimalUnary<TDecimalMinus> {
2929
static NUdf::TUnboxedValuePod Execute(const NUdf::TUnboxedValuePod& arg) {
3030
const auto v = arg.GetInt128();
3131
return NYql::NDecimal::IsComparable(v) ? NUdf::TUnboxedValuePod(-v) : arg;
@@ -52,7 +52,9 @@ void RegisterMinus(IBuiltinFunctionRegistry& registry) {
5252
}
5353

5454
void RegisterMinus(TKernelFamilyMap& kernelFamilyMap) {
55-
kernelFamilyMap["Minus"] = std::make_unique<TUnaryNumericKernelFamily<TMinus>>();
55+
auto family = std::make_unique<TUnaryNumericKernelFamily<TMinus>>();
56+
AddUnaryDecimalKernels<TDecimalMinus>(*family);
57+
kernelFamilyMap["Minus"] = std::move(family);
5658
}
5759

5860
} // namespace NMiniKQL

ydb/library/yql/tests/sql/dq_file/part7/canondata/result.json

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -534,6 +534,28 @@
534534
}
535535
],
536536
"test.test[blocks-date_sub_interval_scalar--Results]": [],
537+
"test.test[blocks-decimal_unary--Analyze]": [
538+
{
539+
"checksum": "84e84352daef0d01b5fc6884ec3ebf29",
540+
"size": 3719,
541+
"uri": "https://{canondata_backend}/1775319/1b7c1e5298ad827e3c0e08d1d3f96ba4f42d8217/resource.tar.gz#test.test_blocks-decimal_unary--Analyze_/plan.txt"
542+
}
543+
],
544+
"test.test[blocks-decimal_unary--Debug]": [
545+
{
546+
"checksum": "b04ec7f3b1dfbc6dc36a41426b29bf3e",
547+
"size": 1586,
548+
"uri": "https://{canondata_backend}/1775319/1b7c1e5298ad827e3c0e08d1d3f96ba4f42d8217/resource.tar.gz#test.test_blocks-decimal_unary--Debug_/opt.yql_patched"
549+
}
550+
],
551+
"test.test[blocks-decimal_unary--Plan]": [
552+
{
553+
"checksum": "84e84352daef0d01b5fc6884ec3ebf29",
554+
"size": 3719,
555+
"uri": "https://{canondata_backend}/1775319/1b7c1e5298ad827e3c0e08d1d3f96ba4f42d8217/resource.tar.gz#test.test_blocks-decimal_unary--Plan_/plan.txt"
556+
}
557+
],
558+
"test.test[blocks-decimal_unary--Results]": [],
537559
"test.test[blocks-interval_add_interval--Analyze]": [
538560
{
539561
"checksum": "9c2458b2facc126f97eff00c30912d08",

ydb/library/yql/tests/sql/hybrid_file/part1/canondata/result.json

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -601,6 +601,20 @@
601601
"uri": "https://{canondata_backend}/1942525/c5f89e47e168e72c1372db8b5d6d6ab7387f91cf/resource.tar.gz#test.test_blocks-date_less_or_equal--Plan_/plan.txt"
602602
}
603603
],
604+
"test.test[blocks-decimal_unary--Debug]": [
605+
{
606+
"checksum": "c6286d593f3ce2581337dee33cd574ff",
607+
"size": 2241,
608+
"uri": "https://{canondata_backend}/1775319/b7d43990cdf7ceef10e2a27dcd97d11ed7c5fde7/resource.tar.gz#test.test_blocks-decimal_unary--Debug_/opt.yql_patched"
609+
}
610+
],
611+
"test.test[blocks-decimal_unary--Plan]": [
612+
{
613+
"checksum": "e65c86ba6e55033fe5f47711b65bba6d",
614+
"size": 4093,
615+
"uri": "https://{canondata_backend}/1775319/b7d43990cdf7ceef10e2a27dcd97d11ed7c5fde7/resource.tar.gz#test.test_blocks-decimal_unary--Plan_/plan.txt"
616+
}
617+
],
604618
"test.test[blocks-distinct_mixed_keys--Debug]": [
605619
{
606620
"checksum": "8a1c4eaf681e49661dbe802d96a52122",

ydb/library/yql/tests/sql/sql2yql/canondata/result.json

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3793,6 +3793,13 @@
37933793
"uri": "https://{canondata_backend}/1916746/24ebd7e8cc87e236525be92f05f3986e32fcda57/resource.tar.gz#test_sql2yql.test_blocks-decimal_op_decimal_scalar_/sql.yql"
37943794
}
37953795
],
3796+
"test_sql2yql.test[blocks-decimal_unary]": [
3797+
{
3798+
"checksum": "0b6acc0b2de96c09957216267b9b2c12",
3799+
"size": 1323,
3800+
"uri": "https://{canondata_backend}/1775319/787a3a5920e3b54b1be55e16dd1267b5dff04052/resource.tar.gz#test_sql2yql.test_blocks-decimal_unary_/sql.yql"
3801+
}
3802+
],
37963803
"test_sql2yql.test[blocks-distinct_mixed_all]": [
37973804
{
37983805
"checksum": "079fb36cbdcf679cc0b27308369a2109",
@@ -23323,6 +23330,13 @@
2332323330
"uri": "https://{canondata_backend}/1916746/24ebd7e8cc87e236525be92f05f3986e32fcda57/resource.tar.gz#test_sql_format.test_blocks-decimal_op_decimal_scalar_/formatted.sql"
2332423331
}
2332523332
],
23333+
"test_sql_format.test[blocks-decimal_unary]": [
23334+
{
23335+
"checksum": "39c91354e139d6f6125a1291a75ba269",
23336+
"size": 76,
23337+
"uri": "https://{canondata_backend}/1775319/787a3a5920e3b54b1be55e16dd1267b5dff04052/resource.tar.gz#test_sql_format.test_blocks-decimal_unary_/formatted.sql"
23338+
}
23339+
],
2332623340
"test_sql_format.test[blocks-distinct_mixed_all]": [
2332723341
{
2332823342
"checksum": "97c994cd539ad7338eccbf32126c0e46",
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
in Input input_decimal.txt
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
USE plato;
2+
3+
SELECT
4+
-cs_ext_list_price,
5+
abs(cs_ext_tax),
6+
FROM Input;
7+

ydb/library/yql/tests/sql/yt_native_file/part7/canondata/result.json

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -539,6 +539,34 @@
539539
"uri": "https://{canondata_backend}/1889210/bd004525e88f1c6e39560bd470cf49f93d1e64d3/resource.tar.gz#test.test_blocks-date_sub_interval_scalar--Results_/results.txt"
540540
}
541541
],
542+
"test.test[blocks-decimal_unary--Debug]": [
543+
{
544+
"checksum": "e98c063457aae4bfa904c3e9dfe9bde4",
545+
"size": 1641,
546+
"uri": "https://{canondata_backend}/1937424/c545560a57a71f4626711573be536163931d255d/resource.tar.gz#test.test_blocks-decimal_unary--Debug_/opt.yql"
547+
}
548+
],
549+
"test.test[blocks-decimal_unary--Peephole]": [
550+
{
551+
"checksum": "b5e004bc82ab1eaa68ca168f6d7d3b7e",
552+
"size": 1635,
553+
"uri": "https://{canondata_backend}/1937424/c545560a57a71f4626711573be536163931d255d/resource.tar.gz#test.test_blocks-decimal_unary--Peephole_/opt.yql"
554+
}
555+
],
556+
"test.test[blocks-decimal_unary--Plan]": [
557+
{
558+
"checksum": "4578f9a1b54bf43aa6ebe0c94d50b4c1",
559+
"size": 4059,
560+
"uri": "https://{canondata_backend}/1937424/c545560a57a71f4626711573be536163931d255d/resource.tar.gz#test.test_blocks-decimal_unary--Plan_/plan.txt"
561+
}
562+
],
563+
"test.test[blocks-decimal_unary--Results]": [
564+
{
565+
"checksum": "f4a5b7a689e94b666f2deb2b367ce4e5",
566+
"size": 3586,
567+
"uri": "https://{canondata_backend}/1937424/c545560a57a71f4626711573be536163931d255d/resource.tar.gz#test.test_blocks-decimal_unary--Results_/results.txt"
568+
}
569+
],
542570
"test.test[blocks-interval_add_interval--Debug]": [
543571
{
544572
"checksum": "e53cd73632ad1097239089830422d5a3",

0 commit comments

Comments
 (0)