mkql_tobytes.cpp 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. #include "mkql_tobytes.h"
  2. #include <yql/essentials/minikql/computation/mkql_computation_node_codegen.h> // Y_IGNORE
  3. #include <yql/essentials/minikql/mkql_node_cast.h>
  4. #include <yql/essentials/minikql/mkql_node_builder.h>
  5. #include <yql/essentials/minikql/mkql_string_util.h>
  6. #include <yql/essentials/utils/swap_bytes.h>
  7. namespace NKikimr {
  8. namespace NMiniKQL {
  9. namespace {
  10. template<bool IsOptional, typename Type>
  11. class TToBytesPrimitiveTypeWrapper : public TDecoratorCodegeneratorNode<TToBytesPrimitiveTypeWrapper<IsOptional, Type>> {
  12. using TBaseComputation = TDecoratorCodegeneratorNode<TToBytesPrimitiveTypeWrapper<IsOptional, Type>>;
  13. public:
  14. TToBytesPrimitiveTypeWrapper(IComputationNode* data)
  15. : TBaseComputation(data)
  16. {}
  17. NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const {
  18. if (IsOptional && !value)
  19. return NUdf::TUnboxedValuePod();
  20. const auto& v = value.Get<Type>();
  21. return NUdf::TUnboxedValuePod::Embedded(NUdf::TStringRef(reinterpret_cast<const char*>(&v), sizeof(v)));
  22. }
  23. #ifndef MKQL_DISABLE_CODEGEN
  24. Value* DoGenerateGetValue(const TCodegenContext& ctx, Value* value, BasicBlock*& block) const {
  25. auto& context = ctx.Codegen.GetContext();
  26. const uint64_t one[] = {0ULL, sizeof(Type) << 48ULL};
  27. const auto size = ConstantInt::get(value->getType(), APInt(128, 2, one));
  28. const uint64_t two[] = {0xFFFFFFFFFFFFFFFFULL, 0xFF00FFFFFFFFFFFFULL};
  29. const auto mask = ConstantInt::get(value->getType(), APInt(128, 2, two));
  30. const auto result = BinaryOperator::CreateOr(BinaryOperator::CreateAnd(value, mask, "and", block), size, "or", block);
  31. if constexpr (IsOptional)
  32. return SelectInst::Create(IsExists(value, block, context), result, GetEmpty(context), "select", block);
  33. return result;
  34. }
  35. #endif
  36. };
  37. template<bool IsOptional>
  38. class TToBytesDecimalWrapper : public TMutableComputationNode<TToBytesDecimalWrapper<IsOptional>> {
  39. using TBaseComputation = TMutableComputationNode<TToBytesDecimalWrapper<IsOptional>>;
  40. public:
  41. TToBytesDecimalWrapper(TComputationMutables& mutables, IComputationNode* data)
  42. : TBaseComputation(mutables, EValueRepresentation::String)
  43. , Data(data)
  44. {}
  45. NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
  46. const auto& value = this->Data->GetValue(ctx);
  47. if (IsOptional && !value)
  48. return NUdf::TUnboxedValuePod();
  49. return MakeString(NUdf::TStringRef(reinterpret_cast<const char*>(&value), 15));
  50. }
  51. void RegisterDependencies() const final {
  52. this->DependsOn(Data);
  53. }
  54. IComputationNode* const Data;
  55. };
  56. template<bool IsOptional, typename Type>
  57. class TToBytesTzTypeWrapper : public TDecoratorComputationNode<TToBytesTzTypeWrapper<IsOptional, Type>> {
  58. using TBaseComputation = TDecoratorComputationNode<TToBytesTzTypeWrapper<IsOptional, Type>>;
  59. public:
  60. TToBytesTzTypeWrapper(IComputationNode* data)
  61. : TBaseComputation(data)
  62. {}
  63. NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const {
  64. if (IsOptional && !value)
  65. return NUdf::TUnboxedValuePod();
  66. const auto v = NYql::SwapBytes(value.Get<Type>());
  67. const auto tzId = NYql::SwapBytes(value.GetTimezoneId());
  68. char buf[sizeof(Type) + sizeof(ui16)];
  69. std::memcpy(buf, &v, sizeof(v));
  70. std::memcpy(buf + sizeof(Type), &tzId, sizeof(tzId));
  71. return NUdf::TUnboxedValuePod::Embedded(NUdf::TStringRef(buf, sizeof(buf)));
  72. }
  73. };
  74. class TToBytesWrapper : public TDecoratorCodegeneratorNode<TToBytesWrapper> {
  75. using TBaseComputation = TDecoratorCodegeneratorNode<TToBytesWrapper>;
  76. public:
  77. TToBytesWrapper(IComputationNode* optional)
  78. : TBaseComputation(optional)
  79. {}
  80. NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const { return value; }
  81. #ifndef MKQL_DISABLE_CODEGEN
  82. Value* DoGenerateGetValue(const TCodegenContext&, Value* value, BasicBlock*&) const { return value; }
  83. #endif
  84. };
  85. }
  86. IComputationNode* WrapToBytes(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
  87. MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");
  88. bool isOptional;
  89. const auto dataType = UnpackOptionalData(callable.GetInput(0), isOptional);
  90. const auto data = LocateNode(ctx.NodeLocator, callable, 0);
  91. switch(dataType->GetSchemeType()) {
  92. #define MAKE_PRIMITIVE_TYPE_BYTES(type, layout) \
  93. case NUdf::TDataType<type>::Id: \
  94. if (isOptional) \
  95. return new TToBytesPrimitiveTypeWrapper<true, layout>(data); \
  96. else \
  97. return new TToBytesPrimitiveTypeWrapper<false, layout>(data);
  98. KNOWN_FIXED_VALUE_TYPES(MAKE_PRIMITIVE_TYPE_BYTES)
  99. #undef MAKE_PRIMITIVE_TYPE_BYTES
  100. #define MAKE_TZ_TYPE_BYTES(type, layout) \
  101. case NUdf::TDataType<type>::Id: \
  102. if (isOptional) \
  103. return new TToBytesTzTypeWrapper<true, layout>(data); \
  104. else \
  105. return new TToBytesTzTypeWrapper<false, layout>(data);
  106. MAKE_TZ_TYPE_BYTES(NUdf::TTzDate, ui16);
  107. MAKE_TZ_TYPE_BYTES(NUdf::TTzDatetime, ui32);
  108. MAKE_TZ_TYPE_BYTES(NUdf::TTzTimestamp, ui64);
  109. MAKE_TZ_TYPE_BYTES(NUdf::TTzDate32, i32);
  110. MAKE_TZ_TYPE_BYTES(NUdf::TTzDatetime64, i64);
  111. MAKE_TZ_TYPE_BYTES(NUdf::TTzTimestamp64, i64);
  112. #undef MAKE_TZ_TYPE_BYTES
  113. case NUdf::TDataType<NUdf::TDecimal>::Id: {
  114. if (isOptional) \
  115. return new TToBytesDecimalWrapper<true>(ctx.Mutables, data);
  116. else
  117. return new TToBytesDecimalWrapper<false>(ctx.Mutables, data);
  118. }
  119. default:
  120. break;
  121. }
  122. return new TToBytesWrapper(data);
  123. }
  124. }
  125. }