mkql_node_builder.cpp 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399
  1. #include "mkql_node_builder.h"
  2. #include "mkql_node_printer.h"
  3. #include <util/generic/algorithm.h>
  4. namespace NKikimr {
  5. namespace NMiniKQL {
  6. using namespace NDetail;
  7. namespace {
  8. TNode* BuildCallableTypePayload(const TVector<TStringBuf>& argNames,
  9. const TVector<ui64>& argFlags, TStringBuf payload, const TTypeEnvironment& env) {
  10. TStructTypeBuilder itemTypeBuilder(env);
  11. itemTypeBuilder.Add("Name", TDataType::Create(NUdf::TDataType<char*>::Id, env));
  12. itemTypeBuilder.Add("Flags", TDataType::Create(NUdf::TDataType<ui64>::Id, env));
  13. TType* itemType = itemTypeBuilder.Build();
  14. bool startedNames = false;
  15. THashSet<TStringBuf> usedNames;
  16. TListLiteralBuilder argsListBuilder(env, itemType);
  17. for (ui32 i = 0; i < argNames.size(); ++i) {
  18. TStructLiteralBuilder itemBuilder(env);
  19. auto name = argNames[i];
  20. auto flags = argFlags[i];
  21. bool hasName = name.size() > 0;
  22. if (startedNames) {
  23. MKQL_ENSURE(hasName, "Named arguments already started");
  24. } else {
  25. startedNames = hasName;
  26. }
  27. if (hasName) {
  28. MKQL_ENSURE(usedNames.insert(name).second, "Duplication of argument name: " << name);
  29. }
  30. itemBuilder.Add("Name", TRuntimeNode(BuildDataLiteral(name, NUdf::EDataSlot::String, env), true));
  31. itemBuilder.Add("Flags", TRuntimeNode(BuildDataLiteral(NUdf::TUnboxedValuePod(flags), NUdf::EDataSlot::Uint64, env), true));
  32. argsListBuilder.Add(TRuntimeNode(itemBuilder.Build(), true));
  33. }
  34. TRuntimeNode argsList = TRuntimeNode(argsListBuilder.Build(), true);
  35. TStructLiteralBuilder payloadBuilder(env);
  36. payloadBuilder.Add("Args", argsList);
  37. payloadBuilder.Add("Payload", TRuntimeNode(BuildDataLiteral(payload, NUdf::EDataSlot::String, env), true));
  38. return payloadBuilder.Build();
  39. }
  40. }
  41. TDataLiteral* BuildDataLiteral(const NUdf::TUnboxedValuePod& value, NUdf::TDataTypeId type, const TTypeEnvironment& env) {
  42. return TDataLiteral::Create(value, TDataType::Create(type, env), env);
  43. }
  44. TDataLiteral* BuildDataLiteral(const NUdf::TStringRef& data, NUdf::TDataTypeId type, const TTypeEnvironment& env) {
  45. return BuildDataLiteral(env.NewStringValue(data), type, env);
  46. }
  47. TOptionalLiteral* BuildOptionalLiteral(TRuntimeNode value, const TTypeEnvironment& env) {
  48. auto type = TOptionalType::Create(value.GetStaticType(), env);
  49. return TOptionalLiteral::Create(value, type, env);
  50. }
  51. TOptionalLiteral* BuildEmptyOptionalLiteral(TType *itemType, const TTypeEnvironment& env) {
  52. auto type = TOptionalType::Create(itemType, env);
  53. return TOptionalLiteral::Create(type, env);
  54. }
  55. TOptionalLiteral* BuildEmptyOptionalDataLiteral(NUdf::TDataTypeId schemeType, const TTypeEnvironment& env) {
  56. return BuildEmptyOptionalLiteral(TDataType::Create(schemeType, env), env);
  57. }
  58. TType* UnpackOptional(TRuntimeNode data, bool& isOptional) {
  59. return UnpackOptional(data.GetStaticType(), isOptional);
  60. }
  61. TType* UnpackOptional(TType* type, bool& isOptional) {
  62. isOptional = type->IsOptional();
  63. if (!isOptional)
  64. return type;
  65. auto optType = static_cast<TOptionalType*>(type);
  66. return optType->GetItemType();
  67. }
  68. TDataType* UnpackOptionalData(TRuntimeNode data, bool& isOptional) {
  69. return UnpackOptionalData(data.GetStaticType(), isOptional);
  70. }
  71. TDataType* UnpackOptionalData(TType* type, bool& isOptional) {
  72. auto unpackedType = UnpackOptional(type, isOptional);
  73. MKQL_ENSURE(unpackedType->IsData(),
  74. "Expected data or optional of data, actual: " << PrintNode(type, true));
  75. return static_cast<TDataType*>(unpackedType);
  76. }
  77. TBlockType::EShape GetResultShape(const TVector<TType*>& types) {
  78. TBlockType::EShape result = TBlockType::EShape::Scalar;
  79. for (auto& type : types) {
  80. MKQL_ENSURE(type->IsBlock(),
  81. "Expected block, actual: " << PrintNode(type, true));
  82. if (static_cast<TBlockType*>(type)->GetShape() == TBlockType::EShape::Many) {
  83. result = TBlockType::EShape::Many;
  84. }
  85. }
  86. return result;
  87. }
  88. TTupleLiteralBuilder::TTupleLiteralBuilder(const TTypeEnvironment& env) : Env(env)
  89. {}
  90. void TTupleLiteralBuilder::Reserve(ui32 size) {
  91. Types.reserve(size);
  92. Values.reserve(size);
  93. }
  94. TTupleLiteralBuilder& TTupleLiteralBuilder::Add(TRuntimeNode value) {
  95. Types.push_back(value.GetRuntimeType());
  96. Values.push_back(value);
  97. return *this;
  98. }
  99. TTupleLiteral* TTupleLiteralBuilder::Build() {
  100. const auto& type = TTupleType::Create(Types.size(), Types.data(), Env);
  101. return TTupleLiteral::Create(Values.size(), Values.data(), type, Env);
  102. }
  103. void TTupleLiteralBuilder::Clear() {
  104. Values.clear();
  105. Types.clear();
  106. }
  107. TStructTypeBuilder::TStructTypeBuilder(const TTypeEnvironment& env)
  108. : Env(&env)
  109. {
  110. }
  111. void TStructTypeBuilder::Reserve(ui32 size) {
  112. Members.reserve(size);
  113. }
  114. TStructTypeBuilder& TStructTypeBuilder::Add(const TStringBuf& name, TType* type, ui32* index) {
  115. Members.push_back(TStructMember(Env->InternName(name).Str(), type, index));
  116. return *this;
  117. }
  118. TStructType* TStructTypeBuilder::Build() {
  119. if (Members.empty())
  120. return Env->GetEmptyStructLazy()->GetType();
  121. Sort(Members.begin(), Members.end());
  122. return TStructType::Create(Members.size(), Members.data(), *Env);
  123. }
  124. void TStructTypeBuilder::FillIndexes() {
  125. ui32 index = 0;
  126. for (const TStructMember& member: Members) {
  127. if (member.Index) {
  128. *(member.Index) = index++;
  129. }
  130. }
  131. }
  132. void TStructTypeBuilder::Clear() {
  133. Members.clear();
  134. }
  135. TStructLiteralBuilder::TStructLiteralBuilder(const TTypeEnvironment& env)
  136. : Env(&env)
  137. {
  138. }
  139. void TStructLiteralBuilder::Reserve(ui32 size) {
  140. Members.reserve(size);
  141. Values.reserve(size);
  142. }
  143. TStructLiteralBuilder& TStructLiteralBuilder::Add(const TStringBuf& name, TRuntimeNode value) {
  144. TType* valueType = value.GetStaticType();
  145. Members.push_back(TStructMember(Env->InternName(name).Str(), valueType));
  146. Values.push_back(value);
  147. return *this;
  148. }
  149. TStructLiteral* TStructLiteralBuilder::Build() {
  150. Y_DEBUG_ABORT_UNLESS(Members.size() == Values.size());
  151. if (Members.empty())
  152. return Env->GetEmptyStructLazy();
  153. TVector<std::pair<TStringBuf, ui32>> sortedIndicies(Members.size());
  154. for (ui32 i = 0, e = Members.size(); i < e; ++i) {
  155. sortedIndicies[i] = std::make_pair(Members[i].Name, i);
  156. }
  157. Sort(sortedIndicies.begin(), sortedIndicies.end(),
  158. [](const std::pair<TStringBuf, ui32>& x, const std::pair<TStringBuf, ui32>& y) {
  159. return x.first < y.first;
  160. });
  161. TVector<TStructMember> sortedMembers(Members.size());
  162. TVector<TRuntimeNode> sortedValues(Members.size());
  163. for (ui32 i = 0, e = Members.size(); i < e; ++i) {
  164. sortedMembers[i] = Members[sortedIndicies[i].second];
  165. sortedValues[i] = Values[sortedIndicies[i].second];
  166. }
  167. auto type = TStructType::Create(sortedMembers.size(), sortedMembers.data(), *Env);
  168. return TStructLiteral::Create(sortedValues.size(), sortedValues.data(), type, *Env);
  169. }
  170. void TStructLiteralBuilder::Clear() {
  171. Members.clear();
  172. }
  173. TListLiteralBuilder::TListLiteralBuilder(const TTypeEnvironment& env, TType* type)
  174. : Env(&env)
  175. , Type(type)
  176. {}
  177. TListLiteralBuilder& TListLiteralBuilder::Add(TRuntimeNode item) {
  178. Items.push_back(item);
  179. return *this;
  180. }
  181. TListLiteral* TListLiteralBuilder::Build() {
  182. auto type = TListType::Create(Type, *Env);
  183. return TListLiteral::Create(Items.data(), Items.size(), type, *Env);
  184. }
  185. void TListLiteralBuilder::Clear() {
  186. Items.clear();
  187. }
  188. TDictLiteralBuilder::TDictLiteralBuilder(const TTypeEnvironment& env, TType* keyType, TType* payloadType)
  189. : Env(&env)
  190. , KeyType(keyType)
  191. , PayloadType(payloadType)
  192. {
  193. }
  194. void TDictLiteralBuilder::Reserve(ui32 size) {
  195. Items.reserve(size);
  196. }
  197. TDictLiteralBuilder& TDictLiteralBuilder::Add(TRuntimeNode key, TRuntimeNode payload) {
  198. Items.push_back(std::make_pair(key, payload));
  199. return *this;
  200. }
  201. TDictLiteral* TDictLiteralBuilder::Build() {
  202. auto type = TDictType::Create(KeyType, PayloadType, *Env);
  203. return TDictLiteral::Create(Items.size(), Items.data(), type, *Env);
  204. }
  205. void TDictLiteralBuilder::Clear() {
  206. Items.clear();
  207. }
  208. TCallableTypeBuilder::TCallableTypeBuilder(const TTypeEnvironment& env, const TStringBuf& name, TType* returnType)
  209. : Env(&env)
  210. , Name(Env->InternName(name))
  211. , ReturnType(returnType)
  212. , OptionalArgsCount(0)
  213. , HasPayload(false)
  214. {}
  215. void TCallableTypeBuilder::Reserve(ui32 size) {
  216. Arguments.reserve(size);
  217. ArgNames.reserve(size);
  218. ArgFlags.reserve(size);
  219. }
  220. TCallableTypeBuilder& TCallableTypeBuilder::Add(TType *type) {
  221. Arguments.push_back(type);
  222. ArgNames.emplace_back();
  223. ArgFlags.emplace_back();
  224. return *this;
  225. }
  226. TCallableTypeBuilder& TCallableTypeBuilder::SetArgumentName(const TStringBuf& name) {
  227. MKQL_ENSURE(!ArgNames.empty(), "No arguments");
  228. HasPayload = true;
  229. ArgNames.back() = name;
  230. return *this;
  231. }
  232. TCallableTypeBuilder& TCallableTypeBuilder::SetArgumentFlags(ui64 flags) {
  233. MKQL_ENSURE(!ArgFlags.empty(), "No arguments");
  234. HasPayload = true;
  235. ArgFlags.back() = flags;
  236. return *this;
  237. }
  238. TCallableTypeBuilder& TCallableTypeBuilder::SetOptionalArgs(ui32 count) {
  239. OptionalArgsCount = count;
  240. return *this;
  241. }
  242. TCallableTypeBuilder& TCallableTypeBuilder::SetPayload(const TStringBuf& data) {
  243. HasPayload = true;
  244. FuncPayload = data;
  245. return *this;
  246. }
  247. TCallableType* TCallableTypeBuilder::Build() {
  248. TNode* payload = nullptr;
  249. if (HasPayload) {
  250. payload = BuildCallableTypePayload(ArgNames, ArgFlags, FuncPayload, *Env);
  251. }
  252. auto ret = TCallableType::Create(ReturnType, Name.Str(), Arguments.size(), Arguments.data(), payload, *Env);
  253. ret->SetOptionalArgumentsCount(OptionalArgsCount);
  254. return ret;
  255. }
  256. void TCallableTypeBuilder::Clear() {
  257. Arguments.clear();
  258. ArgNames.clear();
  259. ArgFlags.clear();
  260. OptionalArgsCount = 0;
  261. HasPayload = false;
  262. FuncPayload = TStringBuf();
  263. }
  264. TCallableBuilder::TCallableBuilder(const TTypeEnvironment& env, const TStringBuf& name, TType* returnType, bool disableMerge)
  265. : Env(&env)
  266. , Name(Env->InternName(name))
  267. , ReturnType(returnType)
  268. , DisableMerge(disableMerge)
  269. , OptionalArgsCount(0)
  270. , HasPayload(false)
  271. {}
  272. void TCallableBuilder::Reserve(ui32 size) {
  273. Arguments.reserve(size);
  274. Inputs.reserve(size);
  275. ArgNames.reserve(size);
  276. ArgFlags.reserve(size);
  277. }
  278. TCallableBuilder& TCallableBuilder::Add(TRuntimeNode input) {
  279. TType* inputType = input.GetStaticType();
  280. Arguments.push_back(inputType);
  281. Inputs.push_back(input);
  282. ArgNames.emplace_back();
  283. ArgFlags.emplace_back();
  284. return *this;
  285. }
  286. TCallableBuilder& TCallableBuilder::SetOptionalArgs(ui32 count) {
  287. OptionalArgsCount = count;
  288. return *this;
  289. }
  290. TCallableBuilder& TCallableBuilder::SetTypePayload(const TStringBuf& data) {
  291. HasPayload = true;
  292. FuncPayload = data;
  293. return *this;
  294. }
  295. TCallableBuilder& TCallableBuilder::SetArgumentName(const TStringBuf& name) {
  296. MKQL_ENSURE(!ArgNames.empty(), "No arguments");
  297. HasPayload = true;
  298. ArgNames.back() = name;
  299. return *this;
  300. }
  301. TCallableBuilder& TCallableBuilder::SetArgumentFlags(ui64 flags) {
  302. MKQL_ENSURE(!ArgFlags.empty(), "No arguments");
  303. HasPayload = true;
  304. ArgFlags.back() = flags;
  305. return *this;
  306. }
  307. TCallable* TCallableBuilder::Build() {
  308. TNode* payload = nullptr;
  309. if (HasPayload) {
  310. payload = BuildCallableTypePayload(ArgNames, ArgFlags, FuncPayload, *Env);
  311. }
  312. auto type = TCallableType::Create(ReturnType, Name.Str(), Arguments.size(), Arguments.data(), payload, *Env);
  313. type->SetOptionalArgumentsCount(OptionalArgsCount);
  314. if (DisableMerge)
  315. type->DisableMerge();
  316. return TCallable::Create(Inputs.size(), Inputs.data(), type, *Env);
  317. }
  318. void TCallableBuilder::Clear() {
  319. Arguments.clear();
  320. Inputs.clear();
  321. ArgNames.clear();
  322. ArgFlags.clear();
  323. OptionalArgsCount = 0;
  324. FuncPayload = TStringBuf();
  325. }
  326. }
  327. }