yql_result_provider.cpp 72 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576
  1. #include "yql_result_provider.h"
  2. #include <yql/essentials/providers/result/expr_nodes/yql_res_expr_nodes.h>
  3. #include <yql/essentials/providers/common/provider/yql_provider.h>
  4. #include <yql/essentials/providers/common/provider/yql_provider_names.h>
  5. #include <yql/essentials/providers/common/provider/yql_data_provider_impl.h>
  6. #include <yql/essentials/providers/common/mkql/yql_type_mkql.h>
  7. #include <yql/essentials/providers/common/codec/yql_codec.h>
  8. #include <yql/essentials/core/yql_execution.h>
  9. #include <yql/essentials/core/yql_expr_optimize.h>
  10. #include <yql/essentials/core/yql_expr_type_annotation.h>
  11. #include <yql/essentials/core/yql_opt_utils.h>
  12. #include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
  13. #include <yql/essentials/ast/yql_gc_nodes.h>
  14. #include <yql/essentials/utils/log/log.h>
  15. #include <library/cpp/yson/node/node_io.h>
  16. #include <util/string/cast.h>
  17. namespace NYql {
  18. namespace {
  19. using namespace NKikimr;
  20. using namespace NKikimr::NMiniKQL;
  21. using namespace NNodes;
  22. class TYsonResultWriter : public IResultWriter
  23. {
  24. public:
  25. TYsonResultWriter(NYson::EYsonFormat format)
  26. : Writer(new NYson::TYsonWriter(&PartialStream, format, ::NYson::EYsonType::Node, true))
  27. {}
  28. void Init(bool discard, const TString& label, TMaybe<TPosition> pos, bool unordered) override {
  29. Discard = discard;
  30. Unordered = unordered;
  31. if (!Discard) {
  32. Writer->OnBeginMap();
  33. if (label) {
  34. Writer->OnKeyedItem("Label");
  35. Writer->OnStringScalar(label);
  36. }
  37. if (pos) {
  38. Writer->OnKeyedItem("Position");
  39. Writer->OnBeginMap();
  40. Writer->OnKeyedItem("File");
  41. Writer->OnStringScalar(pos->File ? pos->File : "<main>");
  42. Writer->OnKeyedItem("Row");
  43. Writer->OnInt64Scalar(pos->Row);
  44. Writer->OnKeyedItem("Column");
  45. Writer->OnInt64Scalar(pos->Column);
  46. Writer->OnEndMap();
  47. }
  48. Writer->OnKeyedItem("Write");
  49. Writer->OnBeginList();
  50. }
  51. }
  52. void Write(const TStringBuf& resultData) override {
  53. if (!Discard) {
  54. Writer->OnListItem();
  55. Writer->OnRaw(resultData);
  56. }
  57. }
  58. void Commit(bool overflow) override {
  59. if (!Discard) {
  60. Writer->OnEndList();
  61. if (overflow) {
  62. Writer->OnKeyedItem("Truncated");
  63. Writer->OnBooleanScalar(true);
  64. }
  65. if (Unordered) {
  66. Writer->OnKeyedItem("Unordered");
  67. Writer->OnBooleanScalar(true);
  68. }
  69. Writer->OnEndMap();
  70. }
  71. }
  72. bool IsDiscard() const override {
  73. return Discard;
  74. }
  75. TStringBuf Str() override {
  76. return PartialStream.Str();
  77. }
  78. ui64 Size() override {
  79. return PartialStream.Size();
  80. }
  81. private:
  82. TStringStream PartialStream;
  83. TAutoPtr<NYson::TYsonWriter> Writer;
  84. bool Discard = false;
  85. bool Unordered = false;
  86. };
  87. IGraphTransformer::TStatus ValidateColumns(TExprNode::TPtr& columns, const TTypeAnnotationNode* listType, TExprContext& ctx) {
  88. bool hasPrefixes = false;
  89. bool hasAutoNames = false;
  90. for (auto& child : columns->Children()) {
  91. if (HasError(child->GetTypeAnn(), ctx)) {
  92. return IGraphTransformer::TStatus::Error;
  93. }
  94. if (!child->IsAtom() && !child->IsList()) {
  95. ctx.AddError(TIssue(ctx.GetPosition(child->Pos()), "either atom or tuple is expected"));
  96. return IGraphTransformer::TStatus::Error;
  97. }
  98. if (child->IsList()) {
  99. if (!EnsureTupleMinSize(*child, 1, ctx)) {
  100. return IGraphTransformer::TStatus::Error;
  101. }
  102. if (!EnsureAtom(*child->Child(0), ctx)) {
  103. return IGraphTransformer::TStatus::Error;
  104. }
  105. if (child->Child(0)->Content() == "prefix") {
  106. if (!EnsureTupleSize(*child, 2, ctx)) {
  107. return IGraphTransformer::TStatus::Error;
  108. }
  109. if (!EnsureAtom(*child->Child(1), ctx)) {
  110. return IGraphTransformer::TStatus::Error;
  111. }
  112. hasPrefixes = true;
  113. } else if (child->Child(0)->Content() == "auto") {
  114. if (!EnsureTupleSize(*child, 1, ctx)) {
  115. return IGraphTransformer::TStatus::Error;
  116. }
  117. hasAutoNames = true;
  118. } else {
  119. ctx.AddError(TIssue(ctx.GetPosition(child->Pos()), TStringBuilder() <<
  120. "Expected 'prefix' or 'auto', but got: " << child->Child(0)->Content()));
  121. return IGraphTransformer::TStatus::Error;
  122. }
  123. }
  124. }
  125. if (listType->GetKind() == ETypeAnnotationKind::EmptyList) {
  126. return IGraphTransformer::TStatus::Ok;
  127. }
  128. if (listType->GetKind() != ETypeAnnotationKind::List) {
  129. ctx.AddError(TIssue(ctx.GetPosition(columns->Pos()), "columns requires list of struct"));
  130. return IGraphTransformer::TStatus::Error;
  131. }
  132. auto itemType = listType->Cast<TListExprType>()->GetItemType();
  133. if (itemType->GetKind() != ETypeAnnotationKind::Struct) {
  134. ctx.AddError(TIssue(ctx.GetPosition(columns->Pos()), "columns requires list of struct"));
  135. return IGraphTransformer::TStatus::Error;
  136. }
  137. auto structType = itemType->Cast<TStructExprType>();
  138. TSet<TString> usedFields;
  139. TExprNode::TListType orderedFields;
  140. TColumnOrder order;
  141. for (size_t i = 0; i < columns->ChildrenSize(); ++i) {
  142. auto child = columns->ChildPtr(i);
  143. if (child->IsAtom()) {
  144. orderedFields.push_back(child);
  145. auto rightName = order.AddColumn(TString(child->Content()));
  146. if (!structType->FindItem(rightName)) {
  147. if (hasAutoNames) {
  148. columns = {};
  149. return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
  150. }
  151. ctx.AddError(TIssue(ctx.GetPosition(child->Pos()), TStringBuilder() <<
  152. "Unknown field in hint: " << child->Content()));
  153. return IGraphTransformer::TStatus::Error;
  154. }
  155. if (!usedFields.insert(rightName).second) {
  156. if (hasAutoNames) {
  157. columns = {};
  158. return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
  159. }
  160. ctx.AddError(TIssue(ctx.GetPosition(child->Pos()), TStringBuilder() <<
  161. "Duplicate field in hint: " << rightName));
  162. return IGraphTransformer::TStatus::Error;
  163. }
  164. } else if (child->Child(0)->Content() == "auto") {
  165. TString columnName = "column" + ToString(i);
  166. if (!structType->FindItem(columnName) || !usedFields.insert(columnName).second) {
  167. columns = {};
  168. return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
  169. }
  170. orderedFields.push_back(ctx.NewAtom(child->Pos(), columnName));
  171. } else {
  172. auto prefix = child->Child(1)->Content();
  173. for (auto& x : structType->GetItems()) {
  174. if (x->GetName().StartsWith(prefix)) {
  175. orderedFields.push_back(ctx.NewAtom(child->Pos(), x->GetName()));
  176. if (!usedFields.insert(TString(x->GetName())).second) {
  177. if (hasAutoNames) {
  178. columns = {};
  179. return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
  180. }
  181. ctx.AddError(TIssue(ctx.GetPosition(child->Pos()), TStringBuilder() <<
  182. "Duplicate field in hint: " << x->GetName()));
  183. return IGraphTransformer::TStatus::Error;
  184. }
  185. }
  186. }
  187. }
  188. }
  189. if (usedFields.size() != structType->GetSize()) {
  190. if (hasAutoNames) {
  191. columns = {};
  192. return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
  193. }
  194. ctx.AddError(TIssue(ctx.GetPosition(columns->Pos()), TStringBuilder() <<
  195. "Mismatch of fields in hint and in the struct, columns fields: " << usedFields.size()
  196. << ", struct fields:" << structType->GetSize()));
  197. return IGraphTransformer::TStatus::Error;
  198. }
  199. if (hasPrefixes || hasAutoNames) {
  200. columns = ctx.NewList(columns->Pos(), std::move(orderedFields));
  201. return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
  202. }
  203. return IGraphTransformer::TStatus::Ok;
  204. }
  205. class TResultCallableExecutionTransformer : public TGraphTransformerBase {
  206. public:
  207. TResultCallableExecutionTransformer(const TIntrusivePtr<TResultProviderConfig>& config)
  208. : Config(config)
  209. {
  210. YQL_ENSURE(!Config->Types.AvailablePureResultDataSources.empty());
  211. }
  212. TStatus DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final {
  213. output = input;
  214. TString uniqId = TStringBuilder() << '#' << input->UniqueId();
  215. YQL_LOG_CTX_SCOPE(uniqId);
  216. YQL_ENSURE(input->Type() == TExprNode::Callable);
  217. TExprBase node(input);
  218. if (node.Maybe<TResFill>() || node.Maybe<TResPull>() || node.Maybe<TResIf>() || node.Maybe<TResFor>()) {
  219. auto provider = Config->Types.DataSourceMap.FindPtr(input->Child(5)->Content());
  220. Y_ENSURE(provider, "DataSource not exist: " << input->Child(5)->Content());
  221. if (node.Maybe<TResPull>()) {
  222. return HandleFillOrPull<TPull>(node, output, ctx, *(*provider));
  223. } else {
  224. return HandleFillOrPull<TResult>(node, output, ctx, *(*provider));
  225. }
  226. }
  227. if (input->Content() == CommitName) {
  228. if (ResultWriter) {
  229. TExprBase commitChild(input->ChildPtr(0));
  230. bool overflow = commitChild.Maybe<TResPull>() ? PullOverflow : FillOverflow;
  231. ui64& committedSize = commitChild.Maybe<TResPull>() ? CommittedPullSize : CommittedFillSize;
  232. if (!ResultWriter->IsDiscard()) {
  233. ResultWriter->Commit(overflow);
  234. Config->CommittedResults.push_back(TString(ResultWriter->Str()));
  235. committedSize += Config->CommittedResults.back().size();
  236. }
  237. ResultWriter.Reset();
  238. }
  239. input->SetState(TExprNode::EState::ExecutionComplete);
  240. input->SetResult(ctx.NewWorld (input->Pos()));
  241. return TStatus::Ok;
  242. }
  243. if (input->Content() == ConfigureName) {
  244. input->SetState(TExprNode::EState::ExecutionComplete);
  245. input->SetResult(ctx.NewWorld(input->Pos()));
  246. return TStatus::Ok;
  247. }
  248. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << "Failed to execute node: " << input->Content()));
  249. return TStatus::Ok;
  250. }
  251. NThreading::TFuture<void> DoGetAsyncFuture(const TExprNode& input) final {
  252. Y_UNUSED(input);
  253. YQL_ENSURE(DelegatedProvider);
  254. YQL_ENSURE(DelegatedNode);
  255. YQL_ENSURE(DelegatedNodeOutput);
  256. return DelegatedProvider->GetCallableExecutionTransformer()
  257. .GetAsyncFuture(*DelegatedNode);
  258. }
  259. TStatus DoApplyAsyncChanges(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final {
  260. output = input;
  261. YQL_ENSURE(DelegatedProvider);
  262. YQL_ENSURE(DelegatedNode);
  263. YQL_ENSURE(DelegatedNodeOutput);
  264. auto status = DelegatedProvider->GetCallableExecutionTransformer()
  265. .ApplyAsyncChanges(DelegatedNode, DelegatedNodeOutput, ctx);
  266. if (status == TStatus::Repeat && input != DelegatedNodeOutput->TailPtr()) {
  267. output = DelegatedNodeOutput->TailPtr();
  268. } else {
  269. FinishNode(*input, ctx, status);
  270. }
  271. return status;
  272. }
  273. void Rewind() final {
  274. DelegatedProvider = nullptr;
  275. DelegatedNode = nullptr;
  276. DelegatedNodeOutput = nullptr;
  277. CommittedPullSize = 0;
  278. PullOverflow = false;
  279. CommittedFillSize = 0;
  280. FillOverflow = false;
  281. ResultWriter.Drop();
  282. }
  283. private:
  284. template <class TTarget>
  285. bool& GetOverflowFlagAndCommitedSize(ui64& committed);
  286. template <class TTarget>
  287. TStatus HandleFillOrPull(TExprBase input, TExprNode::TPtr& output, TExprContext& ctx, IDataProvider& provider) {
  288. auto requireWorld = RequireChild(input.Ref(), TResBase::idx_World);
  289. auto requireData = input.Maybe<TResPull>() ? RequireChild(input.Ref(), TResPull::idx_Data) : IGraphTransformer::TStatus::Ok;
  290. auto requireStatus = requireWorld.Combine(requireData);
  291. if (requireStatus.Level != IGraphTransformer::TStatus::Ok) {
  292. return requireStatus;
  293. }
  294. bool needWriter = true;
  295. TMaybeNode<TExprBase> dataNode;
  296. TMaybeNode<TCoNameValueTupleList> options;
  297. if (input.Maybe<TResIf>()) {
  298. if (input.Ref().HasResult()) {
  299. auto resultYsonString = input.Ref().GetResult().Content();
  300. auto resultNode = NYT::NodeFromYsonString(TString(resultYsonString), ::NYson::EYsonType::Node);
  301. YQL_ENSURE(resultNode.IsMap());
  302. auto resultBoolNode = resultNode.AsMap()["Data"];
  303. YQL_ENSURE(resultBoolNode.IsBool());
  304. const bool predicate = resultBoolNode.AsBool();
  305. auto branchStatus = RequireChild(input.Ref(), predicate ? TResIf::idx_Then : TResIf::idx_Else);
  306. if (branchStatus.Level != IGraphTransformer::TStatus::Ok) {
  307. return branchStatus;
  308. }
  309. input.Ptr()->SetResult(ctx.NewWorld(input.Pos()));
  310. input.Ptr()->SetState(TExprNode::EState::ExecutionComplete);
  311. return TStatus::Ok;
  312. }
  313. needWriter = false;
  314. dataNode = input.Cast<TResIf>().Condition();
  315. options = input.Cast<TResIf>().Settings();
  316. } else if (input.Maybe<TResFor>()) {
  317. const auto& forNode = input.Cast<TResFor>();
  318. if (forNode.Total().Value()) {
  319. // has total, run active node if any
  320. const auto& items = forNode.Items();
  321. auto total = FromString<ui32>(forNode.Total().Value());
  322. auto current = FromString<ui32>(forNode.Current().Value());
  323. if ((current > total) || (total && current >= total)) {
  324. auto zero = ctx.NewAtom(TPositionHandle(), "0", TNodeFlags::Default);
  325. zero->SetTypeAnn(ctx.MakeType<TUnitExprType>());
  326. zero->SetState(TExprNode::EState::ConstrComplete);
  327. zero->SetDependencyScope(nullptr, nullptr); // HOTFIX for CSEE
  328. input.Ptr()->ChildRef(TResFor::idx_Current) = std::move(zero); //FIXME: Don't use ChilfRef
  329. input.Ptr()->SetResult(ctx.NewWorld(input.Pos()));
  330. input.Ptr()->SetState(TExprNode::EState::ExecutionComplete);
  331. return TStatus::Ok;
  332. }
  333. const auto& active = forNode.Active();
  334. if (active.Ref().IsCallable("Void")) {
  335. // select new active node
  336. TExprNode::TPtr active;
  337. if (total == 0) {
  338. // use else block
  339. const auto& elseLambda = forNode.Else();
  340. active = ctx.ReplaceNode(elseLambda.Body().Ptr(), elseLambda.Args().Arg(0).Ref(), forNode.World().Ptr());
  341. } else {
  342. // use some list item
  343. auto listNode = items.Raw();
  344. if (listNode->IsCallable("Just")) {
  345. listNode = listNode->Child(0);
  346. }
  347. ui32 index = current;
  348. if (listNode->IsCallable("List")) {
  349. ++index;
  350. }
  351. auto listElem = listNode->Child(index);
  352. const auto& iterLambda = forNode.Iter();
  353. active = ctx.ReplaceNodes(iterLambda.Body().Ptr(), { {
  354. iterLambda.Args().Arg(0).Raw(), forNode.World().Ptr()
  355. }, { iterLambda.Args().Arg(1).Raw(), listElem } });
  356. }
  357. output = Build<TResFor>(ctx, forNode.Pos())
  358. .World(forNode.World())
  359. .DataSink(forNode.DataSink())
  360. .Items(forNode.Items())
  361. .Iter(forNode.Iter())
  362. .Else(forNode.Else())
  363. .DelegatedSource(forNode.DelegatedSource())
  364. .Settings(forNode.Settings())
  365. .Total(forNode.Total())
  366. .Current()
  367. .Value(ToString(current))
  368. .Build()
  369. .Active(active)
  370. .Done().Ptr();
  371. return IGraphTransformer::TStatus(TStatus::Repeat, true);
  372. } else {
  373. auto status = RequireChild(input.Ref(), TResFor::idx_Active);
  374. if (status.Level != IGraphTransformer::TStatus::Ok) {
  375. return status;
  376. }
  377. ++current;
  378. // active node complete, drop it
  379. output = Build<TResFor>(ctx, forNode.Pos())
  380. .World(forNode.World())
  381. .DataSink(forNode.DataSink())
  382. .Items(forNode.Items())
  383. .Iter(forNode.Iter())
  384. .Else(forNode.Else())
  385. .DelegatedSource(forNode.DelegatedSource())
  386. .Settings(forNode.Settings())
  387. .Total(forNode.Total())
  388. .Current()
  389. .Value(ToString(current))
  390. .Build()
  391. .template Active<TCoVoid>()
  392. .Build()
  393. .Done().Ptr();
  394. return IGraphTransformer::TStatus(TStatus::Repeat, true);
  395. }
  396. } else if (input.Ref().HasResult()) {
  397. // parse list
  398. auto resultYsonString = input.Ref().GetResult().Content();
  399. auto resultNode = NYT::NodeFromYsonString(TString(resultYsonString), ::NYson::EYsonType::Node);
  400. YQL_ENSURE(resultNode.IsMap());
  401. auto resultDataNode = resultNode.AsMap()["Data"];
  402. const auto& itemsNode = forNode.Items().Ref();
  403. TScopedAlloc alloc(__LOCATION__);
  404. TTypeEnvironment env(alloc);
  405. TStringStream err;
  406. NKikimr::NMiniKQL::TTypeBuilder typeBuilder(env);
  407. TType* mkqlType = NCommon::BuildType(*itemsNode.GetTypeAnn(), typeBuilder, err);
  408. if (!mkqlType) {
  409. ctx.AddError(TIssue(ctx.GetPosition(itemsNode.Pos()), TStringBuilder() << "Failed to process type: " << err.Str()));
  410. return TStatus::Error;
  411. }
  412. TMemoryUsageInfo memInfo("Eval");
  413. THolderFactory holderFactory(alloc.Ref(), memInfo);
  414. auto value = NCommon::ParseYsonNodeInResultFormat(holderFactory, resultDataNode, mkqlType, &err);
  415. if (!value) {
  416. ctx.AddError(TIssue(ctx.GetPosition(itemsNode.Pos()), TStringBuilder() << "Failed to parse data: " << err.Str()));
  417. return TStatus::Error;
  418. }
  419. // build expr literal & total/active
  420. auto itemsLiteral = NCommon::ValueToExprLiteral(itemsNode.GetTypeAnn(), *value, ctx, itemsNode.Pos());
  421. ui32 totalItems = 0;
  422. auto listNode = itemsLiteral;
  423. if (listNode->IsCallable("Just")) {
  424. listNode = listNode->Child(0);
  425. }
  426. if (listNode->IsCallable("AsList")) {
  427. totalItems = listNode->ChildrenSize();
  428. } else if (listNode->IsCallable("List")) {
  429. totalItems = listNode->ChildrenSize() - 1;
  430. }
  431. output = Build<TResFor>(ctx, forNode.Pos())
  432. .World(forNode.World())
  433. .DataSink(forNode.DataSink())
  434. .Items(itemsLiteral)
  435. .Iter(forNode.Iter())
  436. .Else(forNode.Else())
  437. .DelegatedSource(forNode.DelegatedSource())
  438. .Settings(forNode.Settings())
  439. .Total()
  440. .Value(ToString(totalItems))
  441. .Build()
  442. .Current()
  443. .Value("0")
  444. .Build()
  445. .Active<TCoVoid>()
  446. .Build()
  447. .Done().Ptr();
  448. return IGraphTransformer::TStatus(TStatus::Repeat, true);
  449. }
  450. needWriter = false;
  451. dataNode = forNode.Items();
  452. options = forNode.Settings();
  453. } else {
  454. dataNode = input.Cast<TResWriteBase>().Data();
  455. options = input.Cast<TResWriteBase>().Settings();
  456. }
  457. DelegatedProvider = &provider;
  458. auto fillSettings = Config->FillSettings;
  459. auto resultSize = ResultWriter ? ResultWriter->Size() : 0;
  460. ui64 committedSize;
  461. bool& overflow = GetOverflowFlagAndCommitedSize<TTarget>(committedSize);
  462. if (fillSettings.AllResultsBytesLimit && committedSize + resultSize >= *fillSettings.AllResultsBytesLimit) {
  463. overflow = true;
  464. }
  465. if (fillSettings.AllResultsBytesLimit) {
  466. if (!overflow && committedSize <= *fillSettings.AllResultsBytesLimit) {
  467. *fillSettings.AllResultsBytesLimit -= committedSize;
  468. } else {
  469. *fillSettings.AllResultsBytesLimit = 0;
  470. }
  471. }
  472. auto atomType = ctx.MakeType<TUnitExprType>();
  473. auto rowsLimit = fillSettings.RowsLimitPerWrite;
  474. bool discard = false;
  475. TString label;
  476. bool unordered = false;
  477. for (auto setting : options.Cast()) {
  478. if (setting.Name().Value() == "take") {
  479. auto value = FromString<ui64>(setting.Value().Cast<TCoAtom>().Value());
  480. if (rowsLimit) {
  481. rowsLimit = Min(*rowsLimit, value);
  482. } else {
  483. rowsLimit = value;
  484. }
  485. } else if (setting.Name().Value() == "discard") {
  486. discard = true;
  487. } else if (setting.Name().Value() == "label") {
  488. label = TString(setting.Value().Cast<TCoAtom>().Value());
  489. } else if (setting.Name().Value() == "unordered") {
  490. unordered = true;
  491. }
  492. }
  493. TString publicId;
  494. if (auto id = Config->Types.TranslateOperationId(input.Ref().UniqueId())) {
  495. publicId = ToString(*id);
  496. }
  497. if (needWriter && !ResultWriter) {
  498. YQL_ENSURE(Config->WriterFactory);
  499. ResultWriter = Config->WriterFactory();
  500. ResultWriter->Init(discard, label, Config->SupportsResultPosition ?
  501. TMaybe<TPosition>(ctx.GetPosition(input.Pos())) : Nothing(), unordered);
  502. }
  503. if (input.Maybe<TResIf>() || input.Maybe<TResFor>()) {
  504. fillSettings = IDataProvider::TFillSettings();
  505. fillSettings.AllResultsBytesLimit.Clear();
  506. discard = fillSettings.Discard = false;
  507. fillSettings.Format = IDataProvider::EResultFormat::Yson;
  508. fillSettings.FormatDetails = ToString((ui32)NYson::EYsonFormat::Binary);
  509. fillSettings.RowsLimitPerWrite.Clear();
  510. rowsLimit.Clear();
  511. }
  512. DelegatedNode = Build<TTarget>(ctx, input.Pos())
  513. .Input(dataNode.Cast())
  514. .BytesLimit()
  515. .Value(fillSettings.AllResultsBytesLimit ? ToString(*fillSettings.AllResultsBytesLimit) : TString())
  516. .Build()
  517. .RowsLimit()
  518. .Value(rowsLimit ? ToString(*rowsLimit) : TString())
  519. .Build()
  520. .FormatDetails()
  521. .Value(fillSettings.FormatDetails)
  522. .Build()
  523. .Settings(options.Cast())
  524. .Format()
  525. .Value(ToString((ui32)fillSettings.Format))
  526. .Build()
  527. .PublicId()
  528. .Value(publicId)
  529. .Build()
  530. .Discard()
  531. .Value(ToString(discard))
  532. .Build()
  533. .Origin(input)
  534. .Done().Ptr();
  535. for (auto idx: {TResOrPullBase::idx_BytesLimit, TResOrPullBase::idx_RowsLimit, TResOrPullBase::idx_FormatDetails,
  536. TResOrPullBase::idx_Format, TResOrPullBase::idx_PublicId, TResOrPullBase::idx_Discard }) {
  537. DelegatedNode->Child(idx)->SetTypeAnn(atomType);
  538. DelegatedNode->Child(idx)->SetState(TExprNode::EState::ConstrComplete);
  539. }
  540. DelegatedNode->SetTypeAnn(input.Ref().GetTypeAnn());
  541. DelegatedNode->SetState(TExprNode::EState::ConstrComplete);
  542. input.Ptr()->SetState(TExprNode::EState::ExecutionInProgress);
  543. auto status = DelegatedProvider->GetCallableExecutionTransformer().Transform(DelegatedNode, DelegatedNodeOutput, ctx);
  544. if (status.Level != TStatus::Async) {
  545. status = FinishNode(*input.Ptr(), ctx, status);
  546. }
  547. return status;
  548. }
  549. IGraphTransformer::TStatus FinishNode(TExprNode& input, TExprContext& ctx, IGraphTransformer::TStatus status) {
  550. if (status.Level == TStatus::Ok) {
  551. auto data = DelegatedNode->GetResult().Content();
  552. const bool needWriter = input.Content() != TResIf::CallableName()
  553. && input.Content() != TResFor::CallableName();
  554. if (needWriter) {
  555. ResultWriter->Write(data);
  556. input.SetResult(ctx.NewAtom(input.Pos(), ""));
  557. input.SetState(TExprNode::EState::ExecutionComplete);
  558. } else {
  559. input.SetResult(ctx.NewAtom(input.Pos(), data));
  560. input.SetState(TExprNode::EState::ExecutionRequired);
  561. status = IGraphTransformer::TStatus::Repeat;
  562. }
  563. } else if (status.Level == TStatus::Error) {
  564. if (const auto issies = ctx.AssociativeIssues.extract(DelegatedNode.Get())) {
  565. ctx.IssueManager.RaiseIssues(issies.mapped());
  566. }
  567. } else {
  568. input.SetState(TExprNode::EState::ExecutionRequired);
  569. }
  570. DelegatedProvider = nullptr;
  571. DelegatedNode = nullptr;
  572. DelegatedNodeOutput = nullptr;
  573. return status;
  574. }
  575. private:
  576. const TIntrusivePtr<TResultProviderConfig> Config;
  577. IDataProvider* DelegatedProvider = nullptr;
  578. TExprNode::TPtr DelegatedNode;
  579. TExprNode::TPtr DelegatedNodeOutput;
  580. ui64 CommittedPullSize = 0;
  581. bool PullOverflow = false;
  582. ui64 CommittedFillSize = 0;
  583. bool FillOverflow = false;
  584. TIntrusivePtr<IResultWriter> ResultWriter;
  585. };
  586. template <class TTarget>
  587. bool& TResultCallableExecutionTransformer::GetOverflowFlagAndCommitedSize(ui64& committed) {
  588. committed = CommittedFillSize;
  589. return FillOverflow;
  590. }
  591. template<>
  592. bool& TResultCallableExecutionTransformer::GetOverflowFlagAndCommitedSize<TPull>(ui64& committed) {
  593. committed = CommittedPullSize;
  594. return PullOverflow;
  595. }
  596. class TResultTrackableNodeProcessor : public TTrackableNodeProcessorBase {
  597. public:
  598. TResultTrackableNodeProcessor(const TIntrusivePtr<TResultProviderConfig>& config)
  599. : Config(config)
  600. {}
  601. void GetUsedNodes(const TExprNode& input, TVector<TString>& usedNodeIds) override {
  602. usedNodeIds.clear();
  603. if (TMaybeNode<TResFill>(&input) || TMaybeNode<TResPull>(&input) || TMaybeNode<TResIf>(&input)
  604. || TMaybeNode<TResFor>(&input)) {
  605. auto provider = Config->Types.DataSourceMap.FindPtr(input.Child(5)->Content());
  606. Y_ENSURE(provider, "DataSource not exist: " << input.Child(5)->Content());
  607. (*provider)->GetTrackableNodeProcessor().GetUsedNodes(input, usedNodeIds);
  608. }
  609. }
  610. private:
  611. const TIntrusivePtr<TResultProviderConfig> Config;
  612. };
  613. class TPhysicalFinalizingTransformer final : public TSyncTransformerBase {
  614. public:
  615. TPhysicalFinalizingTransformer(const TIntrusivePtr<TResultProviderConfig>& config)
  616. : Config(config) {}
  617. TStatus DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final {
  618. TOptimizeExprSettings settings(&Config->Types);
  619. settings.ProcessedNodes = &PhysicalOptProcessedNodes;
  620. TStatus status = OptimizeExprEx(input, output,
  621. [&](const TExprNode::TPtr& node, TExprContext& ctx, IOptimizationContext& optCtx) -> TExprNode::TPtr {
  622. auto ret = node;
  623. if (auto maybeWrite = TMaybeNode<TResWrite>(node)) {
  624. auto resWrite = maybeWrite.Cast();
  625. bool isRef = false;
  626. bool isAutoRef = false;
  627. for (auto child: resWrite.Settings()) {
  628. if (child.Name().Value() == "ref") {
  629. isRef = true;
  630. }
  631. if (child.Name().Value() == "autoref") {
  632. isAutoRef = true;
  633. }
  634. }
  635. auto writeInput = resWrite.Data();
  636. for (auto& source : Config->Types.DataSources) {
  637. TSyncMap syncList;
  638. bool canRef;
  639. if (source->CanPullResult(writeInput.Ref(), syncList, canRef)) {
  640. auto newInput = writeInput;
  641. if (isRef && !canRef) {
  642. ctx.AddError(TIssue(ctx.GetPosition(writeInput.Pos()), TStringBuilder() <<
  643. "RefSelect mode isn't supported by provider: " << source->GetName()));
  644. return nullptr;
  645. }
  646. auto fillSettings = Config->FillSettings;
  647. if (!isRef && (!isAutoRef || !canRef)) {
  648. for (auto setting: resWrite.Settings()) {
  649. if (setting.Name().Value() == "take") {
  650. auto value = FromString<ui64>(setting.Value().Cast<TCoAtom>().Value());
  651. if (fillSettings.RowsLimitPerWrite) {
  652. fillSettings.RowsLimitPerWrite = Min(*fillSettings.RowsLimitPerWrite, value);
  653. } else {
  654. fillSettings.RowsLimitPerWrite = value;
  655. }
  656. }
  657. }
  658. if (fillSettings.RowsLimitPerWrite) {
  659. *fillSettings.RowsLimitPerWrite += 1;
  660. }
  661. } else {
  662. fillSettings.RowsLimitPerWrite.Clear();
  663. }
  664. newInput = TExprBase(source->OptimizePull(newInput.Ptr(), fillSettings, ctx, optCtx));
  665. ret = Build<TResPull>(ctx, resWrite.Pos())
  666. .World(ApplySyncListToWorld(resWrite.World().Ptr(), syncList, ctx))
  667. .DataSink(resWrite.DataSink())
  668. .Key(resWrite.Key())
  669. .Data(newInput)
  670. .Settings(resWrite.Settings())
  671. .DelegatedSource()
  672. .Value(source->GetName())
  673. .Build()
  674. .Done().Ptr();
  675. YQL_CLOG(INFO, ProviderResult) << "ResPull";
  676. return ret;
  677. }
  678. }
  679. if (!isRef) {
  680. auto data = resWrite.Data();
  681. if (auto unordered = data.Maybe<TCoUnorderedBase>()) {
  682. data = unordered.Cast().Input();
  683. }
  684. TSyncMap syncList;
  685. if (IsPureIsolatedLambda(writeInput.Ref(), &syncList)) {
  686. auto cleanup = DefaultCleanupWorld(data.Ptr(), ctx);
  687. if (!cleanup) {
  688. return nullptr;
  689. }
  690. ret = Build<TResFill>(ctx, resWrite.Pos())
  691. .World(ApplySyncListToWorld(resWrite.World().Ptr(), syncList, ctx))
  692. .DataSink(resWrite.DataSink())
  693. .Key(resWrite.Key())
  694. .Data(cleanup)
  695. .Settings(resWrite.Settings())
  696. .DelegatedSource()
  697. .Value(Config->Types.GetDefaultDataSource())
  698. .Build()
  699. .Done().Ptr();
  700. YQL_CLOG(INFO, ProviderResult) << "ResFill";
  701. return ret;
  702. }
  703. for (auto& source : Config->Types.DataSources) {
  704. TSyncMap syncList;
  705. if (source->CanBuildResult(writeInput.Ref(), syncList)) {
  706. auto cleanup = source->CleanupWorld(data.Ptr(), ctx);
  707. if (!cleanup) {
  708. return nullptr;
  709. }
  710. ret = Build<TResFill>(ctx, resWrite.Pos())
  711. .World(ApplySyncListToWorld(resWrite.World().Ptr(), syncList, ctx))
  712. .DataSink(resWrite.DataSink())
  713. .Key(resWrite.Key())
  714. .Data(cleanup)
  715. .Settings(resWrite.Settings())
  716. .DelegatedSource()
  717. .Value(source->GetName())
  718. .Build()
  719. .Done().Ptr();
  720. YQL_CLOG(INFO, ProviderResult) << "ResFill";
  721. return ret;
  722. }
  723. }
  724. }
  725. } else if (node->Content() == IfName) {
  726. TSyncMap syncList;
  727. auto foundDataSource = FindDataSource(*node->Child(1), syncList);
  728. if (!foundDataSource.empty()) {
  729. auto provider = Config->Types.DataSourceMap.FindPtr(foundDataSource);
  730. Y_ENSURE(provider, "DataSource doesn't exist: " << foundDataSource);
  731. auto cleanup = (*provider)->CleanupWorld(node->ChildPtr(1), ctx);
  732. if (!cleanup) {
  733. return nullptr;
  734. }
  735. ret = Build<TResIf>(ctx, node->Pos())
  736. .World(ApplySyncListToWorld(node->ChildPtr(0), syncList, ctx))
  737. .DataSink()
  738. .Build()
  739. .Condition(cleanup)
  740. .Then(node->ChildPtr(2))
  741. .Else(node->ChildPtr(3))
  742. .DelegatedSource()
  743. .Value(foundDataSource)
  744. .Build()
  745. .Settings()
  746. .Build()
  747. .Done().Ptr();
  748. YQL_CLOG(INFO, ProviderResult) << "ResIf";
  749. return ret;
  750. }
  751. } else if (node->Content() == ForName) {
  752. TSyncMap syncList;
  753. auto foundDataSource = FindDataSource(*node->Child(1), syncList);
  754. if (!foundDataSource.empty()) {
  755. auto provider = Config->Types.DataSourceMap.FindPtr(foundDataSource);
  756. Y_ENSURE(provider, "DataSource doesn't exist: " << foundDataSource);
  757. auto cleanup = (*provider)->CleanupWorld(node->ChildPtr(1), ctx);
  758. if (!cleanup) {
  759. return nullptr;
  760. }
  761. ret = Build<TResFor>(ctx, node->Pos())
  762. .World(ApplySyncListToWorld(node->ChildPtr(0), syncList, ctx))
  763. .DataSink()
  764. .Build()
  765. .Items(cleanup)
  766. .Iter(node->ChildPtr(2))
  767. .Else(node->ChildPtr(3))
  768. .DelegatedSource()
  769. .Value(foundDataSource)
  770. .Build()
  771. .Settings()
  772. .Build()
  773. .Total()
  774. .Value("")
  775. .Build()
  776. .Current()
  777. .Value("")
  778. .Build()
  779. .Active<TCoVoid>()
  780. .Build()
  781. .Done().Ptr();
  782. YQL_CLOG(INFO, ProviderResult) << "ResFor";
  783. return ret;
  784. }
  785. }
  786. return ret;
  787. }, ctx, settings);
  788. return status;
  789. }
  790. void Rewind() final {
  791. PhysicalOptProcessedNodes.clear();
  792. }
  793. private:
  794. TString FindDataSource(const TExprNode& node, TSyncMap& syncList) const {
  795. syncList.clear();
  796. TString foundDataSource;
  797. if (IsPureIsolatedLambda(node)) {
  798. foundDataSource = Config->Types.GetDefaultDataSource();
  799. }
  800. if (foundDataSource.empty()) {
  801. for (auto& source : Config->Types.DataSources) {
  802. syncList.clear();
  803. if (source->CanBuildResult(node, syncList)) {
  804. foundDataSource = TString(source->GetName());
  805. break;
  806. }
  807. }
  808. }
  809. return foundDataSource;
  810. }
  811. private:
  812. const TIntrusivePtr<TResultProviderConfig> Config;
  813. TProcessedNodesSet PhysicalOptProcessedNodes;
  814. };
  815. class TResultProvider : public TDataProviderBase {
  816. public:
  817. struct TFunctions {
  818. THashSet<TStringBuf> Names;
  819. TFunctions() {
  820. Names.insert(TResWrite::CallableName());
  821. Names.insert(TResFill::CallableName());
  822. Names.insert(TResPull::CallableName());
  823. Names.insert(TResIf::CallableName());
  824. Names.insert(TResFor::CallableName());
  825. }
  826. };
  827. TResultProvider(const TIntrusivePtr<TResultProviderConfig>& config)
  828. : Config(config)
  829. , TrackableNodeProcessor(config)
  830. {}
  831. TStringBuf GetName() const override {
  832. return ResultProviderName;
  833. }
  834. bool ValidateParameters(TExprNode& node, TExprContext& ctx, TMaybe<TString>& cluster) override {
  835. if (!EnsureArgsCount(node, 1, ctx)) {
  836. return false;
  837. }
  838. cluster = Nothing();
  839. return true;
  840. }
  841. bool CanParse(const TExprNode& node) override {
  842. return ResultProviderFunctions().contains(node.Content()) || node.Content() == ConfigureName;
  843. }
  844. void FillModifyCallables(THashSet<TStringBuf>& callables) override {
  845. callables.insert(TResWrite::CallableName());
  846. }
  847. IGraphTransformer& GetTypeAnnotationTransformer(bool instantOnly) override {
  848. Y_UNUSED(instantOnly);
  849. if (!TypeAnnotationTransformer) {
  850. TypeAnnotationTransformer = CreateFunctorTransformer(
  851. [&](const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx)->IGraphTransformer::TStatus {
  852. output = input;
  853. if (auto maybeRes = TMaybeNode<TResWriteBase>(input)) {
  854. auto res = maybeRes.Cast();
  855. if (!EnsureArgsCount(*input, res.Maybe<TResWrite>() ? 5 : 6, ctx)) {
  856. return IGraphTransformer::TStatus::Error;
  857. }
  858. if (!EnsureWorldType(*res.Ref().Child(TResWriteBase::idx_World), ctx)) {
  859. return IGraphTransformer::TStatus::Error;
  860. }
  861. if (!EnsureSpecificDataSink(*res.Ref().Child(TResWriteBase::idx_DataSink), ResultProviderName, ctx)) {
  862. return IGraphTransformer::TStatus::Error;
  863. }
  864. if (!res.Ref().Child(TResWriteBase::idx_Key)->IsCallable("Key") || res.Ref().Child(TResWriteBase::idx_Key)->ChildrenSize() > 0) {
  865. ctx.AddError(TIssue(ctx.GetPosition(res.Ref().Child(TResWriteBase::idx_Key)->Pos()), "Expected empty key"));
  866. return IGraphTransformer::TStatus::Error;
  867. }
  868. if (!EnsurePersistable(res.Data().Ref(), ctx)) {
  869. return IGraphTransformer::TStatus::Error;
  870. }
  871. auto settings = res.Ref().Child(TResWriteBase::idx_Settings);
  872. if (!EnsureTuple(*settings, ctx)) {
  873. return IGraphTransformer::TStatus::Error;
  874. }
  875. TExprNode::TPtr columns;
  876. bool hasRef = false;
  877. bool hasAutoRef = false;
  878. ui32 settingPos = 0;
  879. for (auto& setting : settings->Children()) {
  880. if (!EnsureTupleMinSize(*setting, 1, ctx)) {
  881. return IGraphTransformer::TStatus::Error;
  882. }
  883. if (!EnsureAtom(*setting->Child(0), ctx)) {
  884. return IGraphTransformer::TStatus::Error;
  885. }
  886. auto content = setting->Child(0)->Content();
  887. if (content == "ref") {
  888. hasRef = true;
  889. if (!EnsureTupleMaxSize(*setting, 1, ctx)) {
  890. return IGraphTransformer::TStatus::Error;
  891. }
  892. } else if (content == "autoref") {
  893. hasAutoRef = true;
  894. if (!EnsureTupleMaxSize(*setting, 1, ctx)) {
  895. return IGraphTransformer::TStatus::Error;
  896. }
  897. } else if (content == "type") {
  898. if (!EnsureTupleMaxSize(*setting, 1, ctx)) {
  899. return IGraphTransformer::TStatus::Error;
  900. }
  901. } else if (content == "take") {
  902. if (!EnsureTupleMaxSize(*setting, 2, ctx)) {
  903. return IGraphTransformer::TStatus::Error;
  904. }
  905. if (!EnsureAtom(*setting->Child(1), ctx)) {
  906. return IGraphTransformer::TStatus::Error;
  907. }
  908. ui64 limit = 0;
  909. if (!TryFromString(setting->Child(1)->Content(), limit)) {
  910. ctx.AddError(TIssue(ctx.GetPosition(setting->Child(1)->Pos()), "Expected unsigned integer"));
  911. return IGraphTransformer::TStatus::Error;
  912. }
  913. } else if (content == "columns") {
  914. if (columns) {
  915. ctx.AddError(TIssue(ctx.GetPosition(setting->Pos()), "columns is already used"));
  916. return IGraphTransformer::TStatus::Error;
  917. }
  918. if (!EnsureTupleMaxSize(*setting, 2, ctx)) {
  919. return IGraphTransformer::TStatus::Error;
  920. }
  921. columns = setting->ChildPtr(1);
  922. auto status = ValidateColumns(columns, res.Data().Ref().GetTypeAnn(), ctx);
  923. if (status.Level != IGraphTransformer::TStatus::Ok) {
  924. if (status.Level == IGraphTransformer::TStatus::Repeat) {
  925. if (!columns) {
  926. output = ctx.ChangeChild(*input, 4, RemoveSetting(*input->Child(4), "columns", ctx));
  927. } else {
  928. auto newSetting = ctx.ChangeChild(*setting, 1, std::move(columns));
  929. auto newSettings = ctx.ChangeChild(*settings, settingPos, std::move(newSetting));
  930. output = ctx.ChangeChild(*input, 4, std::move(newSettings));
  931. }
  932. }
  933. return status;
  934. }
  935. } else if (content == "freezeColumns") {
  936. if (!EnsureTupleMaxSize(*setting, 1, ctx)) {
  937. return IGraphTransformer::TStatus::Error;
  938. }
  939. } else if (content == "discard") {
  940. if (!EnsureTupleMaxSize(*setting, 1, ctx)) {
  941. return IGraphTransformer::TStatus::Error;
  942. }
  943. } else if (content == "label") {
  944. if (!EnsureTupleMaxSize(*setting, 2, ctx)) {
  945. return IGraphTransformer::TStatus::Error;
  946. }
  947. if (!EnsureAtom(*setting->Child(1), ctx)) {
  948. return IGraphTransformer::TStatus::Error;
  949. }
  950. } else if (content == "unordered") {
  951. if (!EnsureTupleMaxSize(*setting, 1, ctx)) {
  952. return IGraphTransformer::TStatus::Error;
  953. }
  954. } else {
  955. ctx.AddError(TIssue(ctx.GetPosition(setting->Pos()), "Expected label,discard,ref,autoref,type,unordered,take or columns atom"));
  956. return IGraphTransformer::TStatus::Error;
  957. }
  958. if (hasRef && hasAutoRef) {
  959. ctx.AddError(TIssue(ctx.GetPosition(setting->Pos()), "Either ref or autoref may be specified, not both"));
  960. return IGraphTransformer::TStatus::Error;
  961. }
  962. ++settingPos;
  963. }
  964. if (auto right = res.Data().Maybe<TCoRight>()) {
  965. auto source = right.Cast().Input();
  966. if (!source.Maybe<TCoCons>()) {
  967. const TIntrusivePtr<IDataProvider>* provider = nullptr;
  968. if (source.Ref().Type() == TExprNode::Callable || source.Ref().ChildrenSize() >= 2) {
  969. if (source.Ref().Child(1)->IsCallable("DataSource")) {
  970. auto name = source.Ref().Child(1)->Child(0)->Content();
  971. provider = Config->Types.DataSourceMap.FindPtr(name);
  972. Y_ENSURE(provider, "DataSource doesn't exist: " << name);
  973. }
  974. if (source.Ref().Child(1)->IsCallable("DataSink")) {
  975. auto name = source.Ref().Child(1)->Child(0)->Content();
  976. provider = Config->Types.DataSinkMap.FindPtr(name);
  977. Y_ENSURE(provider, "DataSink doesn't exist: " << name);
  978. }
  979. }
  980. if (!provider) {
  981. ctx.AddError(TIssue(ctx.GetPosition(res.Data().Pos()), "Expected Right! over Datasource or Datasink"));
  982. return IGraphTransformer::TStatus::Error;
  983. }
  984. }
  985. }
  986. if (res.Maybe<TResTransientBase>()) {
  987. auto resTransient = res.Cast<TResTransientBase>();
  988. if (!EnsureAtom(*resTransient.Ref().Child(TResTransientBase::idx_DelegatedSource), ctx)) {
  989. return IGraphTransformer::TStatus::Error;
  990. }
  991. if (!Config->Types.DataSourceMap.FindPtr(resTransient.DelegatedSource().Value())) {
  992. ctx.AddError(TIssue(ctx.GetPosition(resTransient.DelegatedSource().Pos()),
  993. TStringBuilder() << "DataSource is not found: " << resTransient.DelegatedSource().Value()));
  994. return IGraphTransformer::TStatus::Error;
  995. }
  996. }
  997. if (res.Data().Ref().IsCallable("AssumeColumnOrder")) {
  998. if (!HasSetting(res.Settings().Ref(), "freezeColumns")) {
  999. auto dataOrder = Config->Types.LookupColumnOrder(res.Data().Ref());
  1000. YQL_ENSURE(dataOrder);
  1001. YQL_CLOG(INFO, ProviderResult) << "Setting result column order: " << FormatColumnOrder(dataOrder);
  1002. auto settings = RemoveSetting(res.Settings().Ref(), "columns", ctx);
  1003. TExprNodeList columnsList;
  1004. for (auto& [col, gen_col] : *dataOrder) {
  1005. columnsList.push_back(ctx.NewAtom(settings->Pos(), col));
  1006. }
  1007. settings = AddSetting(*settings, settings->Pos(), "columns", ctx.NewList(settings->Pos(), std::move(columnsList)), ctx);
  1008. settings = AddSetting(*settings, settings->Pos(), "freezeColumns", nullptr, ctx);
  1009. output = ctx.ChangeChild(*input, TResWriteBase::idx_Settings, std::move(settings));
  1010. }
  1011. output = ctx.ChangeChild(*output, TResWriteBase::idx_Data, res.Data().Ref().HeadPtr());
  1012. return IGraphTransformer::TStatus::Repeat;
  1013. }
  1014. input->SetTypeAnn(res.World().Ref().GetTypeAnn());
  1015. return IGraphTransformer::TStatus::Ok;
  1016. }
  1017. else if (auto maybeIf = TMaybeNode<TResIf>(input)) {
  1018. if (!EnsureArgsCount(*input, 7, ctx)) {
  1019. return IGraphTransformer::TStatus::Error;
  1020. }
  1021. if (!EnsureWorldType(*input->Child(TResIf::idx_World), ctx)) {
  1022. return IGraphTransformer::TStatus::Error;
  1023. }
  1024. if (!EnsureSpecificDataSink(*input->Child(TResIf::idx_DataSink), ResultProviderName, ctx)) {
  1025. return IGraphTransformer::TStatus::Error;
  1026. }
  1027. if (!EnsureSpecificDataType(*input->Child(TResIf::idx_Condition), EDataSlot::Bool, ctx)) {
  1028. return IGraphTransformer::TStatus::Error;
  1029. }
  1030. if (!EnsureWorldType(*input->Child(TResIf::idx_Then), ctx)) {
  1031. return IGraphTransformer::TStatus::Error;
  1032. }
  1033. if (!EnsureWorldType(*input->Child(TResIf::idx_Else), ctx)) {
  1034. return IGraphTransformer::TStatus::Error;
  1035. }
  1036. if (!EnsureAtom(*input->Child(TResIf::idx_DelegatedSource), ctx)) {
  1037. return IGraphTransformer::TStatus::Error;
  1038. }
  1039. if (!Config->Types.DataSourceMap.FindPtr(input->Child(TResIf::idx_DelegatedSource)->Content())) {
  1040. ctx.AddError(TIssue(ctx.GetPosition(input->Child(TResIf::idx_DelegatedSource)->Pos()),
  1041. TStringBuilder() << "DataSource is not found: " << input->Child(TResIf::idx_DelegatedSource)->Content()));
  1042. return IGraphTransformer::TStatus::Error;
  1043. }
  1044. auto settings = input->Child(TResIf::idx_Settings);
  1045. if (!EnsureTupleSize(*settings, 0, ctx)) {
  1046. return IGraphTransformer::TStatus::Error;
  1047. }
  1048. input->SetTypeAnn(input->Child(TResIf::idx_World)->GetTypeAnn());
  1049. return IGraphTransformer::TStatus::Ok;
  1050. }
  1051. else if (auto maybeFor = TMaybeNode<TResFor>(input)) {
  1052. if (!EnsureArgsCount(*input, 10, ctx)) {
  1053. return IGraphTransformer::TStatus::Error;
  1054. }
  1055. if (!EnsureWorldType(*input->Child(TResFor::idx_World), ctx)) {
  1056. return IGraphTransformer::TStatus::Error;
  1057. }
  1058. if (!EnsureSpecificDataSink(*input->Child(TResFor::idx_DataSink), ResultProviderName, ctx)) {
  1059. return IGraphTransformer::TStatus::Error;
  1060. }
  1061. if (!EnsurePersistable(*input->Child(TResFor::idx_Items), ctx)) {
  1062. return IGraphTransformer::TStatus::Error;
  1063. }
  1064. auto listType = RemoveOptionalType(input->Child(TResFor::idx_Items)->GetTypeAnn());
  1065. if (!EnsureListType(input->Child(TResFor::idx_Items)->Pos(), *listType, ctx)) {
  1066. return IGraphTransformer::TStatus::Error;
  1067. }
  1068. auto itemType = listType->Cast<TListExprType>()->GetItemType();
  1069. auto status = ConvertToLambda(input->ChildRef(TResFor::idx_Iter), ctx, 2);
  1070. if (status.Level != IGraphTransformer::TStatus::Ok) {
  1071. return status;
  1072. }
  1073. auto& lambda1 = input->ChildRef(TResFor::idx_Iter);
  1074. if (!UpdateLambdaAllArgumentsTypes(lambda1, { ctx.MakeType<TWorldExprType>(), itemType }, ctx)) {
  1075. return IGraphTransformer::TStatus::Error;
  1076. }
  1077. if (!lambda1->GetTypeAnn()) {
  1078. return IGraphTransformer::TStatus::Repeat;
  1079. }
  1080. if (!EnsureWorldType(*lambda1->Child(1), ctx)) {
  1081. return IGraphTransformer::TStatus::Error;
  1082. }
  1083. status = ConvertToLambda(input->ChildRef(TResFor::idx_Else), ctx, 1);
  1084. if (status.Level != IGraphTransformer::TStatus::Ok) {
  1085. return status;
  1086. }
  1087. auto& lambda2 = input->ChildRef(TResFor::idx_Else);
  1088. if (!UpdateLambdaAllArgumentsTypes(lambda2, { ctx.MakeType<TWorldExprType>() }, ctx)) {
  1089. return IGraphTransformer::TStatus::Error;
  1090. }
  1091. if (!lambda2->GetTypeAnn()) {
  1092. return IGraphTransformer::TStatus::Repeat;
  1093. }
  1094. if (!EnsureWorldType(*lambda2->Child(1), ctx)) {
  1095. return IGraphTransformer::TStatus::Error;
  1096. }
  1097. if (!EnsureAtom(*input->Child(TResFor::idx_DelegatedSource), ctx)) {
  1098. return IGraphTransformer::TStatus::Error;
  1099. }
  1100. if (!Config->Types.DataSourceMap.FindPtr(input->Child(TResFor::idx_DelegatedSource)->Content())) {
  1101. ctx.AddError(TIssue(ctx.GetPosition(input->Child(TResFor::idx_DelegatedSource)->Pos()),
  1102. TStringBuilder() << "DataSource is not found: " << input->Child(TResFor::idx_DelegatedSource)->Content()));
  1103. return IGraphTransformer::TStatus::Error;
  1104. }
  1105. auto settings = input->Child(TResFor::idx_Settings);
  1106. if (!EnsureTupleSize(*settings, 0, ctx)) {
  1107. return IGraphTransformer::TStatus::Error;
  1108. }
  1109. if (!EnsureAtom(*input->Child(TResFor::idx_Total), ctx)) {
  1110. return IGraphTransformer::TStatus::Error;
  1111. }
  1112. auto total = input->Child(TResFor::idx_Total)->Content();
  1113. ui32 totalValue = 0;
  1114. if (total && !TryFromString(total, totalValue)) {
  1115. ctx.AddError(TIssue(ctx.GetPosition(input->Child(TResFor::idx_Total)->Pos()),
  1116. TStringBuilder() << "Expected number, but got: " << total));
  1117. return IGraphTransformer::TStatus::Error;
  1118. }
  1119. if (!EnsureAtom(*input->Child(TResFor::idx_Current), ctx)) {
  1120. return IGraphTransformer::TStatus::Error;
  1121. }
  1122. auto current = input->Child(TResFor::idx_Current)->Content();
  1123. ui32 currentValue = 0;
  1124. if (current && !TryFromString(current, currentValue)) {
  1125. ctx.AddError(TIssue(ctx.GetPosition(input->Child(TResFor::idx_Current)->Pos()),
  1126. TStringBuilder() << "Expected number, but got: " << current));
  1127. return IGraphTransformer::TStatus::Error;
  1128. }
  1129. if (!total != !current) {
  1130. ctx.AddError(TIssue(ctx.GetPosition(input->Child(TResFor::idx_Current)->Pos()),
  1131. TStringBuilder() << "Current value should be set simultaneously with total value"));
  1132. return IGraphTransformer::TStatus::Error;
  1133. }
  1134. const auto& active = *input->Child(TResFor::idx_Active);
  1135. if (!active.IsCallable("Void")) {
  1136. if (!EnsureWorldType(active, ctx)) {
  1137. return IGraphTransformer::TStatus::Error;
  1138. }
  1139. }
  1140. input->SetTypeAnn(input->Child(TResIf::idx_World)->GetTypeAnn());
  1141. return IGraphTransformer::TStatus::Ok;
  1142. }
  1143. else if (auto maybeCommit = TMaybeNode<TCoCommit>(input)) {
  1144. auto commit = maybeCommit.Cast();
  1145. auto settings = NCommon::ParseCommitSettings(commit, ctx);
  1146. if (!settings.EnsureModeEmpty(ctx)) {
  1147. return IGraphTransformer::TStatus::Error;
  1148. }
  1149. if (!settings.EnsureEpochEmpty(ctx)) {
  1150. return IGraphTransformer::TStatus::Error;
  1151. }
  1152. if (!settings.EnsureOtherEmpty(ctx)) {
  1153. return IGraphTransformer::TStatus::Error;
  1154. }
  1155. input->SetTypeAnn(commit.World().Ref().GetTypeAnn());
  1156. return IGraphTransformer::TStatus::Ok;
  1157. }
  1158. else if (input->Content() == ConfigureName) {
  1159. if (!EnsureMinArgsCount(*input, 3, ctx)) {
  1160. return IGraphTransformer::TStatus::Error;
  1161. }
  1162. if (!EnsureWorldType(*input->Child(0), ctx)) {
  1163. return IGraphTransformer::TStatus::Error;
  1164. }
  1165. if (!EnsureSpecificDataSink(*input->Child(1), ResultProviderName, ctx)) {
  1166. return IGraphTransformer::TStatus::Error;
  1167. }
  1168. if (!EnsureAtom(*input->Child(2), ctx)) {
  1169. return IGraphTransformer::TStatus::Error;
  1170. }
  1171. auto command = input->Child(2)->Content();
  1172. if (command == "SizeLimit") {
  1173. if (!EnsureArgsCount(*input, 4, ctx)) {
  1174. return IGraphTransformer::TStatus::Error;
  1175. }
  1176. if (!EnsureAtom(*input->Child(3), ctx)) {
  1177. return IGraphTransformer::TStatus::Error;
  1178. }
  1179. const auto limitStr = input->Child(3)->Content();
  1180. ui64 limit;
  1181. if (!TryFromString(limitStr, limit)) {
  1182. ctx.AddError(TIssue(ctx.GetPosition(input->Child(3)->Pos()), TStringBuilder() << "expected integer, but got: " << limitStr));
  1183. return IGraphTransformer::TStatus::Error;
  1184. }
  1185. if (Config->FillSettings.AllResultsBytesLimit) {
  1186. Config->FillSettings.AllResultsBytesLimit = Min(*Config->FillSettings.AllResultsBytesLimit, limit);
  1187. } else {
  1188. Config->FillSettings.AllResultsBytesLimit = limit;
  1189. }
  1190. } else {
  1191. ctx.AddError(TIssue(ctx.GetPosition(input->Child(2)->Pos()), TStringBuilder() << "Unsupported command: " << command));
  1192. return IGraphTransformer::TStatus::Error;
  1193. }
  1194. input->SetTypeAnn(input->Child(0)->GetTypeAnn());
  1195. return IGraphTransformer::TStatus::Ok;
  1196. }
  1197. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << "(Result) Unsupported function: " << input->Content()));
  1198. return IGraphTransformer::TStatus::Error;
  1199. });
  1200. }
  1201. return *TypeAnnotationTransformer;
  1202. }
  1203. TExprNode::TPtr RewriteIO(const TExprNode::TPtr& node, TExprContext& ctx) override {
  1204. auto ret = node;
  1205. if (node->Content() == WriteName) {
  1206. ret = ctx.RenameNode(*ret, TResWrite::CallableName());
  1207. ret = ctx.ChangeChild(*ret, TResWrite::idx_Data,
  1208. ctx.Builder(node->Pos())
  1209. .Callable("RemovePrefixMembers")
  1210. .Add(0, node->ChildPtr(TResWrite::idx_Data))
  1211. .List(1)
  1212. .Atom(0, "_yql_sys_", TNodeFlags::Default)
  1213. .Seal()
  1214. .Seal()
  1215. .Build()
  1216. );
  1217. }
  1218. else {
  1219. YQL_ENSURE(false, "Expected Write!");
  1220. }
  1221. YQL_CLOG(INFO, ProviderResult) << "RewriteIO";
  1222. return ret;
  1223. }
  1224. IGraphTransformer& GetPhysicalFinalizingTransformer() override {
  1225. if (!PhysicalFinalizingTransformer) {
  1226. PhysicalFinalizingTransformer = new TPhysicalFinalizingTransformer(Config);
  1227. }
  1228. return *PhysicalFinalizingTransformer;
  1229. }
  1230. bool CanExecute(const TExprNode& node) override {
  1231. if (node.Content() == TResFill::CallableName()) {
  1232. return true;
  1233. }
  1234. if (node.Content() == TResPull::CallableName()) {
  1235. return true;
  1236. }
  1237. if (node.Content() == TResIf::CallableName()) {
  1238. return true;
  1239. }
  1240. if (node.Content() == TResFor::CallableName()) {
  1241. return true;
  1242. }
  1243. if (node.Content() == ConfigureName) {
  1244. return true;
  1245. }
  1246. return false;
  1247. }
  1248. bool ValidateExecution(const TExprNode& node, TExprContext& ctx) override {
  1249. auto getDataProvider = [&]() {
  1250. auto provider = Config->Types.DataSourceMap.FindPtr(node.Child(5)->Content());
  1251. Y_ENSURE(provider, "DataSource doesn't exist: " << node.Child(5)->Content());
  1252. return *provider;
  1253. };
  1254. if (TResTransientBase::Match(&node)) {
  1255. return getDataProvider()->ValidateExecution(TResTransientBase(&node).Data().Ref(), ctx);
  1256. }
  1257. if (TResIf::Match(&node)) {
  1258. return getDataProvider()->ValidateExecution(TResIf(&node).Condition().Ref(), ctx);
  1259. }
  1260. if (TResFor::Match(&node)) {
  1261. return getDataProvider()->ValidateExecution(TResFor(&node).Items().Ref(), ctx);
  1262. }
  1263. return true;
  1264. }
  1265. IGraphTransformer& GetCallableExecutionTransformer() override {
  1266. if (!CallableExecutionTransformer) {
  1267. CallableExecutionTransformer = new TResultCallableExecutionTransformer(Config);
  1268. }
  1269. return *CallableExecutionTransformer;
  1270. }
  1271. void Reset() final {
  1272. TDataProviderBase::Reset();
  1273. if (CallableExecutionTransformer) {
  1274. CallableExecutionTransformer.Reset();
  1275. }
  1276. Config->CommittedResults.clear();
  1277. }
  1278. bool GetDependencies(const TExprNode& node, TExprNode::TListType& children, bool compact) override {
  1279. if (CanExecute(node)) {
  1280. children.push_back(node.ChildPtr(0));
  1281. if (auto resPull = TMaybeNode<TResPull>(&node)) {
  1282. children.push_back(resPull.Cast().Data().Ptr());
  1283. } else if (auto resIf = TMaybeNode<TResIf>(&node)) {
  1284. children.push_back(resIf.Cast().Then().Ptr());
  1285. children.push_back(resIf.Cast().Else().Ptr());
  1286. } else if (auto resFor = TMaybeNode<TResFor>(&node)) {
  1287. auto active = resFor.Cast().Active().Ptr();
  1288. if (!active->IsCallable("Void")) {
  1289. children.push_back(active);
  1290. }
  1291. } else if (auto resFill = TMaybeNode<TResFill>(&node)) {
  1292. const auto provider = Config->Types.DataSourceMap.FindPtr(resFill.Cast().DelegatedSource().Value());
  1293. Y_ENSURE(provider, "DataSource not exist: " << resFill.Cast().DelegatedSource().Value());
  1294. (*provider)->GetPlanFormatter().GetResultDependencies(resFill.Cast().Data().Ptr(), children, compact);
  1295. }
  1296. return true;
  1297. }
  1298. return false;
  1299. }
  1300. void WritePlanDetails(const TExprNode& node, NYson::TYsonWriter& writer, bool withLimits) override {
  1301. Y_UNUSED(withLimits);
  1302. if (auto resPull = TMaybeNode<TResPull>(&node)) {
  1303. auto dataSourceName = resPull.Cast().DelegatedSource().Value();
  1304. auto dataSource = Config->Types.DataSourceMap.FindPtr(dataSourceName);
  1305. YQL_ENSURE(dataSource);
  1306. (*dataSource)->GetPlanFormatter().WritePullDetails(resPull.Cast().Data().Ref(), writer);
  1307. }
  1308. }
  1309. TString GetProviderPath(const TExprNode& node) override {
  1310. Y_UNUSED(node);
  1311. return "result";
  1312. }
  1313. TString GetOperationDisplayName(const TExprNode& node) override {
  1314. if (node.Content() == CommitName) {
  1315. return TString::Join(node.Content(), " on result");
  1316. }
  1317. if (auto maybeResFor = TMaybeNode<TResFor>(&node)) {
  1318. auto resFor = maybeResFor.Cast();
  1319. TStringBuilder res;
  1320. res << node.Content();
  1321. if (resFor.Total().Value() && resFor.Total().Value() != "0") {
  1322. res << ", " << (node.GetState() == TExprNode::EState::ExecutionComplete ?
  1323. resFor.Total().Value() : resFor.Current().Value()) << "/" << resFor.Total().Value();
  1324. }
  1325. return res;
  1326. }
  1327. return TString(node.Content());
  1328. }
  1329. ITrackableNodeProcessor& GetTrackableNodeProcessor() override {
  1330. return TrackableNodeProcessor;
  1331. }
  1332. private:
  1333. const TIntrusivePtr<TResultProviderConfig> Config;
  1334. TResultTrackableNodeProcessor TrackableNodeProcessor;
  1335. TAutoPtr<IGraphTransformer> TypeAnnotationTransformer;
  1336. TAutoPtr<IGraphTransformer> PhysicalFinalizingTransformer;
  1337. TAutoPtr<IGraphTransformer> CallableExecutionTransformer;
  1338. };
  1339. }
  1340. TIntrusivePtr<IResultWriter> CreateYsonResultWriter(NYson::EYsonFormat format) {
  1341. return MakeIntrusive<TYsonResultWriter>(format);
  1342. }
  1343. TIntrusivePtr<IDataProvider> CreateResultProvider(const TIntrusivePtr<TResultProviderConfig>& config) {
  1344. return new TResultProvider(config);
  1345. }
  1346. const THashSet<TStringBuf>& ResultProviderFunctions() {
  1347. return Singleton<TResultProvider::TFunctions>()->Names;
  1348. }
  1349. }