structs_udf.cpp 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182
  1. #include <yql/essentials/public/udf/udf_value.h>
  2. #include <yql/essentials/public/udf/udf_value_builder.h>
  3. #include <yql/essentials/public/udf/udf_type_builder.h>
  4. #include <yql/essentials/public/udf/udf_registrator.h>
  5. #include <util/generic/yexception.h>
  6. #include <util/string/cast.h>
  7. using namespace NKikimr;
  8. using namespace NUdf;
  9. namespace {
  10. struct TPersonInfo
  11. {
  12. ui32 FirstName = 0;
  13. ui32 LastName = 0;
  14. ui32 Age = 0;
  15. bool RemapKSV = false;
  16. ui32 Key = 0;
  17. ui32 Subkey = 0;
  18. ui32 Value = 0;
  19. static constexpr ui32 FieldsCount = 3U;
  20. };
  21. //////////////////////////////////////////////////////////////////////////////
  22. // TPersonMember
  23. //////////////////////////////////////////////////////////////////////////////
  24. class TPersonMember: public TBoxedValue
  25. {
  26. public:
  27. explicit TPersonMember(ui32 memberIndex)
  28. : MemberIndex_(memberIndex)
  29. {
  30. }
  31. private:
  32. TUnboxedValue Run(
  33. const IValueBuilder* valueBuilder,
  34. const TUnboxedValuePod* args) const override
  35. {
  36. Y_UNUSED(valueBuilder);
  37. return args[0].GetElement(MemberIndex_);
  38. }
  39. const ui32 MemberIndex_;
  40. };
  41. //////////////////////////////////////////////////////////////////////////////
  42. // TNewPerson
  43. //////////////////////////////////////////////////////////////////////////////
  44. class TNewPerson: public TBoxedValue
  45. {
  46. public:
  47. explicit TNewPerson(const TPersonInfo& personIndexes)
  48. : Info_(personIndexes)
  49. {
  50. }
  51. private:
  52. TUnboxedValue Run(
  53. const IValueBuilder* valueBuilder,
  54. const TUnboxedValuePod* args) const override
  55. {
  56. TUnboxedValue name, surname, age;
  57. if (Info_.RemapKSV) {
  58. name = args->GetElement(Info_.Key);
  59. surname = args->GetElement(Info_.Subkey);
  60. const auto ageStr = args->GetElement(Info_.Value);
  61. const ui32 ageNum = FromString(ageStr.AsStringRef().Data(), ageStr.AsStringRef().Size());
  62. age = TUnboxedValuePod(ageNum);
  63. } else {
  64. name = TUnboxedValuePod(args[0]);
  65. surname = TUnboxedValuePod(args[1]);
  66. age = TUnboxedValuePod(args[2]);
  67. }
  68. TUnboxedValue* items = nullptr;
  69. auto result = valueBuilder->NewArray(Info_.FieldsCount, items);
  70. items[Info_.FirstName] = std::move(name);
  71. items[Info_.LastName] = std::move(surname);
  72. items[Info_.Age] = std::move(age);
  73. return result;
  74. }
  75. const TPersonInfo Info_;
  76. };
  77. //////////////////////////////////////////////////////////////////////////////
  78. // TPersonModule
  79. //////////////////////////////////////////////////////////////////////////////
  80. class TPersonModule: public IUdfModule
  81. {
  82. public:
  83. TStringRef Name() const {
  84. return TStringRef::Of("Person");
  85. }
  86. void CleanupOnTerminate() const final {}
  87. void GetAllFunctions(IFunctionsSink& sink) const final {
  88. sink.Add(TStringRef::Of("FirstName"));
  89. sink.Add(TStringRef::Of("LastName"));
  90. sink.Add(TStringRef::Of("Age"));
  91. sink.Add(TStringRef::Of("New"))->SetTypeAwareness();
  92. }
  93. void BuildFunctionTypeInfo(
  94. const TStringRef& name,
  95. TType* userType,
  96. const TStringRef& typeConfig,
  97. ui32 flags,
  98. IFunctionTypeInfoBuilder& builder) const final
  99. {
  100. Y_UNUSED(userType);
  101. try {
  102. bool typesOnly = (flags & TFlags::TypesOnly);
  103. TPersonInfo personInfo;
  104. auto personType = builder.Struct(personInfo.FieldsCount)->
  105. AddField<char*>("FirstName", &personInfo.FirstName)
  106. .AddField<char*>("LastName", &personInfo.LastName)
  107. .AddField<ui32>("Age", &personInfo.Age)
  108. .Build();
  109. if (TStringRef::Of("FirstName") == name) {
  110. // function signature: String FirstName(PersonStruct p)
  111. // runConfig: void
  112. builder.Returns<char*>().Args()->Add(personType).Done();
  113. if (!typesOnly) {
  114. builder.Implementation(new TPersonMember(personInfo.FirstName));
  115. }
  116. }
  117. else if (TStringRef::Of("LastName") == name) {
  118. // function signature: String LastName(PersonStruct p)
  119. // runConfig: void
  120. builder.Returns<char*>().Args()->Add(personType).Done();
  121. if (!typesOnly) {
  122. builder.Implementation(new TPersonMember(personInfo.LastName));
  123. }
  124. }
  125. else if (TStringRef::Of("Age") == name) {
  126. // function signature: ui32 Age(PersonStruct p)
  127. // runConfig: void
  128. builder.Returns<ui32>().Args()->Add(personType).Done();
  129. if (!typesOnly) {
  130. builder.Implementation(new TPersonMember(personInfo.Age));
  131. }
  132. }
  133. else if (TStringRef::Of("New") == name) {
  134. // function signature:
  135. // PersonStruct New(String firstName, String lastName, ui32 age)
  136. // runConfig: void
  137. builder.Returns(personType);
  138. if (TStringRef::Of("RemapKSV") == typeConfig) {
  139. personInfo.RemapKSV = true;
  140. auto inputType = builder.Struct(personInfo.FieldsCount)->
  141. AddField<char*>("key", &personInfo.Key)
  142. .AddField<char*>("subkey", &personInfo.Subkey)
  143. .AddField<char*>("value", &personInfo.Value)
  144. .Build();
  145. builder.Args()->Add(inputType);
  146. } else {
  147. builder.Args()->Add<char*>().Add<char*>().Add<ui32>();
  148. }
  149. builder.RunConfig<void>(); // this is optional because default runConfigType is Void
  150. if (!typesOnly) {
  151. builder.Implementation(new TNewPerson(personInfo));
  152. }
  153. }
  154. } catch (const std::exception& e) {
  155. builder.SetError(CurrentExceptionMessage());
  156. }
  157. }
  158. };
  159. } // namespace
  160. REGISTER_MODULES(TPersonModule)