diff --git a/ydb/library/yql/core/arrow_kernels/registry/ut/registry_ut.cpp b/ydb/library/yql/core/arrow_kernels/registry/ut/registry_ut.cpp index 670b3a8020b9..f66c4dbe97a1 100644 --- a/ydb/library/yql/core/arrow_kernels/registry/ut/registry_ut.cpp +++ b/ydb/library/yql/core/arrow_kernels/registry/ut/registry_ut.cpp @@ -120,6 +120,20 @@ Y_UNIT_TEST_SUITE(TKernelRegistryTest) { }); } + Y_UNIT_TEST(TestMinus) { + TestOne([](auto& b,auto& ctx) { + auto blockInt32Type = ctx.template MakeType(ctx.template MakeType(EDataSlot::Int32)); + return b.AddUnaryOp(TKernelRequestBuilder::EUnaryOp::Minus, blockInt32Type, blockInt32Type); + }); + } + + Y_UNIT_TEST(TestAbs) { + TestOne([](auto& b,auto& ctx) { + auto blockInt32Type = ctx.template MakeType(ctx.template MakeType(EDataSlot::Int32)); + return b.AddUnaryOp(TKernelRequestBuilder::EUnaryOp::Abs, blockInt32Type, blockInt32Type); + }); + } + Y_UNIT_TEST(TestCoalesece) { TestOne([](auto& b,auto& ctx) { auto blockStringType = ctx.template MakeType(ctx.template MakeType(EDataSlot::String)); diff --git a/ydb/library/yql/core/arrow_kernels/request/request.cpp b/ydb/library/yql/core/arrow_kernels/request/request.cpp index fa1b7f6f3306..0099b3fb80e9 100644 --- a/ydb/library/yql/core/arrow_kernels/request/request.cpp +++ b/ydb/library/yql/core/arrow_kernels/request/request.cpp @@ -29,6 +29,8 @@ ui32 TKernelRequestBuilder::AddUnaryOp(EUnaryOp op, const TTypeAnnotationNode* a Items_.emplace_back(Pb_.BlockNot(arg)); break; case EUnaryOp::Size: + case EUnaryOp::Minus: + case EUnaryOp::Abs: Items_.emplace_back(Pb_.BlockFunc(ToString(op), returnType, { arg })); break; } diff --git a/ydb/library/yql/core/arrow_kernels/request/request.h b/ydb/library/yql/core/arrow_kernels/request/request.h index cd84d3db467d..c6398b459106 100644 --- a/ydb/library/yql/core/arrow_kernels/request/request.h +++ b/ydb/library/yql/core/arrow_kernels/request/request.h @@ -11,6 +11,8 @@ class TKernelRequestBuilder { enum class EUnaryOp { Not, Size, + Minus, + Abs }; enum class EBinaryOp { diff --git a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json index 63d5117c782b..1c65ddc67b21 100644 --- a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json +++ b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json @@ -1285,6 +1285,25 @@ "VarArgBase": "TExprBase", "Match": {"Type": "Callable", "Name": "Xor"} }, + { + "Name": "TCoUnaryArithmetic", + "Base": "TCallable", + "Match": {"Type": "CallableBase"}, + "Builder": {"Generate": "None"}, + "Children": [ + {"Index": 0, "Name": "Arg", "Type": "TExprBase"} + ] + }, + { + "Name": "TCoMinus", + "Base": "TCoUnaryArithmetic", + "Match": {"Type": "Callable", "Name": "Minus"} + }, + { + "Name": "TCoAbs", + "Base": "TCoUnaryArithmetic", + "Match": {"Type": "Callable", "Name": "Abs"} + }, { "Name": "TCoBinaryArithmetic", "Base": "TCallable", diff --git a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins.cpp b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins.cpp index 987f6c282d0f..afa2ff99c28c 100644 --- a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins.cpp +++ b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins.cpp @@ -101,9 +101,11 @@ void RegisterDefaultOperations(IBuiltinFunctionRegistry& registry, TKernelFamily RegisterRotRight(registry); RegisterPlus(registry); RegisterMinus(registry); + RegisterMinus(kernelFamilyMap); RegisterBitNot(registry); RegisterCountBits(registry); RegisterAbs(registry); + RegisterAbs(kernelFamilyMap); RegisterConvert(registry); RegisterConcat(registry); RegisterSubstring(registry); diff --git a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_abs.cpp b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_abs.cpp index e5d7731fc2b5..6be9662115be 100644 --- a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_abs.cpp +++ b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_abs.cpp @@ -79,5 +79,9 @@ void RegisterAbs(IBuiltinFunctionRegistry& registry) { NDecimal::RegisterUnaryFunction(registry, "Abs"); } +void RegisterAbs(TKernelFamilyMap& kernelFamilyMap) { + kernelFamilyMap["Abs"] = std::make_unique>(); +} + } // namespace NMiniKQL } // namespace NKikimr diff --git a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_impl.h.txt b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_impl.h.txt index 9f505aee3d3f..b6ab6fcfe685 100644 --- a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_impl.h.txt +++ b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_impl.h.txt @@ -800,9 +800,11 @@ void RegisterRotLeft(IBuiltinFunctionRegistry& registry); void RegisterRotRight(IBuiltinFunctionRegistry& registry); void RegisterPlus(IBuiltinFunctionRegistry& registry); void RegisterMinus(IBuiltinFunctionRegistry& registry); +void RegisterMinus(TKernelFamilyMap& kernelFamilyMap); void RegisterBitNot(IBuiltinFunctionRegistry& registry); void RegisterCountBits(IBuiltinFunctionRegistry& registry); void RegisterAbs(IBuiltinFunctionRegistry& registry); +void RegisterAbs(TKernelFamilyMap& kernelFamilyMap); void RegisterConvert(IBuiltinFunctionRegistry& registry); void RegisterConcat(IBuiltinFunctionRegistry& registry); void RegisterSubstring(IBuiltinFunctionRegistry& registry); @@ -1087,6 +1089,54 @@ private: const arrow::compute::ScalarKernel ArrowKernel; }; +template +struct TUnaryKernelExecs : TUnaryKernelExecsBase> +{ + static arrow::Status ExecScalar(arrow::compute::KernelContext*, const arrow::compute::ExecBatch& batch, arrow::Datum* res) { + if (const auto& arg = batch.values.front(); !arg.scalar()->is_valid) { + *res = arrow::MakeNullScalar(GetPrimitiveDataType()); + } else { + const auto val = GetPrimitiveScalarValue(*arg.scalar()); + *res = MakeScalarDatum(TFuncInstance::Do(val)); + } + return arrow::Status::OK(); + } + + static arrow::Status ExecArray(arrow::compute::KernelContext*, const arrow::compute::ExecBatch& batch, arrow::Datum* res) { + static_assert(!std::is_same::value); + + const auto& arg = batch.values.front(); + auto& resArr = *res->array(); + + const auto& arr = *arg.array(); + auto length = arr.length; + const auto values = arr.GetValues(1); + auto resValues = resArr.GetMutableValues(1); + for (int64_t i = 0; i < length; ++i) { + resValues[i] = TFuncInstance::Do(values[i]); + } + + return arrow::Status::OK(); + } +}; + +template class TFunc> +void AddUnaryKernel(TKernelFamilyBase& owner) { + using TInputLayout = typename TInput::TLayout; + using TOutputLayout = typename TOutput::TLayout; + + using TFuncInstance = TFunc; + using TExecs = TUnaryKernelExecs; + + std::vector argTypes({ TInput::Id }); + NUdf::TDataTypeId returnType = TOutput::Id; + + arrow::compute::ScalarKernel k({ GetPrimitiveInputArrowType() }, GetPrimitiveOutputArrowType(), &TExecs::Exec); + k.null_handling = owner.NullMode == TKernelFamily::ENullMode::Default ? arrow::compute::NullHandling::INTERSECTION : arrow::compute::NullHandling::COMPUTED_PREALLOCATE; + owner.Adopt(argTypes, returnType, std::make_unique(owner, argTypes, returnType, k)); +} + template class TFunc> void AddBinaryKernel(TKernelFamilyBase& owner) { @@ -1125,6 +1175,19 @@ void AddBinaryKernelPoly(TKernelFamilyBase& owner) { owner.Adopt(argTypes, returnType, std::make_unique(owner, argTypes, returnType, k)); } +template class TFunc> +void AddUnaryIntegralKernels(TKernelFamilyBase& owner) { + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); + AddUnaryKernel, NUdf::TDataType, TFunc>(owner); +} + template class TFunc> void AddBinaryIntegralKernels(TKernelFamilyBase& owner) { AddBinaryKernel, NUdf::TDataType, NUdf::TDataType, TFunc>(owner); @@ -1210,6 +1273,16 @@ public: } }; +template class TFunc> +class TUnaryNumericKernelFamily : public TKernelFamilyBase { +public: + TUnaryNumericKernelFamily(TKernelFamily::ENullMode nullMode = TKernelFamily::ENullMode::Default) + : TKernelFamilyBase(nullMode) + { + AddUnaryIntegralKernels(*this); + } +}; + template class TFunc> void AddBinaryPredicateKernel(TKernelFamilyBase& owner) { diff --git a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_minus.cpp b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_minus.cpp index aad3e5ccd748..e6f267718c3c 100644 --- a/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_minus.cpp +++ b/ydb/library/yql/minikql/invoke_builtins/mkql_builtins_minus.cpp @@ -48,5 +48,9 @@ void RegisterMinus(IBuiltinFunctionRegistry& registry) { RegisterFunctionUnOpt, NUdf::TDataType, TMinus, TUnaryArgsOpt>(registry, "Minus"); } +void RegisterMinus(TKernelFamilyMap& kernelFamilyMap) { + kernelFamilyMap["Minus"] = std::make_unique>(); +} + } // namespace NMiniKQL } // namespace NKikimr