Skip to content

Commit d468e6d

Browse files
resetiusresetius
authored
Add decimal comparison kernels (#8271)
1 parent f807bc5 commit d468e6d

File tree

14 files changed

+170
-15
lines changed

14 files changed

+170
-15
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,7 @@ void RegisterEquals(TKernelFamilyMap& kernelFamilyMap) {
302302

303303
AddNumericComparisonKernels<TEqualsOp>(*family);
304304
AddDateComparisonKernels<TDiffDateEqualsOp>(*family);
305+
AddDecimalComparisonKernels<TDecimalEquals>(*family);
305306
RegisterStringKernelEquals(*family);
306307

307308
kernelFamilyMap["Equals"] = std::move(family);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,7 @@ void RegisterGreater(TKernelFamilyMap& kernelFamilyMap) {
299299

300300
AddNumericComparisonKernels<TGreaterOp>(*family);
301301
AddDateComparisonKernels<TDiffDateGreaterOp>(*family);
302+
AddDecimalComparisonKernels<TDecimalGreater>(*family);
302303
RegisterStringKernelGreater(*family);
303304

304305
kernelFamilyMap["Greater"] = std::move(family);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,7 @@ void RegisterGreaterOrEqual(TKernelFamilyMap& kernelFamilyMap) {
299299

300300
AddNumericComparisonKernels<TGreaterOrEqualOp>(*family);
301301
AddDateComparisonKernels<TDiffDateGreaterOrEqualOp>(*family);
302+
AddDecimalComparisonKernels<TDecimalGreaterOrEqual>(*family);
302303
RegisterStringKernelGreaterOrEqual(*family);
303304

304305
kernelFamilyMap["GreaterOrEqual"] = std::move(family);

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

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -296,21 +296,18 @@ std::shared_ptr<arrow::compute::ScalarKernel> TDecimalKernel::MakeArrowKernel(co
296296

297297
MKQL_ENSURE(*dataType1->GetDataSlot() == NUdf::EDataSlot::Decimal, "Require decimal");
298298
MKQL_ENSURE(*dataType2->GetDataSlot() == NUdf::EDataSlot::Decimal, "Require decimal");
299-
MKQL_ENSURE(*dataResultType->GetDataSlot() == NUdf::EDataSlot::Decimal, "Require decimal");
300299

301300
auto decimalType1 = static_cast<TDataDecimalType*>(dataType1);
302301
auto decimalType2 = static_cast<TDataDecimalType*>(dataType2);
303-
auto decimalResultType = static_cast<TDataDecimalType*>(dataResultType);
304302

305303
MKQL_ENSURE(decimalType1->GetParams() == decimalType2->GetParams(), "Require same precision/scale");
306-
MKQL_ENSURE(decimalType1->GetParams() == decimalResultType->GetParams(), "Require same precision/scale");
307304

308305
ui8 precision = decimalType1->GetParams().first;
309306
MKQL_ENSURE(precision >= 1&& precision <= 35, TStringBuilder() << "Wrong precision: " << (int)precision);
310307

311308
auto k = std::make_shared<arrow::compute::ScalarKernel>(std::vector<arrow::compute::InputType>{
312309
GetPrimitiveInputArrowType(NUdf::EDataSlot::Decimal), GetPrimitiveInputArrowType(NUdf::EDataSlot::Decimal)
313-
}, GetPrimitiveOutputArrowType(NUdf::EDataSlot::Decimal), Exec);
310+
}, GetPrimitiveOutputArrowType(*dataResultType->GetDataSlot()), Exec);
314311
k->null_handling = arrow::compute::NullHandling::INTERSECTION;
315312
k->init = [precision](arrow::compute::KernelContext*, const arrow::compute::KernelInitArgs&) {
316313
auto state = std::make_unique<TDecimalKernel::TKernelState>();
@@ -713,7 +710,8 @@ arrow::Status ExecDecimalScalarArrayOptImpl(const arrow::compute::ExecBatch& bat
713710

714711
arrow::Status ExecDecimalScalarScalarOptImpl(arrow::compute::KernelContext* kernelCtx,
715712
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
716-
TPrimitiveDataTypeGetter typeGetter, TUntypedBinaryScalarOptFuncPtr func) {
713+
TPrimitiveDataTypeGetter typeGetter, TPrimitiveDataScalarGetterWithMemPool scalarGetter,
714+
TUntypedBinaryScalarOptFuncPtr func) {
717715
MKQL_ENSURE(batch.values.size() == 2, "Expected 2 args");
718716
const auto& arg1 = batch.values[0];
719717
const auto& arg2 = batch.values[1];
@@ -722,9 +720,9 @@ arrow::Status ExecDecimalScalarScalarOptImpl(arrow::compute::KernelContext* kern
722720
} else {
723721
const auto val1Ptr = GetStringScalarValue(*arg1.scalar());
724722
const auto val2Ptr = GetStringScalarValue(*arg2.scalar());
725-
std::shared_ptr<arrow::Buffer> buffer(ARROW_RESULT(arrow::AllocateBuffer(16, kernelCtx->memory_pool())));
726-
auto resDatum = arrow::Datum(std::make_shared<TPrimitiveDataType<NYql::NDecimal::TInt128>::TScalarResult>(buffer));
727-
if (!func(val1Ptr.data(), val2Ptr.data(), buffer->mutable_data())) {
723+
void* resMem;
724+
auto resDatum = scalarGetter(&resMem, kernelCtx->memory_pool());
725+
if (!func(val1Ptr.data(), val2Ptr.data(), resMem)) {
728726
*res = arrow::MakeNullScalar(typeGetter());
729727
} else {
730728
*res = resDatum.scalar();
@@ -736,7 +734,7 @@ arrow::Status ExecDecimalScalarScalarOptImpl(arrow::compute::KernelContext* kern
736734

737735
arrow::Status ExecDecimalBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
738736
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
739-
TPrimitiveDataTypeGetter typeGetter,
737+
TPrimitiveDataTypeGetter typeGetter, TPrimitiveDataScalarGetterWithMemPool scalarGetter,
740738
size_t outputSizeOf,
741739
TUntypedBinaryScalarOptFuncPtr scalarScalarFunc,
742740
TUntypedBinaryArrayOptFuncPtr scalarArrayFunc,
@@ -747,7 +745,7 @@ arrow::Status ExecDecimalBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
747745
const auto& arg2 = batch.values[1];
748746
if (arg1.is_scalar()) {
749747
if (arg2.is_scalar()) {
750-
return ExecDecimalScalarScalarOptImpl(kernelCtx, batch, res, typeGetter, scalarScalarFunc);
748+
return ExecDecimalScalarScalarOptImpl(kernelCtx, batch, res, typeGetter, scalarGetter, scalarScalarFunc);
751749
} else {
752750
return ExecDecimalScalarArrayOptImpl(batch, res, scalarArrayFunc);
753751
}
@@ -769,8 +767,7 @@ arrow::Status ExecDecimalScalarImpl(arrow::compute::KernelContext* kernelCtx,
769767
const auto valPtr = GetPrimitiveScalarValuePtr(*arg.scalar());
770768
std::shared_ptr<arrow::Buffer> buffer(ARROW_RESULT(arrow::AllocateBuffer(16, kernelCtx->memory_pool())));
771769
auto resDatum = arrow::Datum(std::make_shared<TPrimitiveDataType<NYql::NDecimal::TInt128>::TScalarResult>(buffer));
772-
const auto resPtr = GetPrimitiveScalarValueMutablePtr(*resDatum.scalar());
773-
func(valPtr, resPtr);
770+
func(valPtr, buffer->mutable_data());
774771
*res = resDatum.scalar();
775772
}
776773

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

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -942,6 +942,7 @@ inline std::shared_ptr<arrow::ArrayData> CopyTz(const std::shared_ptr<arrow::Arr
942942

943943
using TPrimitiveDataTypeGetter = std::shared_ptr<arrow::DataType>(*)();
944944
using TPrimitiveDataScalarGetter= arrow::Datum(*)();
945+
using TPrimitiveDataScalarGetterWithMemPool = arrow::Datum(*)(void** result, arrow::MemoryPool*);
945946
using TUntypedBinaryScalarFuncPtr = void(*)(const void*, const void*, void*);
946947
using TUntypedBinaryArrayFuncPtr = void(*)(const void*, const void*, void*, int64_t length, int64_t offset1, int64_t offset2);
947948
using TUntypedBinaryScalarOptFuncPtr = bool(*)(const void*, const void*, void*);
@@ -1023,7 +1024,7 @@ arrow::Status ExecBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
10231024

10241025
arrow::Status ExecDecimalBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
10251026
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
1026-
TPrimitiveDataTypeGetter typeGetter,
1027+
TPrimitiveDataTypeGetter typeGetter, TPrimitiveDataScalarGetterWithMemPool scalarGetter,
10271028
size_t outputSizeOf,
10281029
TUntypedBinaryScalarOptFuncPtr scalarScalarFunc,
10291030
TUntypedBinaryArrayOptFuncPtr scalarArrayFunc,
@@ -1032,7 +1033,7 @@ arrow::Status ExecDecimalBinaryOptImpl(arrow::compute::KernelContext* kernelCtx,
10321033

10331034
arrow::Status ExecDecimalUnaryImpl(arrow::compute::KernelContext* kernelCtx,
10341035
const arrow::compute::ExecBatch& batch, arrow::Datum* res,
1035-
TPrimitiveDataTypeGetter typeGetter,
1036+
TPrimitiveDataTypeGetter typeGetter,
10361037
TUntypedUnaryScalarFuncPtr scalarFunc, TUntypedUnaryArrayFuncPtr arrayFunc);
10371038

10381039
template<typename TInput1, bool Tz1, typename TInput2, bool Tz2, typename TOutput, EPropagateTz PropagateTz,
@@ -1664,6 +1665,12 @@ struct TDecimalKernelExecs
16641665
using TInput2 = NYql::NDecimal::TInt128;
16651666
using TOutput = NYql::NDecimal::TInt128;
16661667

1668+
static arrow::Datum ScalarGetter(void** result, arrow::MemoryPool* memory_pool) {
1669+
std::shared_ptr<arrow::Buffer> buffer(ARROW_RESULT(arrow::AllocateBuffer(16, memory_pool)));
1670+
*result = buffer->mutable_data();
1671+
return arrow::Datum(std::make_shared<TPrimitiveDataType<NYql::NDecimal::TInt128>::TScalarResult>(buffer));
1672+
}
1673+
16671674
template<ui8 precision>
16681675
static arrow::Status ExecImpl(arrow::compute::KernelContext* kernelCtx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
16691676
auto scalarScalarFunc = &TBinaryKernelOptExecsImpl<TInput1, TInput2, TOutput, TFuncInstance<precision>>::ScalarScalarCoreOpt;
@@ -1672,7 +1679,7 @@ struct TDecimalKernelExecs
16721679
auto arrayArrayFunc = &TBinaryKernelOptExecsImpl<TInput1, TInput2, TOutput, TFuncInstance<precision>>::ArrayArrayCoreOpt;
16731680

16741681
return ExecDecimalBinaryOptImpl(kernelCtx, batch, res,
1675-
&GetPrimitiveDataType<TOutput>,
1682+
&GetPrimitiveDataType<TOutput>, &ScalarGetter,
16761683
sizeof(TOutput),
16771684
(TUntypedBinaryScalarOptFuncPtr)scalarScalarFunc,
16781685
(TUntypedBinaryArrayOptFuncPtr)scalarArrayFunc,
@@ -1709,5 +1716,46 @@ void AddBinaryDecimalKernels(TKernelFamilyBase& owner) {
17091716
owner.Adopt(argTypes, returnType, std::move(kernel));
17101717
}
17111718

1719+
template<class TFuncInstance>
1720+
struct TDecimalComparisonKernelExecs
1721+
{
1722+
using TInput1 = NYql::NDecimal::TInt128;
1723+
using TInput2 = NYql::NDecimal::TInt128;
1724+
using TOutput = bool;
1725+
1726+
static arrow::Datum ScalarGetter(void** resMem, arrow::MemoryPool*) {
1727+
auto result = MakeDefaultScalarDatum<TOutput>();
1728+
*resMem = GetPrimitiveScalarValueMutablePtr(*result.scalar());
1729+
return result;
1730+
}
1731+
1732+
static arrow::Status Exec(arrow::compute::KernelContext* kernelCtx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
1733+
auto scalarScalarFunc = &TBinaryKernelOptExecsImpl<TInput1, TInput2, TOutput, TFuncInstance>::ScalarScalarCoreOpt;
1734+
auto scalarArrayFunc = &TBinaryKernelOptExecsImpl<TInput1, TInput2, TOutput, TFuncInstance>::ScalarArrayCoreOpt;
1735+
auto arrayScalarFunc = &TBinaryKernelOptExecsImpl<TInput1, TInput2, TOutput, TFuncInstance>::ArrayScalarCoreOpt;
1736+
auto arrayArrayFunc = &TBinaryKernelOptExecsImpl<TInput1, TInput2, TOutput, TFuncInstance>::ArrayArrayCoreOpt;
1737+
1738+
return ExecDecimalBinaryOptImpl(kernelCtx, batch, res,
1739+
&GetPrimitiveDataType<TOutput>, &ScalarGetter,
1740+
sizeof(TOutput),
1741+
(TUntypedBinaryScalarOptFuncPtr)scalarScalarFunc,
1742+
(TUntypedBinaryArrayOptFuncPtr)scalarArrayFunc,
1743+
(TUntypedBinaryArrayOptFuncPtr)arrayScalarFunc,
1744+
(TUntypedBinaryArrayOptFuncPtr)arrayArrayFunc);
1745+
}
1746+
};
1747+
1748+
template<class TFunc>
1749+
void AddDecimalComparisonKernels(TKernelFamilyBase& owner) {
1750+
auto type1 = NUdf::GetDataTypeInfo(NUdf::EDataSlot::Decimal).TypeId;
1751+
auto type2 = type1;
1752+
auto returnType = NUdf::GetDataTypeInfo(NUdf::EDataSlot::Bool).TypeId;
1753+
std::vector<NUdf::TDataTypeId> argTypes({ type1, type2 });
1754+
1755+
using Execs = TDecimalComparisonKernelExecs<TFunc>;
1756+
auto kernel = std::make_unique<TDecimalKernel>(owner, argTypes, returnType, &Execs::Exec, TKernel::ENullMode::Default);
1757+
owner.Adopt(argTypes, returnType, std::move(kernel));
1758+
}
1759+
17121760
}
17131761
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,7 @@ void RegisterLess(TKernelFamilyMap& kernelFamilyMap) {
299299

300300
AddNumericComparisonKernels<TLessOp>(*family);
301301
AddDateComparisonKernels<TDiffDateLessOp>(*family);
302+
AddDecimalComparisonKernels<TDecimalLess>(*family);
302303
RegisterStringKernelLess(*family);
303304

304305
kernelFamilyMap["Less"] = std::move(family);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,7 @@ void RegisterLessOrEqual(TKernelFamilyMap& kernelFamilyMap) {
299299

300300
AddNumericComparisonKernels<TLessOrEqualOp>(*family);
301301
AddDateComparisonKernels<TDiffDateLessOrEqualOp>(*family);
302+
AddDecimalComparisonKernels<TDecimalLessOrEqual>(*family);
302303
RegisterStringKernelLessOrEqual(*family);
303304

304305
kernelFamilyMap["LessOrEqual"] = std::move(family);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,7 @@ void RegisterNotEquals(TKernelFamilyMap& kernelFamilyMap) {
302302

303303
AddNumericComparisonKernels<TNotEqualsOp>(*family);
304304
AddDateComparisonKernels<TDiffDateNotEqualsOp>(*family);
305+
AddDecimalComparisonKernels<TDecimalNotEquals>(*family);
305306
RegisterStringKernelNotEquals(*family);
306307

307308
kernelFamilyMap["NotEquals"] = std::move(family);

‎ydb/library/yql/tests/sql/dq_file/part18/canondata/result.json

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -570,6 +570,28 @@
570570
}
571571
],
572572
"test.test[blocks-date_sub_scalar--Results]": [],
573+
"test.test[blocks-decimal_comparison--Analyze]": [
574+
{
575+
"checksum": "84e84352daef0d01b5fc6884ec3ebf29",
576+
"size": 3719,
577+
"uri": "https://{canondata_backend}/1924537/081e3ea5ef34a4fe33a8e971e47d53ea3a5151a4/resource.tar.gz#test.test_blocks-decimal_comparison--Analyze_/plan.txt"
578+
}
579+
],
580+
"test.test[blocks-decimal_comparison--Debug]": [
581+
{
582+
"checksum": "050cf26dbe07c6d9cacb0446ed018c90",
583+
"size": 2343,
584+
"uri": "https://{canondata_backend}/1923547/eaec090e24c98f865902b418cc92d70420b05c0b/resource.tar.gz#test.test_blocks-decimal_comparison--Debug_/opt.yql_patched"
585+
}
586+
],
587+
"test.test[blocks-decimal_comparison--Plan]": [
588+
{
589+
"checksum": "84e84352daef0d01b5fc6884ec3ebf29",
590+
"size": 3719,
591+
"uri": "https://{canondata_backend}/1924537/081e3ea5ef34a4fe33a8e971e47d53ea3a5151a4/resource.tar.gz#test.test_blocks-decimal_comparison--Plan_/plan.txt"
592+
}
593+
],
594+
"test.test[blocks-decimal_comparison--Results]": [],
573595
"test.test[blocks-filter_partial_expr--Analyze]": [
574596
{
575597
"checksum": "e8f201d2a8a9bec0b35c5b959be1c92a",

‎ydb/library/yql/tests/sql/hybrid_file/part4/canondata/result.json

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -489,6 +489,20 @@
489489
"uri": "https://{canondata_backend}/1920236/24fc177ad15ffb34261b5dbf2e1c5c96c4ad7faf/resource.tar.gz#test.test_blocks-date_sub_interval_scalar--Plan_/plan.txt"
490490
}
491491
],
492+
"test.test[blocks-decimal_comparison--Debug]": [
493+
{
494+
"checksum": "2d5a263433c13ff6f5d31cefd40bcc96",
495+
"size": 4161,
496+
"uri": "https://{canondata_backend}/1130705/7b1ee13c4da21e2ddff70372802eda23a04da498/resource.tar.gz#test.test_blocks-decimal_comparison--Debug_/opt.yql_patched"
497+
}
498+
],
499+
"test.test[blocks-decimal_comparison--Plan]": [
500+
{
501+
"checksum": "e65c86ba6e55033fe5f47711b65bba6d",
502+
"size": 4093,
503+
"uri": "https://{canondata_backend}/1784826/2b2e157f3cee0db3cfdb7f3fc538c5f30f0b593e/resource.tar.gz#test.test_blocks-decimal_comparison--Plan_/plan.txt"
504+
}
505+
],
492506
"test.test[blocks-distinct_pure_all--Debug]": [
493507
{
494508
"checksum": "a368ba0efe31dc0503d90694682eb845",

0 commit comments

Comments
 (0)