yql_expr_constraint.cpp 191 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714
  1. #include "yql_expr_constraint.h"
  2. #include "yql_callable_transform.h"
  3. #include "yql_opt_utils.h"
  4. #include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
  5. #include <yql/essentials/core/yql_type_helpers.h>
  6. #include <yql/essentials/core/yql_join.h>
  7. #include <yql/essentials/utils/log/profile.h>
  8. #include <library/cpp/yson/node/node_io.h>
  9. #include <util/generic/scope.h>
  10. #include <util/generic/utility.h>
  11. #include <util/generic/algorithm.h>
  12. #include <util/string/builder.h>
  13. #include <util/string/type.h>
  14. namespace NYql {
  15. using namespace NNodes;
  16. namespace {
  17. template <size_t FromChild, class... Other>
  18. struct TApplyConstraintFromInput;
  19. template <size_t FromChild>
  20. struct TApplyConstraintFromInput<FromChild> {
  21. static void Do(const TExprNode::TPtr&) {
  22. }
  23. };
  24. template <size_t FromChild, class TConstraint, class... Other>
  25. struct TApplyConstraintFromInput<FromChild, TConstraint, Other...> {
  26. static void Do(const TExprNode::TPtr& input) {
  27. if (auto c = input->Child(FromChild)->GetConstraint<TConstraint>()) {
  28. input->AddConstraint(c);
  29. }
  30. TApplyConstraintFromInput<FromChild, Other...>::Do(input);
  31. }
  32. };
  33. template <class TConstraint>
  34. const TConstraint* MakeCommonConstraint(const TExprNode::TPtr& input, size_t from, TExprContext& ctx) {
  35. TVector<const TConstraintSet*> constraints;
  36. for (size_t i = from; i < input->ChildrenSize(); ++i) {
  37. constraints.push_back(&input->Child(i)->GetConstraintSet());
  38. }
  39. return TConstraint::MakeCommon(constraints, ctx);
  40. }
  41. template <class... Other>
  42. struct TApplyCommonConstraint;
  43. template <class TConstraint>
  44. struct TApplyCommonConstraint<TConstraint> {
  45. static void Do(const TExprNode::TPtr& input, const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx) {
  46. if (auto c = TConstraint::MakeCommon(constraints, ctx)) {
  47. input->AddConstraint(c);
  48. }
  49. }
  50. };
  51. template <class TConstraint, class... Other>
  52. struct TApplyCommonConstraint<TConstraint, Other...> {
  53. static void Do(const TExprNode::TPtr& input, const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx) {
  54. if (auto c = TConstraint::MakeCommon(constraints, ctx)) {
  55. input->AddConstraint(c);
  56. }
  57. TApplyCommonConstraint<Other...>::Do(input, constraints, ctx);
  58. }
  59. };
  60. class TCallableConstraintTransformer : public TCallableTransformerBase<TCallableConstraintTransformer> {
  61. using THandler = TStatus(TCallableConstraintTransformer::*)(const TExprNode::TPtr&, TExprNode::TPtr&, TExprContext&) const;
  62. public:
  63. TCallableConstraintTransformer(TTypeAnnotationContext& types, bool instantOnly, bool subGraph)
  64. : TCallableTransformerBase<TCallableConstraintTransformer>(types, instantOnly)
  65. , SubGraph(subGraph)
  66. {
  67. Functions["FailMe"] = &TCallableConstraintTransformer::FailMe;
  68. Functions["Unordered"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  69. Functions["UnorderedSubquery"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  70. Functions["Sort"] = &TCallableConstraintTransformer::SortWrap;
  71. Functions["AssumeSorted"] = &TCallableConstraintTransformer::SortWrap;
  72. Functions["AssumeUnique"] = &TCallableConstraintTransformer::AssumeUniqueWrap<false, true>;
  73. Functions["AssumeDistinct"] = &TCallableConstraintTransformer::AssumeUniqueWrap<true, true>;
  74. Functions["AssumeUniqueHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap<false, false>;
  75. Functions["AssumeDistinctHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap<true, false>;
  76. Functions["AssumeConstraints"] = &TCallableConstraintTransformer::AssumeConstraintsWrap;
  77. Functions["AssumeChopped"] = &TCallableConstraintTransformer::AssumeChoppedWrap;
  78. Functions["AssumeColumnOrder"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  79. Functions["AssumeAllMembersNullableAtOnce"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  80. Functions["Top"] = &TCallableConstraintTransformer::TopWrap<false>;
  81. Functions["TopSort"] = &TCallableConstraintTransformer::TopWrap<true>;
  82. Functions["TakeWhile"] = &TCallableConstraintTransformer::FilterWrap<true>;
  83. Functions["SkipWhile"] = &TCallableConstraintTransformer::FilterWrap<true>;
  84. Functions["TakeWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap<true>;
  85. Functions["SkipWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap<true>;
  86. Functions["WideTakeWhile"] = &TCallableConstraintTransformer::FilterWrap<true>;
  87. Functions["WideSkipWhile"] = &TCallableConstraintTransformer::FilterWrap<true>;
  88. Functions["WideTakeWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap<true>;
  89. Functions["WideSkipWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap<true>;
  90. Functions["Iterator"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  91. Functions["ForwardList"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  92. Functions["LazyList"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  93. Functions["TableSource"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  94. Functions["WideTableSource"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  95. Functions["ToFlow"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  96. Functions["FromFlow"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  97. Functions["ToStream"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  98. Functions["ToSequence"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  99. Functions["Collect"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  100. Functions["FilterNullMembers"] = &TCallableConstraintTransformer::FromFirst<TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode>;
  101. Functions["SkipNullMembers"] = &TCallableConstraintTransformer::FromFirst<TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode>;
  102. Functions["FilterNullElements"] = &TCallableConstraintTransformer::FromFirst<TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode>;
  103. Functions["SkipNullElements"] = &TCallableConstraintTransformer::FromFirst<TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode>;
  104. Functions["Right!"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  105. Functions["Cons!"] = &TCallableConstraintTransformer::CopyAllFrom<1>;
  106. Functions["ExtractMembers"] = &TCallableConstraintTransformer::ExtractMembersWrap;
  107. Functions["RemoveSystemMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap;
  108. Functions["RemovePrefixMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap;
  109. Functions["FlattenMembers"] = &TCallableConstraintTransformer::FlattenMembersWrap;
  110. Functions["SelectMembers"] = &TCallableConstraintTransformer::SelectMembersWrap;
  111. Functions["FilterMembers"] = &TCallableConstraintTransformer::SelectMembersWrap;
  112. Functions["CastStruct"] = &TCallableConstraintTransformer::SelectMembersWrap;
  113. Functions["SafeCast"] = &TCallableConstraintTransformer::CastWrap<false>;
  114. Functions["StrictCast"] = &TCallableConstraintTransformer::CastWrap<true>;
  115. Functions["ToString"] = &TCallableConstraintTransformer::CastWrap<true>;
  116. Functions["ToBytes"] = &TCallableConstraintTransformer::CastWrap<true>;
  117. Functions["DivePrefixMembers"] = &TCallableConstraintTransformer::DivePrefixMembersWrap;
  118. Functions["OrderedFilter"] = &TCallableConstraintTransformer::FilterWrap<true>;
  119. Functions["Filter"] = &TCallableConstraintTransformer::FilterWrap<false>;
  120. Functions["WideFilter"] = &TCallableConstraintTransformer::FilterWrap<true>;
  121. Functions["OrderedMap"] = &TCallableConstraintTransformer::MapWrap<true, false>;
  122. Functions["Map"] = &TCallableConstraintTransformer::MapWrap<false, false>;
  123. Functions["MapNext"] = &TCallableConstraintTransformer::MapWrap<true, false>;
  124. Functions["OrderedFlatMap"] = &TCallableConstraintTransformer::MapWrap<true, true>;
  125. Functions["FlatMap"] = &TCallableConstraintTransformer::MapWrap<false, true>;
  126. Functions["OrderedMultiMap"] = &TCallableConstraintTransformer::MapWrap<true, false>;
  127. Functions["MultiMap"] = &TCallableConstraintTransformer::MapWrap<false, false>;
  128. Functions["ExpandMap"] = &TCallableConstraintTransformer::MapWrap<true, false, false, true>;
  129. Functions["WideMap"] = &TCallableConstraintTransformer::MapWrap<true, false, true, true>;
  130. Functions["NarrowMap"] = &TCallableConstraintTransformer::MapWrap<true, false, true, false>;
  131. Functions["NarrowFlatMap"] = &TCallableConstraintTransformer::MapWrap<true, true, true, false>;
  132. Functions["NarrowMultiMap"] = &TCallableConstraintTransformer::MapWrap<true, false, true, false>;
  133. Functions["OrderedFlatMapToEquiJoin"] = &TCallableConstraintTransformer::MapWrap<true, true>;
  134. Functions["FlatMapToEquiJoin"] = &TCallableConstraintTransformer::MapWrap<false, true>;
  135. Functions["OrderedLMap"] = &TCallableConstraintTransformer::LMapWrap<true>;
  136. Functions["LMap"] = &TCallableConstraintTransformer::LMapWrap<false>;
  137. Functions["Extract"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode>;
  138. Functions["OrderedExtract"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode>;
  139. Functions["OrderedExtend"] = &TCallableConstraintTransformer::ExtendWrap<true>;
  140. Functions["Extend"] = &TCallableConstraintTransformer::ExtendWrap<false>;
  141. Functions["UnionAll"] = &TCallableConstraintTransformer::ExtendWrap<false>;
  142. Functions["Merge"] = &TCallableConstraintTransformer::MergeWrap<false>;
  143. Functions["UnionMerge"] = &TCallableConstraintTransformer::MergeWrap<true>;
  144. Functions["Skip"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  145. Functions["Take"] = &TCallableConstraintTransformer::TakeWrap;
  146. Functions["Limit"] = &TCallableConstraintTransformer::TakeWrap;
  147. Functions["Member"] = &TCallableConstraintTransformer::MemberWrap;
  148. Functions["AsStruct"] = &TCallableConstraintTransformer::AsStructWrap;
  149. Functions["BlockAsStruct"] = &TCallableConstraintTransformer::AsStructWrap;
  150. Functions["Just"] = &TCallableConstraintTransformer::FromFirst<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  151. Functions["Unwrap"] = &TCallableConstraintTransformer::FromFirst<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  152. Functions["Ensure"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  153. Functions["ToList"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  154. Functions["ToOptional"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  155. Functions["Head"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  156. Functions["Last"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  157. Functions["Reverse"] = &TCallableConstraintTransformer::ReverseWrap;
  158. Functions["Replicate"] = &TCallableConstraintTransformer::FromFirst<TVarIndexConstraintNode, TMultiConstraintNode>;
  159. Functions["AddMember"] = &TCallableConstraintTransformer::AddMemberWrap;
  160. Functions["RemoveMember"] = &TCallableConstraintTransformer::RemoveMemberWrap;
  161. Functions["ForceRemoveMember"] = &TCallableConstraintTransformer::RemoveMemberWrap;
  162. Functions["ReplaceMember"] = &TCallableConstraintTransformer::ReplaceMemberWrap;
  163. Functions["AsList"] = &TCallableConstraintTransformer::AsListWrap;
  164. Functions["OptionalIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<false, false>;
  165. Functions["FlatOptionalIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<false, true>;
  166. Functions["ListIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<true, false>;
  167. Functions["FlatListIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<true, true>;
  168. Functions["EmptyIterator"] = &TCallableConstraintTransformer::FromEmpty;
  169. Functions["EmptyFrom"] = &TCallableConstraintTransformer::EmptyFromWrap;
  170. Functions["List"] = &TCallableConstraintTransformer::ListWrap;
  171. Functions["Dict"] = &TCallableConstraintTransformer::DictWrap;
  172. Functions["EmptyList"] = &TCallableConstraintTransformer::FromEmpty;
  173. Functions["EmptyDict"] = &TCallableConstraintTransformer::FromEmpty;
  174. Functions["DictFromKeys"] = &TCallableConstraintTransformer::DictFromKeysWrap;
  175. Functions["If"] = &TCallableConstraintTransformer::IfWrap;
  176. Functions["Nothing"] = &TCallableConstraintTransformer::FromEmpty;
  177. Functions["IfPresent"] = &TCallableConstraintTransformer::IfPresentWrap;
  178. Functions["Coalesce"] = &TCallableConstraintTransformer::CommonFromChildren<0, TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
  179. Functions["CombineByKey"] = &TCallableConstraintTransformer::FromFinalLambda<TCoCombineByKey::idx_FinishHandlerLambda>;
  180. Functions["FinalizeByKey"] = &TCallableConstraintTransformer::FromFinalLambda<TCoFinalizeByKey::idx_FinishHandlerLambda>;
  181. Functions["CombineCore"] = &TCallableConstraintTransformer::FromFinalLambda<TCoCombineCore::idx_FinishHandler>;
  182. Functions["PartitionByKey"] = &TCallableConstraintTransformer::ShuffleByKeysWrap<true>;
  183. Functions["PartitionsByKeys"] = &TCallableConstraintTransformer::ShuffleByKeysWrap<true>;
  184. Functions["ShuffleByKeys"] = &TCallableConstraintTransformer::ShuffleByKeysWrap<false>;
  185. Functions["Switch"] = &TCallableConstraintTransformer::SwitchWrap;
  186. Functions["Visit"] = &TCallableConstraintTransformer::VisitWrap;
  187. Functions["VariantItem"] = &TCallableConstraintTransformer::VariantItemWrap;
  188. Functions["Variant"] = &TCallableConstraintTransformer::VariantWrap;
  189. Functions["DynamicVariant"] = &TCallableConstraintTransformer::DynamicVariantWrap;
  190. Functions["Guess"] = &TCallableConstraintTransformer::GuessWrap;
  191. Functions["Mux"] = &TCallableConstraintTransformer::MuxWrap;
  192. Functions["Nth"] = &TCallableConstraintTransformer::NthWrap;
  193. Functions["EquiJoin"] = &TCallableConstraintTransformer::EquiJoinWrap;
  194. Functions["JoinDict"] = &TCallableConstraintTransformer::JoinDictWrap;
  195. Functions["MapJoinCore"] = &TCallableConstraintTransformer::MapJoinCoreWrap;
  196. Functions["GraceJoinCore"] = &TCallableConstraintTransformer::GraceJoinCoreWrap;
  197. Functions["GraceSelfJoinCore"] = &TCallableConstraintTransformer::GraceSelfJoinCoreWrap;
  198. Functions["CommonJoinCore"] = &TCallableConstraintTransformer::FromFirst<TEmptyConstraintNode>;
  199. Functions["ToDict"] = &TCallableConstraintTransformer::ToDictWrap;
  200. Functions["DictItems"] = &TCallableConstraintTransformer::DictItemsWrap;
  201. Functions["DictKeys"] = &TCallableConstraintTransformer::DictHalfWrap<true>;
  202. Functions["DictPayloads"] = &TCallableConstraintTransformer::DictHalfWrap<false>;
  203. Functions["Chain1Map"] = &TCallableConstraintTransformer::Chain1MapWrap<false>;
  204. Functions["WideChain1Map"] = &TCallableConstraintTransformer::Chain1MapWrap<true>;
  205. Functions["IsKeySwitch"] = &TCallableConstraintTransformer::IsKeySwitchWrap;
  206. Functions["Condense"] = &TCallableConstraintTransformer::CondenseWrap;
  207. Functions["Condense1"] = &TCallableConstraintTransformer::Condense1Wrap<false>;
  208. Functions["GroupingCore"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
  209. Functions["Chopper"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
  210. Functions["WideChopper"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
  211. Functions["WideCombiner"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
  212. Functions["WideCondense1"] = &TCallableConstraintTransformer::Condense1Wrap<true>;
  213. Functions["Aggregate"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  214. Functions["AggregateMergeState"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  215. Functions["AggregateMergeFinalize"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  216. Functions["AggregateMergeManyFinalize"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  217. Functions["AggregateFinalize"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  218. Functions["AggregateCombine"] = &TCallableConstraintTransformer::AggregateWrap<false>;
  219. Functions["AggregateCombineState"] = &TCallableConstraintTransformer::AggregateWrap<false>;
  220. Functions["Fold"] = &TCallableConstraintTransformer::FoldWrap;
  221. Functions["Fold1"] = &TCallableConstraintTransformer::FoldWrap;
  222. Functions["WithContext"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  223. Functions["WithWorld"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  224. Functions["WideTop"] = &TCallableConstraintTransformer::WideTopWrap<false>;
  225. Functions["WideTopSort"] = &TCallableConstraintTransformer::WideTopWrap<true>;
  226. Functions["WideSort"] = &TCallableConstraintTransformer::WideTopWrap<true>;
  227. Functions["WideTopBlocks"] = &TCallableConstraintTransformer::WideTopWrap<false>;
  228. Functions["WideTopSortBlocks"] = &TCallableConstraintTransformer::WideTopWrap<true>;
  229. Functions["WideSortBlocks"] = &TCallableConstraintTransformer::WideTopWrap<true>;
  230. Functions["WideToBlocks"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  231. Functions["WideFromBlocks"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  232. Functions["ReplicateScalars"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
  233. Functions["BlockMergeFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  234. Functions["BlockMergeManyFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap<true>;
  235. Functions["MultiHoppingCore"] = &TCallableConstraintTransformer::MultiHoppingCoreWrap;
  236. Functions["StablePickle"] = &TCallableConstraintTransformer::FromFirst<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode>;
  237. Functions["Unpickle"] = &TCallableConstraintTransformer::FromSecond<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode>;
  238. }
  239. std::optional<IGraphTransformer::TStatus> ProcessCore(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
  240. if (const auto func = Functions.find(input->Content()); Functions.cend() != func) {
  241. return (this->*func->second)(input, output, ctx);
  242. }
  243. return std::nullopt;
  244. }
  245. std::optional<IGraphTransformer::TStatus> ProcessList(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
  246. if (!input->ChildrenSize() || ETypeAnnotationKind::Tuple != input->GetTypeAnn()->GetKind())
  247. return TStatus::Ok;
  248. return AsTupleWrap(input, output, ctx);
  249. }
  250. TStatus ProcessUnknown(const TExprNode::TPtr& input, TExprContext&) {
  251. return UpdateAllChildLambdasConstraints(*input);
  252. }
  253. TStatus ValidateProviderCommitResult(const TExprNode::TPtr&, TExprContext&) {
  254. return TStatus::Ok;
  255. }
  256. TStatus ValidateProviderReadResult(const TExprNode::TPtr&, TExprContext&) {
  257. return TStatus::Ok;
  258. }
  259. TStatus ValidateProviderWriteResult(const TExprNode::TPtr&, TExprContext&) {
  260. return TStatus::Ok;
  261. }
  262. TStatus ValidateProviderConfigureResult(const TExprNode::TPtr&, TExprContext&) {
  263. return TStatus::Ok;
  264. }
  265. IGraphTransformer& GetTransformer(IDataProvider& provider) const {
  266. return provider.GetConstraintTransformer(InstantOnly, SubGraph);
  267. }
  268. private:
  269. template <size_t Ndx>
  270. TStatus CopyAllFrom(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  271. Y_UNUSED(output);
  272. Y_UNUSED(ctx);
  273. input->CopyConstraints(*input->Child(Ndx));
  274. return TStatus::Ok;
  275. }
  276. template <class... TConstraints>
  277. TStatus FromFirst(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  278. Y_UNUSED(output);
  279. Y_UNUSED(ctx);
  280. TApplyConstraintFromInput<0, TConstraints...>::Do(input);
  281. return TStatus::Ok;
  282. }
  283. template <class... TConstraints>
  284. TStatus FromSecond(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  285. Y_UNUSED(output);
  286. Y_UNUSED(ctx);
  287. TApplyConstraintFromInput<1, TConstraints...>::Do(input);
  288. return TStatus::Ok;
  289. }
  290. template <size_t StartFromChild, class... TConstraints>
  291. TStatus CommonFromChildren(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  292. Y_UNUSED(output);
  293. TVector<const TConstraintSet*> constraints;
  294. for (size_t i = StartFromChild; i < input->ChildrenSize(); ++i) {
  295. constraints.push_back(&input->Child(i)->GetConstraintSet());
  296. }
  297. TApplyCommonConstraint<TConstraints...>::Do(input, constraints, ctx);
  298. return TStatus::Ok;
  299. }
  300. TStatus FromEmpty(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  301. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  302. return TStatus::Ok;
  303. }
  304. TStatus FailMe(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  305. if (input->Child(0)->Content() == "constraint") {
  306. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  307. }
  308. return TStatus::Ok;
  309. }
  310. TStatus EmptyFromWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  311. auto set = input->Head().GetConstraintSet();
  312. set.RemoveConstraint(TEmptyConstraintNode::Name());
  313. if (!set) {
  314. const auto type = input->GetTypeAnn();
  315. output = ctx.NewCallable(input->Pos(), GetEmptyCollectionName(type), {ExpandType(input->Pos(), *type, ctx)});
  316. return TStatus::Repeat;
  317. }
  318. set.AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  319. input->SetConstraints(set);
  320. return TStatus::Ok;
  321. }
  322. template <size_t LambdaIdx>
  323. TStatus FromFinalLambda(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  324. TStatus status = UpdateAllChildLambdasConstraints(*input);
  325. if (status != TStatus::Ok) {
  326. return status;
  327. }
  328. TApplyConstraintFromInput<LambdaIdx, TMultiConstraintNode, TEmptyConstraintNode>::Do(input);
  329. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  330. }
  331. TStatus InheriteEmptyFromInput(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  332. auto status = UpdateAllChildLambdasConstraints(*input);
  333. if (status != TStatus::Ok) {
  334. return status;
  335. }
  336. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  337. }
  338. template<bool Sort>
  339. TStatus WideTopWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  340. if constexpr (Sort) {
  341. TSortedConstraintNode::TContainerType sorted;
  342. sorted.reserve(input->Tail().ChildrenSize());
  343. for (const auto& item : input->Tail().Children()) {
  344. if (item->Tail().IsCallable("Bool"))
  345. sorted.emplace_back(std::make_pair(TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType(1U, item->Head().Content())}, FromString<bool>(item->Tail().Tail().Content())));
  346. else
  347. break;
  348. }
  349. if (!sorted.empty()) {
  350. input->AddConstraint(ctx.MakeConstraint<TSortedConstraintNode>(std::move(sorted)));
  351. }
  352. }
  353. return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  354. }
  355. TStatus SortWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  356. if (const auto status = UpdateLambdaConstraints(input->Tail()); status != TStatus::Ok) {
  357. return status;
  358. }
  359. if (const auto sorted = DeduceSortConstraint(*input->Child(1), *input->Child(2), ctx)) {
  360. input->AddConstraint(sorted->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  361. }
  362. return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  363. }
  364. TStatus AssumeConstraintsWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  365. TConstraintSet set;
  366. try {
  367. set = ctx.MakeConstraintSet(NYT::NodeFromYsonString(input->Tail().Content()));
  368. } catch (...) {
  369. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() <<
  370. "Bad constraints yson-value: " << CurrentExceptionMessage()));
  371. return IGraphTransformer::TStatus::Error;
  372. }
  373. if (!set) {
  374. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), "AssumeConstraints with empty constraints set"));
  375. return IGraphTransformer::TStatus::Error;
  376. }
  377. for (auto constraint: set.GetAllConstraints()) {
  378. if (!constraint->IsApplicableToType(*input->GetTypeAnn())) {
  379. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint
  380. << " is not applicable to " << *input->GetTypeAnn()));
  381. return IGraphTransformer::TStatus::Error;
  382. }
  383. }
  384. for (auto constr: input->Head().GetAllConstraints()) {
  385. if (!constr->GetName().starts_with("PartOf") && !set.GetConstraint(constr->GetName())) {
  386. set.AddConstraint(constr);
  387. }
  388. }
  389. input->SetConstraints(set);
  390. return IGraphTransformer::TStatus::Ok;
  391. }
  392. template<bool Distinct, bool Strict>
  393. TStatus AssumeUniqueWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  394. typename TUniqueConstraintNodeBase<Distinct>::TContentType content;
  395. for (auto i = 1U; i < input->ChildrenSize(); ++i) {
  396. TPartOfConstraintBase::TSetOfSetsType sets;
  397. sets.reserve(input->Child(i)->ChildrenSize());
  398. for (const auto& list : input->Child(i)->Children()) {
  399. if (list->IsAtom())
  400. sets.insert_unique(TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType(1U, list->Content())});
  401. else if (list->IsList()) {
  402. TPartOfConstraintBase::TSetType columns;
  403. columns.reserve(list->ChildrenSize());
  404. for (const auto& column: list->Children()) {
  405. if (column->IsAtom())
  406. columns.insert_unique(TPartOfConstraintBase::TPathType(1U, column->Content()));
  407. else if (column->IsList()) {
  408. TPartOfConstraintBase::TPathType path(column->ChildrenSize());
  409. std::transform(column->Children().cbegin(), column->Children().cend(), path.begin(), [](const TExprNode::TPtr& atom) { return atom->Content(); } );
  410. columns.insert_unique(std::move(path));
  411. }
  412. }
  413. sets.insert_unique(std::move(columns));
  414. }
  415. }
  416. content.insert_unique(std::move(sets));
  417. }
  418. if (content.empty())
  419. content.insert_unique(TPartOfConstraintBase::TSetOfSetsType{TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType()}});
  420. auto constraint = ctx.MakeConstraint<TUniqueConstraintNodeBase<Distinct>>(std::move(content));
  421. if (!constraint->IsApplicableToType(*input->GetTypeAnn())) {
  422. if constexpr (Strict) {
  423. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint
  424. << " is not applicable to " << *input->GetTypeAnn()));
  425. } else {
  426. auto issue = TIssue(ctx.GetPosition(input->Pos()),
  427. TStringBuilder() << (Distinct ? "Distinct" : "Unique") << " sql hint contains invalid column: " << Endl
  428. << *constraint << " is not applicable to " << *input->GetTypeAnn());
  429. SetIssueCode(EYqlIssueCode::TIssuesIds_EIssueCode_YQL_HINT_INVALID_PARAMETERS, issue);
  430. if (ctx.AddWarning(issue)) {
  431. output = input->HeadPtr();
  432. return IGraphTransformer::TStatus::Repeat;
  433. }
  434. }
  435. return IGraphTransformer::TStatus::Error;
  436. }
  437. if constexpr (!Strict) {
  438. output = ctx.RenameNode(*input, Distinct ? "AssumeDistinct" : "AssumeUnique");
  439. return IGraphTransformer::TStatus::Repeat;
  440. }
  441. if (const auto old = input->Head().GetConstraint<TUniqueConstraintNodeBase<Distinct>>()) {
  442. if (old->Includes(*constraint)) {
  443. output = input->HeadPtr();
  444. return TStatus::Repeat;
  445. } else
  446. constraint = TUniqueConstraintNodeBase<Distinct>::Merge(old, constraint, ctx);
  447. }
  448. input->AddConstraint(constraint);
  449. return FromFirst<TSortedConstraintNode, TChoppedConstraintNode, TUniqueConstraintNodeBase<!Distinct>, TEmptyConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  450. }
  451. TStatus AssumeChoppedWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  452. TPartOfConstraintBase::TSetOfSetsType sets;
  453. for (auto i = 1U; i < input->ChildrenSize(); ++i) {
  454. TPartOfConstraintBase::TSetType columns;
  455. columns.reserve(input->Child(i)->ChildrenSize());
  456. for (const auto& column: input->Child(i)->Children()) {
  457. if (column->IsAtom())
  458. columns.insert_unique(TPartOfConstraintBase::TPathType(1U, column->Content()));
  459. else if (column->IsList()) {
  460. TPartOfConstraintBase::TPathType path(column->ChildrenSize());
  461. std::transform(column->Children().cbegin(), column->Children().cend(), path.begin(), [](const TExprNode::TPtr& atom) { return atom->Content(); } );
  462. columns.insert_unique(std::move(path));
  463. }
  464. }
  465. sets.insert_unique(std::move(columns));
  466. }
  467. const auto constraint = ctx.MakeConstraint<TChoppedConstraintNode>(std::move(sets));
  468. if (!constraint->IsApplicableToType(*input->GetTypeAnn())) {
  469. ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint
  470. << " is not applicable to " << *input->GetTypeAnn()));
  471. return IGraphTransformer::TStatus::Error;
  472. }
  473. if (const auto old = input->Head().GetConstraint<TChoppedConstraintNode>()) {
  474. if (old->Equals(*constraint)) {
  475. output = input->HeadPtr();
  476. return TStatus::Repeat;
  477. }
  478. }
  479. input->AddConstraint(constraint);
  480. return FromFirst<TSortedConstraintNode, TDistinctConstraintNode, TUniqueConstraintNode, TEmptyConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  481. }
  482. template <bool UseSort>
  483. TStatus TopWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  484. if (const auto status = UpdateLambdaConstraints(input->Tail()); status != TStatus::Ok) {
  485. return status;
  486. }
  487. if constexpr (UseSort) {
  488. if (const auto sorted = DeduceSortConstraint(*input->Child(2), *input->Child(3), ctx)) {
  489. input->AddConstraint(sorted->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  490. }
  491. }
  492. return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  493. }
  494. template<class TConstraint>
  495. static void FilterFromHead(const TExprNode& input, TConstraintSet& constraints, const TPartOfConstraintBase::TPathFilter& filter, TExprContext& ctx) {
  496. if (const auto source = input.Head().GetConstraint<TConstraint>()) {
  497. if (const auto filtered = source->FilterFields(ctx, filter)) {
  498. constraints.AddConstraint(filtered);
  499. }
  500. }
  501. }
  502. template<class TConstraint, bool Simplify = false>
  503. static void ReduceFromHead(const TExprNode::TPtr& input, const TPartOfConstraintBase::TPathReduce& reduce, TExprContext& ctx) {
  504. if (const auto source = input->Head().GetConstraint<TConstraint>()) {
  505. if (const auto filtered = source->RenameFields(ctx, reduce)) {
  506. if constexpr (Simplify)
  507. input->AddConstraint(filtered->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  508. else
  509. input->AddConstraint(filtered);
  510. }
  511. }
  512. }
  513. template<class TConstraint, bool Simplify = false>
  514. static void FilterFromHead(const TExprNode::TPtr& input, const TPartOfConstraintBase::TPathFilter& filter, TExprContext& ctx) {
  515. if (const auto source = input->Head().GetConstraint<TConstraint>()) {
  516. if (const auto filtered = source->FilterFields(ctx, filter)) {
  517. if constexpr (Simplify)
  518. input->AddConstraint(filtered->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  519. else
  520. input->AddConstraint(filtered);
  521. }
  522. }
  523. }
  524. template<class TConstraint>
  525. static void FilterFromHeadIfMissed(const TExprNode::TPtr& input, const TPartOfConstraintBase::TPathFilter& filter, TExprContext& ctx) {
  526. if (!input->GetConstraint<TConstraint>())
  527. FilterFromHead<TConstraint>(input, filter, ctx);
  528. }
  529. TStatus SelectMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  530. auto outItemType = input->GetTypeAnn();
  531. while (outItemType->GetKind() == ETypeAnnotationKind::Optional) {
  532. outItemType = outItemType->Cast<TOptionalExprType>()->GetItemType();
  533. }
  534. if (outItemType->GetKind() == ETypeAnnotationKind::Variant) {
  535. if (outItemType->Cast<TVariantExprType>()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) {
  536. const auto outSize = outItemType->Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>()->GetSize();
  537. auto multi = input->Head().GetConstraint<TMultiConstraintNode>();
  538. if (multi && multi->GetItems().back().first >= outSize) {
  539. TMultiConstraintNode::TMapType filteredItems;
  540. for (auto& item: multi->GetItems()) {
  541. if (item.first < outSize) {
  542. filteredItems.push_back(item);
  543. }
  544. }
  545. multi = filteredItems.empty() ? nullptr : ctx.MakeConstraint<TMultiConstraintNode>(std::move(filteredItems));
  546. }
  547. if (multi) {
  548. input->AddConstraint(multi);
  549. }
  550. auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>();
  551. if (varIndex && varIndex->GetIndexMapping().back().first >= outSize) {
  552. TVarIndexConstraintNode::TMapType filteredItems;
  553. for (auto& item: varIndex->GetIndexMapping()) {
  554. if (item.first < outSize) {
  555. filteredItems.push_back(item);
  556. }
  557. }
  558. varIndex = filteredItems.empty() ? nullptr : ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(filteredItems));
  559. }
  560. if (varIndex) {
  561. input->AddConstraint(varIndex);
  562. }
  563. }
  564. }
  565. else if (outItemType->GetKind() == ETypeAnnotationKind::Struct) {
  566. const auto filter = [outItemType](const TPartOfConstraintBase::TPathType& path) {
  567. return !path.empty() && TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType);
  568. };
  569. FilterFromHead<TPartOfSortedConstraintNode>(input, filter, ctx);
  570. FilterFromHead<TPartOfChoppedConstraintNode>(input, filter, ctx);
  571. FilterFromHead<TPartOfUniqueConstraintNode>(input, filter, ctx);
  572. FilterFromHead<TPartOfDistinctConstraintNode>(input, filter, ctx);
  573. }
  574. return TStatus::Ok;
  575. }
  576. template <bool Strict>
  577. TStatus CastWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  578. const auto outItemType = input->GetTypeAnn();
  579. const auto inItemType = input->Head().GetTypeAnn();
  580. const auto filter = [inItemType, outItemType, toString = input->IsCallable({"ToString", "ToBytes"})](const TPartOfConstraintBase::TPathType& path) {
  581. if (const auto outType = TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType)) {
  582. const auto inType = TPartOfConstraintBase::GetSubTypeByPath(path, *inItemType);
  583. return (toString && inType->GetKind() == ETypeAnnotationKind::Data && inType->Cast<TDataExprType>()->GetSlot() == EDataSlot::Utf8) || IsSameAnnotation(*outType, *inType);
  584. }
  585. return false;
  586. };
  587. const auto filterForUnique = [inItemType, outItemType](const TPartOfConstraintBase::TPathType& path) {
  588. const auto castResult = CastResult<Strict>(TPartOfConstraintBase::GetSubTypeByPath(path, *inItemType), TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType));
  589. return NUdf::ECastOptions::Complete == castResult || NUdf::ECastOptions::MayFail == castResult;
  590. };
  591. const auto filterForDistinct = [inItemType, outItemType](const TPartOfConstraintBase::TPathType& path) {
  592. return NUdf::ECastOptions::Complete == CastResult<Strict>(TPartOfConstraintBase::GetSubTypeByPath(path, *inItemType), TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType));
  593. };
  594. FilterFromHead<TSortedConstraintNode>(input, filter, ctx);
  595. FilterFromHead<TChoppedConstraintNode>(input, filter, ctx);
  596. FilterFromHead<TUniqueConstraintNode>(input, filterForUnique, ctx);
  597. FilterFromHead<TDistinctConstraintNode>(input, filterForDistinct, ctx);
  598. FilterFromHead<TPartOfSortedConstraintNode>(input, filter, ctx);
  599. FilterFromHead<TPartOfChoppedConstraintNode>(input, filter, ctx);
  600. FilterFromHead<TPartOfUniqueConstraintNode>(input, filterForUnique, ctx);
  601. FilterFromHead<TPartOfDistinctConstraintNode>(input, filterForDistinct, ctx);
  602. const auto unwrapedOutItemType = RemoveOptionalType(outItemType);
  603. const auto unwrapedInItemType = RemoveOptionalType(inItemType);
  604. if (unwrapedInItemType->GetKind() == ETypeAnnotationKind::Variant && unwrapedOutItemType->GetKind() == ETypeAnnotationKind::Variant
  605. && unwrapedOutItemType->Cast<TVariantExprType>()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) {
  606. const auto tupleUnderInType = unwrapedInItemType->Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>();
  607. const auto tupleUnderOutType = unwrapedOutItemType->Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>();
  608. if (auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) {
  609. if (tupleUnderOutType->GetSize() < tupleUnderInType->GetSize()) {
  610. TMultiConstraintNode::TMapType multiItems;
  611. std::copy_if(multi->GetItems().cbegin(), multi->GetItems().cend(),
  612. std::back_inserter(multiItems),
  613. [&](const auto& item) { return item.first < tupleUnderOutType->GetSize(); }
  614. );
  615. if (!multiItems.empty()) {
  616. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(multiItems)));
  617. }
  618. } else {
  619. input->AddConstraint(multi);
  620. }
  621. }
  622. if (auto varItem = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  623. if (tupleUnderOutType->GetSize() < tupleUnderInType->GetSize()) {
  624. TVarIndexConstraintNode::TMapType filteredItems;
  625. std::copy_if(varItem->GetIndexMapping().cbegin(), varItem->GetIndexMapping().cend(),
  626. std::back_inserter(filteredItems),
  627. [&](const auto& item) { return item.second < tupleUnderOutType->GetSize(); }
  628. );
  629. if (!filteredItems.empty()) {
  630. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(filteredItems)));
  631. }
  632. } else {
  633. input->AddConstraint(varItem);
  634. }
  635. }
  636. }
  637. return TStatus::Ok;
  638. }
  639. TStatus DivePrefixMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  640. const auto prefixes = input->Tail().Children();
  641. const auto rename = [&prefixes](const TPartOfConstraintBase::TPathType& path) -> std::vector<TPartOfConstraintBase::TPathType> {
  642. if (path.empty())
  643. return {};
  644. for (const auto& p : prefixes) {
  645. if (const auto& prefix = p->Content(); path.front().starts_with(prefix)) {
  646. auto out = path;
  647. out.front() = out.front().substr(prefix.length());
  648. return {std::move(out)};
  649. }
  650. }
  651. return {};
  652. };
  653. ReduceFromHead<TPartOfSortedConstraintNode>(input, rename, ctx);
  654. ReduceFromHead<TPartOfChoppedConstraintNode>(input, rename, ctx);
  655. ReduceFromHead<TPartOfUniqueConstraintNode>(input, rename, ctx);
  656. ReduceFromHead<TPartOfDistinctConstraintNode>(input, rename, ctx);
  657. return FromFirst<TVarIndexConstraintNode>(input, output, ctx);
  658. }
  659. TStatus ExtractMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  660. const auto outItemType = GetSeqItemType(*input->GetTypeAnn()).Cast<TStructExprType>();
  661. const auto filter = [outItemType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && outItemType->FindItem(path.front()); };
  662. FilterFromHead<TSortedConstraintNode, true>(input, filter, ctx);
  663. FilterFromHead<TChoppedConstraintNode, true>(input, filter, ctx);
  664. FilterFromHead<TUniqueConstraintNode, true>(input, filter, ctx);
  665. FilterFromHead<TDistinctConstraintNode, true>(input, filter, ctx);
  666. FilterFromHead<TPartOfSortedConstraintNode>(input, filter, ctx);
  667. FilterFromHead<TPartOfChoppedConstraintNode>(input, filter, ctx);
  668. FilterFromHead<TPartOfUniqueConstraintNode>(input, filter, ctx);
  669. FilterFromHead<TPartOfDistinctConstraintNode>(input, filter, ctx);
  670. return FromFirst<TEmptyConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  671. }
  672. TStatus RemovePrefixMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  673. const TTypeAnnotationNode* outItemType = GetSeqItemType(input->GetTypeAnn());
  674. if (!outItemType) {
  675. outItemType = input->GetTypeAnn();
  676. }
  677. if (outItemType->GetKind() == ETypeAnnotationKind::Struct) {
  678. const auto outStructType = outItemType->Cast<TStructExprType>();
  679. const auto filter = [outStructType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && outStructType->FindItem(path.front()); };
  680. FilterFromHead<TPartOfSortedConstraintNode>(input, filter, ctx);
  681. FilterFromHead<TPartOfChoppedConstraintNode>(input, filter, ctx);
  682. FilterFromHead<TPartOfUniqueConstraintNode>(input, filter, ctx);
  683. FilterFromHead<TPartOfDistinctConstraintNode>(input, filter, ctx);
  684. }
  685. else if (outItemType->GetKind() == ETypeAnnotationKind::Variant) {
  686. if (auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) {
  687. TMultiConstraintNode::TMapType multiItems;
  688. auto tupleUnderType = outItemType->Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>();
  689. for (auto& item: multi->GetItems()) {
  690. YQL_ENSURE(item.first < tupleUnderType->GetSize());
  691. auto& constr = multiItems[item.first];
  692. const auto outStructType = tupleUnderType->GetItems()[item.first]->Cast<TStructExprType>();
  693. const auto filter = [outStructType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && outStructType->FindItem(path.front()); };
  694. FilterFromHead<TPartOfSortedConstraintNode>(*input, constr, filter, ctx);
  695. FilterFromHead<TPartOfChoppedConstraintNode>(*input, constr, filter, ctx);
  696. FilterFromHead<TPartOfUniqueConstraintNode>(*input, constr, filter, ctx);
  697. FilterFromHead<TPartOfDistinctConstraintNode>(*input, constr, filter, ctx);
  698. }
  699. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(multiItems)));
  700. }
  701. }
  702. return FromFirst<TEmptyConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
  703. }
  704. // TODO: Empty for false condition
  705. template <bool Ordered>
  706. TStatus FilterWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  707. if (const auto status = UpdateLambdaConstraints(*input->Child(1)); status != TStatus::Ok) {
  708. return status;
  709. }
  710. if constexpr (Ordered) {
  711. FromFirst<TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode>(input, output, ctx);
  712. }
  713. return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  714. }
  715. template<class TConstraint>
  716. static const TConstraint* GetConstraintFromWideResultLambda(const TExprNode& lambda, TExprContext& ctx);
  717. template<class TConstraintType>
  718. static const TConstraintType* GetLambdaConstraint(const TExprNode& lambda, TExprContext& ctx) {
  719. if (2U == lambda.ChildrenSize())
  720. return lambda.Tail().GetConstraint<TConstraintType>();
  721. TVector<const TConstraintSet*> constraints;
  722. constraints.reserve(lambda.ChildrenSize() - 1U);
  723. for (size_t i = 1U; i < lambda.ChildrenSize(); ++i) {
  724. constraints.emplace_back(&lambda.Child(i)->GetConstraintSet());
  725. }
  726. return TConstraintType::MakeCommon(constraints, ctx);
  727. }
  728. template<class TConstraintType, bool WideLambda>
  729. static const TConstraintType* GetConstraintFromLambda(const TExprNode& lambda, TExprContext& ctx) {
  730. if constexpr (WideLambda)
  731. return GetConstraintFromWideResultLambda<TConstraintType>(lambda, ctx);
  732. else
  733. return GetLambdaConstraint<TConstraintType>(lambda, ctx);
  734. }
  735. static std::optional<bool> GetDirection(const TExprNode& dir) {
  736. if (dir.IsCallable("Bool"))
  737. return IsTrue(dir.Tail().Content());
  738. if (dir.IsCallable("Not"))
  739. if (const auto d = GetDirection(dir.Head()))
  740. return !*d;
  741. return std::nullopt;
  742. }
  743. static std::vector<std::pair<TPartOfConstraintBase::TPathType, bool>>
  744. ExtractSimpleSortTraits(const TExprNode& sortDirections, const TExprNode& keySelectorLambda) {
  745. const auto& keySelectorBody = keySelectorLambda.Tail();
  746. const auto& keySelectorArg = keySelectorLambda.Head().Head();
  747. std::vector<std::pair<TPartOfConstraintBase::TPathType, bool>> columns;
  748. if (const auto dir = GetDirection(sortDirections))
  749. columns.emplace_back(TPartOfConstraintBase::TPathType(), *dir);
  750. else if (sortDirections.IsList())
  751. if (const auto size = keySelectorBody.ChildrenSize()) {
  752. columns.reserve(size);
  753. for (auto i = 0U; i < size; ++i)
  754. if (const auto dir = GetDirection(*sortDirections.Child(i)))
  755. columns.emplace_back(TPartOfConstraintBase::TPathType(), *dir);
  756. else
  757. return {};
  758. } else
  759. return {};
  760. else
  761. return {};
  762. if (keySelectorBody.IsList())
  763. if (const auto size = keySelectorBody.ChildrenSize()) {
  764. TPartOfConstraintBase::TSetType set;
  765. set.reserve(size);
  766. columns.resize(size, std::make_pair(TPartOfConstraintBase::TPathType(), columns.back().second));
  767. auto it = columns.begin();
  768. for (auto i = 0U; i < size; ++i) {
  769. if (auto path = GetPathToKey<true>(*keySelectorBody.Child(i), keySelectorArg)) {
  770. if (set.insert(*path).second)
  771. it++->first = std::move(*path);
  772. else if (columns.cend() != it)
  773. it = columns.erase(it);
  774. } else {
  775. return {};
  776. }
  777. }
  778. } else
  779. return {};
  780. else if (auto path = GetPathToKey<true>(keySelectorBody, keySelectorArg))
  781. if (columns.size() == 1U)
  782. columns.front().first = std::move(*path);
  783. else
  784. return {};
  785. else
  786. return {};
  787. return columns;
  788. }
  789. template<class TConstraint, bool OrderedMap, class TInput>
  790. static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, bool isSingleItem, TExprContext& ctx) {
  791. constexpr bool isOrderConstraint = std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>();
  792. if (const auto lambda = handler.GetConstraint<TConstraint>()) {
  793. const auto original = input.template GetConstraint<typename TConstraint::TMainConstraint>();
  794. if constexpr (OrderedMap || !isOrderConstraint) {
  795. if (original) {
  796. if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) {
  797. output.AddConstraint(complete);
  798. }
  799. }
  800. }
  801. if (const auto part = input.template GetConstraint<TConstraint>()) {
  802. auto mapping = lambda->GetColumnMapping();
  803. for (auto it = mapping.cbegin(); mapping.cend() != it;) {
  804. if (part->GetColumnMapping().contains(it->first))
  805. ++it;
  806. else
  807. it = mapping.erase(it);
  808. }
  809. if (!mapping.empty()) {
  810. output.AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
  811. }
  812. } else if (isOrderConstraint || isSingleItem) {
  813. if (const auto filtered = lambda->RemoveOriginal(ctx, original))
  814. output.AddConstraint(filtered);
  815. }
  816. }
  817. }
  818. template<class TConstraint, bool OrderedMap, bool WideOutput>
  819. static void GetFromMapLambda(const TExprNode::TPtr& input, bool isSingleItem, TExprContext& ctx) {
  820. constexpr bool isOrderConstraint = std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>();
  821. if (const auto lambda = GetConstraintFromLambda<TConstraint, WideOutput>(input->Tail(), ctx)) {
  822. const auto original = GetDetailed(input->Head().GetConstraint<typename TConstraint::TMainConstraint>(), *input->Head().GetTypeAnn(), ctx);
  823. if constexpr (OrderedMap || !isOrderConstraint) {
  824. if (original) {
  825. if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) {
  826. input->AddConstraint(complete->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  827. }
  828. }
  829. }
  830. if (const auto part = input->Head().GetConstraint<TConstraint>()) {
  831. auto mapping = lambda->GetColumnMapping();
  832. for (auto it = mapping.cbegin(); mapping.cend() != it;) {
  833. if (part->GetColumnMapping().contains(it->first))
  834. ++it;
  835. else
  836. it = mapping.erase(it);
  837. }
  838. if (!mapping.empty()) {
  839. input->AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
  840. }
  841. } else if (isOrderConstraint || isSingleItem) {
  842. if (const auto filtered = lambda->RemoveOriginal(ctx, original))
  843. input->AddConstraint(filtered);
  844. }
  845. }
  846. }
  847. template <bool Ordered, bool Flat, bool WideInput = false, bool WideOutput = false>
  848. TStatus MapWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  849. auto argConstraints = GetConstraintsForInputArgument<Ordered, WideInput>(*input, ctx);
  850. if constexpr (Ordered && !(Flat || WideInput || WideOutput)) {
  851. // TODO: is temporary crutch for MapNext.
  852. if (argConstraints.size() < input->Tail().Head().ChildrenSize())
  853. argConstraints.resize(input->Tail().Head().ChildrenSize(), argConstraints.front());
  854. }
  855. if (const auto status = UpdateLambdaConstraints(input->TailRef(), ctx, argConstraints); status != TStatus::Ok) {
  856. return status;
  857. }
  858. const bool singleItem = ETypeAnnotationKind::Optional == input->GetTypeAnn()->GetKind();
  859. GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
  860. GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
  861. GetFromMapLambda<TPartOfSortedConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
  862. GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
  863. const auto lambdaVarIndex = GetConstraintFromLambda<TVarIndexConstraintNode, WideOutput>(input->Tail(), ctx);
  864. const auto lambdaMulti = GetConstraintFromLambda<TMultiConstraintNode, WideOutput>(input->Tail(), ctx);
  865. const auto inItemType = GetSeqItemType(input->Head().GetTypeAnn());
  866. const bool multiInput = ETypeAnnotationKind::Variant == inItemType->GetKind();
  867. if (const auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  868. if (multiInput) {
  869. if (lambdaVarIndex) {
  870. if (const auto outVarIndex = GetVarIndexOverVarIndexConstraint(*varIndex, *lambdaVarIndex, ctx)) {
  871. input->AddConstraint(outVarIndex);
  872. }
  873. }
  874. } else {
  875. if (lambdaMulti) {
  876. TVarIndexConstraintNode::TMapType remapItems;
  877. for (auto& multiItem: lambdaMulti->GetItems()) {
  878. for (auto& varItem: varIndex->GetIndexMapping()) {
  879. remapItems.push_back(std::make_pair(multiItem.first, varItem.second));
  880. }
  881. }
  882. if (!remapItems.empty()) {
  883. ::SortUnique(remapItems);
  884. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(remapItems)));
  885. }
  886. } else {
  887. input->AddConstraint(varIndex);
  888. }
  889. }
  890. }
  891. const auto inputMulti = input->Head().GetConstraint<TMultiConstraintNode>();
  892. if (lambdaMulti && !input->Head().GetConstraint<TEmptyConstraintNode>()) {
  893. TMultiConstraintNode::TMapType remappedItems;
  894. for (auto& item: lambdaMulti->GetItems()) {
  895. remappedItems.push_back(std::make_pair(item.first, TConstraintSet{}));
  896. if (!multiInput) { // remapping one to many
  897. GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
  898. GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
  899. GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
  900. GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
  901. if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>()) {
  902. remappedItems.pop_back();
  903. }
  904. }
  905. else if (lambdaVarIndex && inputMulti) {
  906. const auto range = lambdaVarIndex->GetIndexMapping().equal_range(item.first);
  907. switch (std::distance(range.first, range.second)) {
  908. case 0: // new index
  909. break;
  910. case 1: // remapping 1 to 1
  911. if (const auto origConstr = inputMulti->GetItem(range.first->second)) {
  912. GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
  913. GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
  914. GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
  915. GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
  916. if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>()) {
  917. remappedItems.pop_back();
  918. }
  919. } else {
  920. remappedItems.pop_back();
  921. }
  922. break;
  923. default: // remapping many to one
  924. {
  925. std::vector<const TConstraintSet*> nonEmpty;
  926. for (auto i = range.first; i != range.second; ++i) {
  927. if (auto origConstr = inputMulti->GetItem(i->second)) {
  928. nonEmpty.push_back(origConstr);
  929. }
  930. }
  931. EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint<TEmptyConstraintNode>(); });
  932. if (nonEmpty.empty()) {
  933. remappedItems.back().second.AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  934. } else if (nonEmpty.size() == 1) {
  935. remappedItems.back().second = std::move(*nonEmpty.front());
  936. }
  937. }
  938. }
  939. } else {
  940. remappedItems.back().second = item.second;
  941. }
  942. }
  943. if (remappedItems) {
  944. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(remappedItems)));
  945. }
  946. }
  947. else if (inputMulti && lambdaVarIndex) { // Many to one
  948. const auto range = lambdaVarIndex->GetIndexMapping().equal_range(0);
  949. std::vector<const TConstraintSet*> nonEmpty;
  950. for (auto i = range.first; i != range.second; ++i) {
  951. if (auto origConstr = inputMulti->GetItem(i->second)) {
  952. nonEmpty.push_back(origConstr);
  953. }
  954. }
  955. EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint<TEmptyConstraintNode>(); });
  956. if (nonEmpty.empty()) {
  957. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  958. } else if (nonEmpty.size() == 1) {
  959. input->SetConstraints(*nonEmpty.front());
  960. }
  961. }
  962. if constexpr (Flat) {
  963. if (const auto lambdaEmpty = GetConstraintFromLambda<TEmptyConstraintNode, WideOutput>(input->Tail(), ctx)) {
  964. input->AddConstraint(lambdaEmpty);
  965. const auto& filter = std::bind(&TPartOfConstraintBase::GetSubTypeByPath, std::placeholders::_1, std::cref(GetSeqItemType(*input->GetTypeAnn())));
  966. FilterFromHeadIfMissed<TUniqueConstraintNode>(input, filter, ctx);
  967. FilterFromHeadIfMissed<TDistinctConstraintNode>(input, filter, ctx);
  968. if constexpr (Ordered) {
  969. FilterFromHeadIfMissed<TSortedConstraintNode>(input, filter, ctx);
  970. FilterFromHeadIfMissed<TChoppedConstraintNode>(input, filter, ctx);
  971. }
  972. }
  973. }
  974. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  975. }
  976. template <bool Ordered>
  977. TStatus LMapWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  978. TConstraintNode::TListType argConstraints;
  979. for (const auto c: input->Head().GetAllConstraints()) {
  980. if (Ordered || (c->GetName() != TSortedConstraintNode::Name() && c->GetName() != TChoppedConstraintNode::Name())) {
  981. argConstraints.push_back(c);
  982. }
  983. }
  984. if (const auto status = UpdateLambdaConstraints(input->TailRef(), ctx, {argConstraints}); status != TStatus::Ok) {
  985. return status;
  986. }
  987. TSet<TStringBuf> except;
  988. if constexpr (!Ordered) {
  989. except.insert(TSortedConstraintNode::Name());
  990. except.insert(TChoppedConstraintNode::Name());
  991. }
  992. if (input->Tail().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional) {
  993. except.insert(TEmptyConstraintNode::Name());
  994. }
  995. CopyExcept(*input, input->Tail(), except);
  996. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  997. }
  998. TStatus AsListWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  999. if (input->ChildrenSize() == 1) {
  1000. if (const auto unique = input->Head().GetConstraint<TUniqueConstraintNode>()) {
  1001. input->AddConstraint(unique);
  1002. }
  1003. if (const auto unique = input->Head().GetConstraint<TDistinctConstraintNode>()) {
  1004. input->AddConstraint(unique);
  1005. }
  1006. if (const auto part = input->Head().GetConstraint<TPartOfUniqueConstraintNode>()) {
  1007. input->AddConstraint(part);
  1008. }
  1009. if (const auto part = input->Head().GetConstraint<TPartOfDistinctConstraintNode>()) {
  1010. input->AddConstraint(part);
  1011. }
  1012. }
  1013. return CommonFromChildren<0, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  1014. }
  1015. template<bool Ordered>
  1016. TStatus ExtendWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1017. if (input->ChildrenSize() == 1) {
  1018. if (const auto unique = input->Head().GetConstraint<TUniqueConstraintNode>()) {
  1019. input->AddConstraint(unique);
  1020. }
  1021. if (const auto part = input->Head().GetConstraint<TPartOfUniqueConstraintNode>()) {
  1022. input->AddConstraint(part);
  1023. }
  1024. if (const auto unique = input->Head().GetConstraint<TDistinctConstraintNode>()) {
  1025. input->AddConstraint(unique);
  1026. }
  1027. if (const auto part = input->Head().GetConstraint<TPartOfDistinctConstraintNode>()) {
  1028. input->AddConstraint(part);
  1029. }
  1030. if constexpr (Ordered) {
  1031. if (const auto sorted = input->Head().GetConstraint<TSortedConstraintNode>()) {
  1032. input->AddConstraint(sorted);
  1033. }
  1034. if (const auto part = input->Head().GetConstraint<TPartOfSortedConstraintNode>()) {
  1035. input->AddConstraint(part);
  1036. }
  1037. if (const auto sorted = input->Head().GetConstraint<TChoppedConstraintNode>()) {
  1038. input->AddConstraint(sorted);
  1039. }
  1040. if (const auto part = input->Head().GetConstraint<TPartOfChoppedConstraintNode>()) {
  1041. input->AddConstraint(part);
  1042. }
  1043. }
  1044. }
  1045. return CommonFromChildren<0, TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  1046. }
  1047. template <bool Union>
  1048. TStatus MergeWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1049. if (auto sort = MakeCommonConstraint<TSortedConstraintNode>(input, 0, ctx)) {
  1050. if (Union && input->ChildrenSize() > 1) {
  1051. // Check and exclude modified keys from final constraint
  1052. const auto resultItemType = input->GetTypeAnn()->Cast<TListExprType>()->GetItemType();
  1053. std::vector<const TTypeAnnotationNode*> inputs;
  1054. for (const auto& child: input->Children()) {
  1055. inputs.emplace_back(child->GetTypeAnn()->Cast<TListExprType>()->GetItemType());
  1056. }
  1057. auto content = sort->GetContent();
  1058. for (auto i = 0U; i < content.size(); ++i) {
  1059. for (auto it = content[i].first.cbegin(); content[i].first.cend() != it;) {
  1060. const auto resultItemSubType = TPartOfConstraintBase::GetSubTypeByPath(*it, *resultItemType);
  1061. YQL_ENSURE(resultItemSubType, "Missing " << *it << " in result type");
  1062. auto childNdx = 0U;
  1063. while (childNdx < input->ChildrenSize()) {
  1064. if (const auto inputItemSubType = TPartOfConstraintBase::GetSubTypeByPath(*it, *inputs[childNdx])) {
  1065. if (IsSameAnnotation(*inputItemSubType, *resultItemSubType)) {
  1066. ++childNdx;
  1067. continue;
  1068. }
  1069. } else {
  1070. YQL_ENSURE(input->Child(childNdx)->GetConstraint<TEmptyConstraintNode>(), "Missing column " << *it << " in non empty input type");
  1071. }
  1072. break;
  1073. }
  1074. if (childNdx < input->ChildrenSize())
  1075. it = content[i].first.erase(it);
  1076. else
  1077. ++it;
  1078. }
  1079. if (content[i].first.empty()) {
  1080. content.resize(i);
  1081. break;
  1082. }
  1083. }
  1084. sort = content.empty() ? nullptr : ctx.MakeConstraint<TSortedConstraintNode>(std::move(content));
  1085. }
  1086. if (sort) {
  1087. input->AddConstraint(sort);
  1088. }
  1089. }
  1090. return ExtendWrap<false>(input, output, ctx);
  1091. }
  1092. TStatus TakeWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1093. if (input->Tail().IsCallable("Uint64") && !FromString<ui64>(input->Tail().Head().Content())) {
  1094. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1095. }
  1096. return CopyAllFrom<0>(input, output, ctx);
  1097. }
  1098. TStatus MemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1099. const auto& memberName = input->Tail().Content();
  1100. const auto& structNode = input->Head();
  1101. if (const auto emptyConstraint = structNode.GetConstraint<TEmptyConstraintNode>()) {
  1102. input->AddConstraint(emptyConstraint);
  1103. } else {
  1104. if (const auto part = structNode.GetConstraint<TPartOfSortedConstraintNode>()) {
  1105. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1106. input->AddConstraint(extracted);
  1107. }
  1108. }
  1109. if (const auto part = structNode.GetConstraint<TPartOfChoppedConstraintNode>()) {
  1110. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1111. input->AddConstraint(extracted);
  1112. }
  1113. }
  1114. if (const auto part = structNode.GetConstraint<TPartOfUniqueConstraintNode>()) {
  1115. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1116. input->AddConstraint(extracted);
  1117. }
  1118. }
  1119. if (const auto part = structNode.GetConstraint<TPartOfDistinctConstraintNode>()) {
  1120. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1121. input->AddConstraint(extracted);
  1122. }
  1123. }
  1124. }
  1125. if (structNode.IsCallable("AsStruct")) {
  1126. for (const auto& child: structNode.Children()) {
  1127. if (child->Head().IsAtom(memberName)) {
  1128. TApplyConstraintFromInput<1, TVarIndexConstraintNode>::Do(child);
  1129. break;
  1130. }
  1131. }
  1132. } else {
  1133. TApplyConstraintFromInput<0, TVarIndexConstraintNode>::Do(input);
  1134. }
  1135. return TStatus::Ok;
  1136. }
  1137. TStatus AsTupleWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1138. TPartOfSortedConstraintNode::TMapType sorted;
  1139. TPartOfChoppedConstraintNode::TMapType chopped;
  1140. TPartOfUniqueConstraintNode::TMapType uniques;
  1141. TPartOfDistinctConstraintNode::TMapType distincts;
  1142. std::vector<const TConstraintSet*> structConstraints;
  1143. for (auto i = 0U; i < input->ChildrenSize(); ++i) {
  1144. const auto child = input->Child(i);
  1145. const auto& name = ctx.GetIndexAsString(i);
  1146. if (const auto part = child->GetConstraint<TPartOfSortedConstraintNode>()) {
  1147. TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(name));
  1148. }
  1149. if (const auto part = child->GetConstraint<TPartOfChoppedConstraintNode>()) {
  1150. TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(name));
  1151. }
  1152. if (const auto part = child->GetConstraint<TPartOfUniqueConstraintNode>()) {
  1153. TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(name));
  1154. }
  1155. if (const auto part = child->GetConstraint<TPartOfDistinctConstraintNode>()) {
  1156. TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(name));
  1157. }
  1158. if (const auto& valueNode = SkipModifiers(child); TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) {
  1159. structConstraints.push_back(&valueNode->Head().GetConstraintSet());
  1160. } else if (valueNode->IsArgument() && ETypeAnnotationKind::Struct != valueNode->GetTypeAnn()->GetKind()) {
  1161. structConstraints.push_back(&valueNode->GetConstraintSet());
  1162. }
  1163. }
  1164. if (!sorted.empty()) {
  1165. input->AddConstraint(ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(sorted)));
  1166. }
  1167. if (!chopped.empty()) {
  1168. input->AddConstraint(ctx.MakeConstraint<TPartOfChoppedConstraintNode>(std::move(chopped)));
  1169. }
  1170. if (!uniques.empty()) {
  1171. input->AddConstraint(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(uniques)));
  1172. }
  1173. if (!distincts.empty()) {
  1174. input->AddConstraint(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(distincts)));
  1175. }
  1176. if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) {
  1177. input->AddConstraint(varIndex);
  1178. }
  1179. return TStatus::Ok;
  1180. }
  1181. TStatus AsStructWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1182. TPartOfSortedConstraintNode::TMapType sorted;
  1183. TPartOfChoppedConstraintNode::TMapType chopped;
  1184. TPartOfUniqueConstraintNode::TMapType uniques;
  1185. TPartOfDistinctConstraintNode::TMapType distincts;
  1186. std::vector<const TConstraintSet*> structConstraints;
  1187. for (const auto& child : input->Children()) {
  1188. const auto& name = child->Head().Content();
  1189. if (const auto part = child->Tail().GetConstraint<TPartOfSortedConstraintNode>()) {
  1190. TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(name));
  1191. }
  1192. if (const auto part = child->Tail().GetConstraint<TPartOfChoppedConstraintNode>()) {
  1193. TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(name));
  1194. }
  1195. if (const auto part = child->Tail().GetConstraint<TPartOfUniqueConstraintNode>()) {
  1196. TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(name));
  1197. }
  1198. if (const auto part = child->Tail().GetConstraint<TPartOfDistinctConstraintNode>()) {
  1199. TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(name));
  1200. }
  1201. if (const auto valueNode = SkipModifiers(&child->Tail()); TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) {
  1202. structConstraints.push_back(&valueNode->Head().GetConstraintSet());
  1203. } else if (valueNode->Type() == TExprNode::Argument) {
  1204. structConstraints.push_back(&valueNode->GetConstraintSet());
  1205. }
  1206. }
  1207. if (!sorted.empty()) {
  1208. input->AddConstraint(ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(sorted)));
  1209. }
  1210. if (!chopped.empty()) {
  1211. input->AddConstraint(ctx.MakeConstraint<TPartOfChoppedConstraintNode>(std::move(chopped)));
  1212. }
  1213. if (!uniques.empty()) {
  1214. input->AddConstraint(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(uniques)));
  1215. }
  1216. if (!distincts.empty()) {
  1217. input->AddConstraint(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(distincts)));
  1218. }
  1219. if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) {
  1220. input->AddConstraint(varIndex);
  1221. }
  1222. return TStatus::Ok;
  1223. }
  1224. TStatus FlattenMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1225. TPartOfSortedConstraintNode::TMapType sorted;
  1226. TPartOfChoppedConstraintNode::TMapType chopped;
  1227. TPartOfUniqueConstraintNode::TMapType uniques;
  1228. TPartOfDistinctConstraintNode::TMapType distincts;
  1229. for (const auto& child : input->Children()) {
  1230. const auto& prefix = child->Head().Content();
  1231. if (const auto part = child->Tail().GetConstraint<TPartOfSortedConstraintNode>()) {
  1232. TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(ctx, prefix));
  1233. }
  1234. if (const auto part = child->Tail().GetConstraint<TPartOfChoppedConstraintNode>()) {
  1235. TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(ctx, prefix));
  1236. }
  1237. if (const auto part = child->Tail().GetConstraint<TPartOfUniqueConstraintNode>()) {
  1238. TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(ctx, prefix));
  1239. }
  1240. if (const auto part = child->Tail().GetConstraint<TPartOfDistinctConstraintNode>()) {
  1241. TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(ctx, prefix));
  1242. }
  1243. }
  1244. if (!sorted.empty()) {
  1245. input->AddConstraint(ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(sorted)));
  1246. }
  1247. if (!chopped.empty()) {
  1248. input->AddConstraint(ctx.MakeConstraint<TPartOfChoppedConstraintNode>(std::move(chopped)));
  1249. }
  1250. if (!uniques.empty()) {
  1251. input->AddConstraint(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(uniques)));
  1252. }
  1253. if (!distincts.empty()) {
  1254. input->AddConstraint(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(distincts)));
  1255. }
  1256. return TStatus::Ok;
  1257. }
  1258. template<class TPartOfConstraint>
  1259. static void AddPartOf(const TExprNode::TPtr& input, TExprContext& ctx) {
  1260. typename TPartOfConstraint::TMapType map;
  1261. if (const auto part = input->Head().GetConstraint<TPartOfConstraint>()) {
  1262. map = part->GetColumnMapping();
  1263. }
  1264. if (const auto part = input->Tail().GetConstraint<TPartOfConstraint>()) {
  1265. TPartOfConstraint::UniqueMerge(map, part->GetColumnMapping(input->Child(1)->Content()));
  1266. }
  1267. if (!map.empty()) {
  1268. input->AddConstraint(ctx.MakeConstraint<TPartOfConstraint>(std::move(map)));
  1269. }
  1270. }
  1271. template<class TPartOfConstraint>
  1272. static void ReplacePartOf(const TExprNode::TPtr& input, TExprContext& ctx) {
  1273. typename TPartOfConstraint::TMapType sorted;
  1274. const auto& name = input->Child(1)->Content();
  1275. if (const auto part = input->Head().GetConstraint<TPartOfConstraint>()) {
  1276. if (const auto filtered = part->FilterFields(ctx, [&name](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && path.front() != name; })) {
  1277. sorted = filtered->GetColumnMapping();
  1278. }
  1279. }
  1280. if (const auto part = input->Tail().GetConstraint<TPartOfConstraint>()) {
  1281. TPartOfConstraint::UniqueMerge(sorted, part->GetColumnMapping(name));
  1282. }
  1283. if (!sorted.empty()) {
  1284. input->AddConstraint(ctx.MakeConstraint<TPartOfConstraint>(std::move(sorted)));
  1285. }
  1286. }
  1287. TStatus AddMemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1288. const auto& addStructNode = input->Head();
  1289. const auto& extraFieldNode = input->Tail();
  1290. if (const auto emptyConstraint = addStructNode.GetConstraint<TEmptyConstraintNode>()) {
  1291. input->AddConstraint(emptyConstraint);
  1292. }
  1293. AddPartOf<TPartOfSortedConstraintNode>(input, ctx);
  1294. AddPartOf<TPartOfChoppedConstraintNode>(input, ctx);
  1295. AddPartOf<TPartOfUniqueConstraintNode>(input, ctx);
  1296. AddPartOf<TPartOfDistinctConstraintNode>(input, ctx);
  1297. TVector<const TConstraintSet*> structConstraints;
  1298. structConstraints.push_back(&addStructNode.GetConstraintSet());
  1299. if (const auto& valueNode = SkipModifiers(&extraFieldNode); TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) {
  1300. structConstraints.push_back(&valueNode->Head().GetConstraintSet());
  1301. } else if (valueNode->Type() == TExprNode::Argument) {
  1302. structConstraints.push_back(&valueNode->GetConstraintSet());
  1303. }
  1304. if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) {
  1305. input->AddConstraint(varIndex);
  1306. }
  1307. return TStatus::Ok;
  1308. }
  1309. TStatus RemoveMemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1310. const auto& name = input->Tail().Content();
  1311. const auto filter = [&name](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && path.front() != name; };
  1312. FilterFromHead<TPartOfSortedConstraintNode>(input, filter, ctx);
  1313. FilterFromHead<TPartOfChoppedConstraintNode>(input, filter, ctx);
  1314. FilterFromHead<TPartOfUniqueConstraintNode>(input, filter, ctx);
  1315. FilterFromHead<TPartOfDistinctConstraintNode>(input, filter, ctx);
  1316. return FromFirst<TVarIndexConstraintNode>(input, output, ctx);
  1317. }
  1318. TStatus ReplaceMemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1319. TVector<const TConstraintSet*> structConstraints;
  1320. structConstraints.push_back(&input->Head().GetConstraintSet());
  1321. ReplacePartOf<TPartOfSortedConstraintNode>(input, ctx);
  1322. ReplacePartOf<TPartOfChoppedConstraintNode>(input, ctx);
  1323. ReplacePartOf<TPartOfUniqueConstraintNode>(input, ctx);
  1324. ReplacePartOf<TPartOfDistinctConstraintNode>(input, ctx);
  1325. if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) {
  1326. input->AddConstraint(varIndex);
  1327. }
  1328. return TStatus::Ok;
  1329. }
  1330. TStatus ListWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1331. switch (input->ChildrenSize()) {
  1332. case 1:
  1333. return FromEmpty(input, output, ctx);
  1334. case 2:
  1335. return FromSecond<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  1336. default:
  1337. break;
  1338. }
  1339. return CommonFromChildren<1, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  1340. }
  1341. template<bool Ordered, bool WideInput, bool WithUnique = true>
  1342. static TSmallVec<TConstraintNode::TListType> GetConstraintsForInputArgument(const TExprNode& node, TExprContext& ctx) {
  1343. TSmallVec<TConstraintNode::TListType> argsConstraints(WideInput ? node.Child(1U)->Head().ChildrenSize() : 1U);
  1344. if constexpr (WideInput) {
  1345. if constexpr (Ordered) {
  1346. if (const auto& mapping = TPartOfSortedConstraintNode::GetCommonMapping(node.Head().GetConstraint<TSortedConstraintNode>(), node.Head().GetConstraint<TPartOfSortedConstraintNode>()); !mapping.empty()) {
  1347. for (ui32 i = 0U; i < argsConstraints.size(); ++i) {
  1348. if (auto extracted = TPartOfSortedConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) {
  1349. argsConstraints[i].emplace_back(ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(extracted)));
  1350. }
  1351. }
  1352. }
  1353. if (const auto& mapping = TPartOfChoppedConstraintNode::GetCommonMapping(node.Head().GetConstraint<TChoppedConstraintNode>(), node.Head().GetConstraint<TPartOfChoppedConstraintNode>()); !mapping.empty()) {
  1354. for (ui32 i = 0U; i < argsConstraints.size(); ++i) {
  1355. if (auto extracted = TPartOfChoppedConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) {
  1356. argsConstraints[i].emplace_back(ctx.MakeConstraint<TPartOfChoppedConstraintNode>(std::move(extracted)));
  1357. }
  1358. }
  1359. }
  1360. }
  1361. if constexpr (WithUnique) {
  1362. if (const auto& mapping = TPartOfUniqueConstraintNode::GetCommonMapping(node.Head().GetConstraint<TUniqueConstraintNode>(), node.Head().GetConstraint<TPartOfUniqueConstraintNode>()); !mapping.empty()) {
  1363. for (ui32 i = 0U; i < argsConstraints.size(); ++i) {
  1364. if (auto extracted = TPartOfUniqueConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) {
  1365. argsConstraints[i].emplace_back(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(extracted)));
  1366. }
  1367. }
  1368. }
  1369. if (const auto& mapping = TPartOfDistinctConstraintNode::GetCommonMapping(node.Head().GetConstraint<TDistinctConstraintNode>(), node.Head().GetConstraint<TPartOfDistinctConstraintNode>()); !mapping.empty()) {
  1370. for (ui32 i = 0U; i < argsConstraints.size(); ++i) {
  1371. if (auto extracted = TPartOfDistinctConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) {
  1372. argsConstraints[i].emplace_back(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(extracted)));
  1373. }
  1374. }
  1375. }
  1376. }
  1377. } else {
  1378. if (const auto inItemType = GetSeqItemType(node.Head().GetTypeAnn())) {
  1379. if (inItemType->GetKind() == ETypeAnnotationKind::Variant) {
  1380. if (inItemType->Cast<TVariantExprType>()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) {
  1381. const auto tupleType = inItemType->Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>();
  1382. argsConstraints.front().push_back(ctx.MakeConstraint<TVarIndexConstraintNode>(*inItemType->Cast<TVariantExprType>()));
  1383. TMultiConstraintNode::TMapType multiItems;
  1384. multiItems.reserve(tupleType->GetSize());
  1385. for (size_t i = 0; i < tupleType->GetSize(); ++i) {
  1386. multiItems.emplace_back(i, TConstraintSet{});
  1387. if constexpr (WithUnique) {
  1388. const auto inputMulti = node.Head().GetConstraint<TMultiConstraintNode>();
  1389. if (const auto inputConstr = inputMulti ? inputMulti->GetItem(i) : nullptr) {
  1390. if (auto mapping = TPartOfUniqueConstraintNode::GetCommonMapping(inputConstr->GetConstraint<TUniqueConstraintNode>(), inputConstr->GetConstraint<TPartOfUniqueConstraintNode>()); !mapping.empty()) {
  1391. multiItems.back().second.AddConstraint(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(mapping)));
  1392. }
  1393. if (auto mapping = TPartOfDistinctConstraintNode::GetCommonMapping(inputConstr->GetConstraint<TDistinctConstraintNode>(), inputConstr->GetConstraint<TPartOfDistinctConstraintNode>()); !mapping.empty()) {
  1394. multiItems.back().second.AddConstraint(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(mapping)));
  1395. }
  1396. }
  1397. }
  1398. }
  1399. if (!multiItems.empty()) {
  1400. argsConstraints.front().emplace_back(ctx.MakeConstraint<TMultiConstraintNode>(std::move(multiItems)));
  1401. }
  1402. }
  1403. } else {
  1404. if constexpr (Ordered) {
  1405. if (auto mapping = TPartOfSortedConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint<TSortedConstraintNode>(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint<TPartOfSortedConstraintNode>()); !mapping.empty()) {
  1406. argsConstraints.front().emplace_back(ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(mapping)));
  1407. }
  1408. if (auto mapping = TPartOfChoppedConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint<TChoppedConstraintNode>(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint<TPartOfChoppedConstraintNode>()); !mapping.empty()) {
  1409. argsConstraints.front().emplace_back(ctx.MakeConstraint<TPartOfChoppedConstraintNode>(std::move(mapping)));
  1410. }
  1411. }
  1412. if constexpr (WithUnique) {
  1413. if (auto mapping = TPartOfUniqueConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint<TUniqueConstraintNode>(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint<TPartOfUniqueConstraintNode>()); !mapping.empty()) {
  1414. argsConstraints.front().emplace_back(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(mapping)));
  1415. }
  1416. if (auto mapping = TPartOfDistinctConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint<TDistinctConstraintNode>(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint<TPartOfDistinctConstraintNode>()); !mapping.empty()) {
  1417. argsConstraints.front().emplace_back(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(mapping)));
  1418. }
  1419. }
  1420. }
  1421. }
  1422. }
  1423. return argsConstraints;
  1424. }
  1425. TStatus DictWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1426. const std::vector<std::string_view> k(1U, ctx.GetIndexAsString(0U));
  1427. input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(k));
  1428. input->AddConstraint(ctx.MakeConstraint<TDistinctConstraintNode>(k));
  1429. if (input->ChildrenSize() == 1) {
  1430. return FromEmpty(input, output, ctx);
  1431. }
  1432. return TStatus::Ok;
  1433. }
  1434. TStatus DictFromKeysWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1435. if (input->Child(1)->ChildrenSize() == 0) {
  1436. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1437. }
  1438. const std::vector<std::string_view> k(1U, ctx.GetIndexAsString(0U));
  1439. input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(k));
  1440. input->AddConstraint(ctx.MakeConstraint<TDistinctConstraintNode>(k));
  1441. return TStatus::Ok;
  1442. }
  1443. template<bool IsList, bool IsFlat>
  1444. TStatus PassOrEmptyWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1445. if (const auto part = input->Tail().GetConstraint<TPartOfSortedConstraintNode>())
  1446. if (const auto filtered = part->CompleteOnly(ctx))
  1447. input->AddConstraint(filtered);
  1448. if (const auto part = input->Tail().GetConstraint<TPartOfChoppedConstraintNode>())
  1449. if (const auto filtered = part->CompleteOnly(ctx))
  1450. input->AddConstraint(filtered);
  1451. if (const auto part = input->Tail().GetConstraint<TPartOfDistinctConstraintNode>())
  1452. if (const auto filtered = part->CompleteOnly(ctx))
  1453. input->AddConstraint(filtered);
  1454. if (const auto part = input->Tail().GetConstraint<TPartOfUniqueConstraintNode>())
  1455. if constexpr (IsList) {
  1456. if (const auto filtered = part->CompleteOnly(ctx))
  1457. input->AddConstraint(filtered);
  1458. } else
  1459. input->AddConstraint(part);
  1460. if constexpr (IsFlat) {
  1461. if (const auto empty = input->Tail().GetConstraint<TEmptyConstraintNode>())
  1462. input->AddConstraint(empty);
  1463. }
  1464. return FromSecond<TUniqueConstraintNode, TDistinctConstraintNode, TSortedConstraintNode, TChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  1465. }
  1466. TStatus IfWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
  1467. std::vector<const TConstraintSet*> constraints;
  1468. constraints.reserve((input->ChildrenSize() << 1U) + 1U);
  1469. constraints.emplace_back(&input->Tail().GetConstraintSet());
  1470. for (auto i = 0U; i < input->ChildrenSize() - 1U; ++i) {
  1471. constraints.emplace_back(&input->Child(++i)->GetConstraintSet());
  1472. }
  1473. if (constraints.empty())
  1474. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1475. else if (1U == constraints.size())
  1476. input->SetConstraints(**constraints.cbegin());
  1477. else
  1478. TApplyCommonConstraint<TSortedConstraintNode
  1479. , TPartOfSortedConstraintNode
  1480. , TChoppedConstraintNode
  1481. , TPartOfChoppedConstraintNode
  1482. , TUniqueConstraintNode
  1483. , TPartOfUniqueConstraintNode
  1484. , TDistinctConstraintNode
  1485. , TPartOfDistinctConstraintNode
  1486. , TEmptyConstraintNode
  1487. , TVarIndexConstraintNode
  1488. , TMultiConstraintNode
  1489. >::Do(input, constraints, ctx);
  1490. return TStatus::Ok;
  1491. }
  1492. TStatus IfPresentWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
  1493. auto optionals = input->ChildrenList();
  1494. const auto lambdaIndex = optionals.size() - 2U;
  1495. auto lambda = std::move(optionals[lambdaIndex]);
  1496. optionals.resize(lambdaIndex);
  1497. std::vector<const TConstraintNode::TListType> constraints;
  1498. constraints.reserve(optionals.size());
  1499. std::transform(optionals.cbegin(), optionals.cend(), std::back_inserter(constraints), [](const TExprNode::TPtr& node){ return node->GetAllConstraints(); });
  1500. if (const auto status = UpdateLambdaConstraints(input->ChildRef(lambdaIndex), ctx, constraints); status != TStatus::Ok) {
  1501. return status;
  1502. }
  1503. if (std::any_of(optionals.cbegin(), optionals.cend(), [] (const TExprNode::TPtr& node) { return bool(node->GetConstraint<TEmptyConstraintNode>()); })) {
  1504. input->CopyConstraints(input->Tail());
  1505. return TStatus::Ok;
  1506. }
  1507. const std::vector<const TConstraintSet*> both = { &lambda->GetConstraintSet(), &input->Tail().GetConstraintSet() };
  1508. TApplyCommonConstraint<TPartOfSortedConstraintNode
  1509. , TPartOfChoppedConstraintNode
  1510. , TPartOfUniqueConstraintNode
  1511. , TPartOfDistinctConstraintNode
  1512. , TEmptyConstraintNode
  1513. , TVarIndexConstraintNode
  1514. , TMultiConstraintNode
  1515. >::Do(input, both, ctx);
  1516. return TStatus::Ok;
  1517. }
  1518. TStatus ReverseWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1519. if (const auto sorted = input->Head().GetConstraint<TSortedConstraintNode>()) {
  1520. auto content = sorted->GetContent();
  1521. std::for_each(content.begin(), content.end(), [](std::pair<TPartOfConstraintBase::TSetType, bool>& pair) { pair.second = !pair.second; });
  1522. input->AddConstraint(ctx.MakeConstraint<TSortedConstraintNode>(std::move(content)));
  1523. }
  1524. return FromFirst<TEmptyConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
  1525. }
  1526. TStatus SwitchWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  1527. TStatus status = TStatus::Ok;
  1528. TDynBitMap outFromChildren; // children, from which take a multi constraint for output
  1529. if (const auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) {
  1530. for (size_t i = 2; i < input->ChildrenSize(); ++i) {
  1531. TMultiConstraintNode::TMapType items;
  1532. ui32 lambdaInputIndex = 0;
  1533. for (auto& child : input->Child(i)->Children()) {
  1534. const ui32 index = FromString<ui32>(child->Content());
  1535. if (auto c = multi->GetItem(index)) {
  1536. items[lambdaInputIndex] = *c;
  1537. outFromChildren.Set(i + 1);
  1538. }
  1539. ++lambdaInputIndex;
  1540. }
  1541. TConstraintNode::TListType argConstraints;
  1542. if (!items.empty()) {
  1543. if (input->Child(i)->ChildrenSize() > 1) {
  1544. argConstraints.push_back(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items)));
  1545. argConstraints.push_back(ctx.MakeConstraint<TVarIndexConstraintNode>(input->Child(i)->ChildrenSize()));
  1546. } else {
  1547. argConstraints = items.front().second.GetAllConstraints();
  1548. }
  1549. }
  1550. status = status.Combine(UpdateLambdaConstraints(input->ChildRef(++i), ctx, {argConstraints}));
  1551. }
  1552. } else {
  1553. const bool inVar = GetSeqItemType(*input->Head().GetTypeAnn()).GetKind() == ETypeAnnotationKind::Variant;
  1554. const TSmallVec<TConstraintNode::TListType> argConstraints(1U, inVar ? TConstraintNode::TListType() : input->Head().GetAllConstraints());
  1555. for (size_t i = 3; i < input->ChildrenSize(); i += 2) {
  1556. status = status.Combine(UpdateLambdaConstraints(input->ChildRef(i), ctx, argConstraints));
  1557. }
  1558. outFromChildren.Set(0, input->ChildrenSize());
  1559. }
  1560. if (status != TStatus::Ok) {
  1561. return status;
  1562. }
  1563. const auto inputVarIndex = input->Head().GetConstraint<TVarIndexConstraintNode>();
  1564. const bool emptyInput = input->Head().GetConstraint<TEmptyConstraintNode>();
  1565. if (GetSeqItemType(*input->GetTypeAnn()).GetKind() == ETypeAnnotationKind::Variant) {
  1566. ui32 outIndexOffset = 0;
  1567. TMultiConstraintNode::TMapType multiItems;
  1568. TVarIndexConstraintNode::TMapType remapItems;
  1569. bool emptyOut = true;
  1570. for (size_t i = 2; i < input->ChildrenSize(); i += 2) {
  1571. const auto lambda = input->Child(i + 1);
  1572. const auto& lambdaItemType = GetSeqItemType(*lambda->GetTypeAnn());
  1573. if (inputVarIndex) {
  1574. if (auto varIndex = lambda->GetConstraint<TVarIndexConstraintNode>()) {
  1575. for (auto& item: varIndex->GetIndexMapping()) {
  1576. YQL_ENSURE(item.second < input->Child(i)->ChildrenSize());
  1577. const auto srcIndex = FromString<size_t>(input->Child(i)->Child(item.second)->Content());
  1578. remapItems.push_back(std::make_pair(outIndexOffset + item.first, srcIndex));
  1579. }
  1580. } else if (lambdaItemType.GetKind() == ETypeAnnotationKind::Variant && input->Child(i)->ChildrenSize() == 1) {
  1581. const auto srcIndex = FromString<size_t>(input->Child(i)->Head().Content());
  1582. for (size_t j = 0; j < lambdaItemType.Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>()->GetSize(); ++j) {
  1583. remapItems.push_back(std::make_pair(outIndexOffset + j, srcIndex));
  1584. }
  1585. } else if (lambdaItemType.GetKind() != ETypeAnnotationKind::Variant && input->Child(i)->ChildrenSize() > 1) {
  1586. for (auto& child : input->Child(i)->Children()) {
  1587. const auto srcIndex = FromString<size_t>(child->Content());
  1588. remapItems.push_back(std::make_pair(outIndexOffset, srcIndex));
  1589. }
  1590. }
  1591. }
  1592. const bool lambdaEmpty = lambda->GetConstraint<TEmptyConstraintNode>();
  1593. if (!lambdaEmpty) {
  1594. emptyOut = false;
  1595. }
  1596. if (lambdaItemType.GetKind() == ETypeAnnotationKind::Variant) {
  1597. if (!emptyInput && outFromChildren.Test(i + 1)) {
  1598. if (auto multi = lambda->GetConstraint<TMultiConstraintNode>()) {
  1599. for (auto& item: multi->GetItems()) {
  1600. multiItems.insert_unique(std::make_pair(outIndexOffset + item.first, item.second));
  1601. }
  1602. }
  1603. }
  1604. outIndexOffset += lambdaItemType.Cast<TVariantExprType>()->GetUnderlyingType()->Cast<TTupleExprType>()->GetSize();
  1605. } else {
  1606. if (!emptyInput && outFromChildren.Test(i + 1) && !lambdaEmpty) {
  1607. multiItems[outIndexOffset] = lambda->GetConstraintSet();
  1608. }
  1609. ++outIndexOffset;
  1610. }
  1611. }
  1612. if (inputVarIndex && !remapItems.empty()) {
  1613. TVarIndexConstraintNode::TMapType result;
  1614. for (auto& item: remapItems) {
  1615. auto range = inputVarIndex->GetIndexMapping().equal_range(item.second);
  1616. for (auto it = range.first; it != range.second; ++it) {
  1617. result.push_back(std::make_pair(item.first, it->second));
  1618. }
  1619. }
  1620. if (!result.empty()) {
  1621. ::Sort(result);
  1622. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(result)));
  1623. }
  1624. }
  1625. if (!multiItems.empty()) {
  1626. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(multiItems)));
  1627. }
  1628. if (emptyOut) {
  1629. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1630. }
  1631. } else {
  1632. YQL_ENSURE(input->ChildrenSize() == 4);
  1633. input->CopyConstraints(*input->Child(3));
  1634. }
  1635. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  1636. }
  1637. TStatus VisitWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
  1638. TStatus status = TStatus::Ok;
  1639. TDynBitMap outFromChildren; // children, from which take a multi constraint for output
  1640. TDynBitMap usedAlts;
  1641. const auto inMulti = input->Head().GetConstraint<TMultiConstraintNode>();
  1642. for (ui32 i = 1; i < input->ChildrenSize(); ++i) {
  1643. if (const auto child = input->Child(i); child->IsAtom()) {
  1644. TSmallVec<TConstraintNode::TListType> argConstraints(1U);
  1645. if (inMulti) {
  1646. const auto index = FromString<ui32>(child->Content());
  1647. usedAlts.Set(index);
  1648. if (const auto c = inMulti->GetItem(index)) {
  1649. argConstraints.front() = c->GetAllConstraints();
  1650. outFromChildren.Set(i + 1U);
  1651. }
  1652. }
  1653. status = status.Combine(UpdateLambdaConstraints(input->ChildRef(++i), ctx, argConstraints));
  1654. } else if (inMulti) { // Check that we can fall to default branch
  1655. for (auto& item: inMulti->GetItems()) {
  1656. if (!usedAlts.Test(item.first)) {
  1657. outFromChildren.Set(i);
  1658. break;
  1659. }
  1660. }
  1661. }
  1662. }
  1663. if (status != TStatus::Ok) {
  1664. return status;
  1665. }
  1666. if (!inMulti) {
  1667. outFromChildren.Set(0, input->ChildrenSize());
  1668. }
  1669. auto outType = input->GetTypeAnn();
  1670. if (auto t = GetSeqItemType(outType)) {
  1671. outType = t;
  1672. }
  1673. if (outType->GetKind() == ETypeAnnotationKind::Variant && outType->Cast<TVariantExprType>()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) {
  1674. TVector<const TConstraintSet*> outConstraints;
  1675. TVarIndexConstraintNode::TMapType remapItems;
  1676. for (ui32 i = 1; i < input->ChildrenSize(); ++i) {
  1677. if (input->Child(i)->IsAtom()) {
  1678. ui32 index = FromString<ui32>(input->Child(i)->Content());
  1679. ++i;
  1680. if (outFromChildren.Test(i)) {
  1681. outConstraints.push_back(&input->Child(i)->GetConstraintSet());
  1682. if (const auto outMulti = input->Child(i)->GetConstraint<TMultiConstraintNode>()) {
  1683. for (auto& item: outMulti->GetItems()) {
  1684. remapItems.push_back(std::make_pair(item.first, index));
  1685. }
  1686. }
  1687. }
  1688. } else {
  1689. if (outFromChildren.Test(i)) {
  1690. outConstraints.push_back(&input->Child(i)->GetConstraintSet());
  1691. const auto outMulti = input->Child(i)->GetConstraint<TMultiConstraintNode>();
  1692. if (outMulti && inMulti) {
  1693. for (auto& outItem: outMulti->GetItems()) {
  1694. for (auto& inItem: inMulti->GetItems()) {
  1695. if (!usedAlts.Test(inItem.first)) {
  1696. remapItems.push_back(std::make_pair(outItem.first, inItem.first));
  1697. }
  1698. }
  1699. }
  1700. }
  1701. }
  1702. }
  1703. }
  1704. if (auto multi = TMultiConstraintNode::MakeCommon(outConstraints, ctx)) {
  1705. input->AddConstraint(multi);
  1706. }
  1707. if (auto empty = TEmptyConstraintNode::MakeCommon(outConstraints, ctx)) {
  1708. input->AddConstraint(empty);
  1709. }
  1710. if (auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  1711. TVarIndexConstraintNode::TMapType varIndexItems;
  1712. for (auto& item: remapItems) {
  1713. const auto range = varIndex->GetIndexMapping().equal_range(item.second);
  1714. for (auto i = range.first; i != range.second; ++i) {
  1715. varIndexItems.push_back(std::make_pair(item.first, i->second));
  1716. }
  1717. }
  1718. if (!varIndexItems.empty()) {
  1719. ::Sort(varIndexItems);
  1720. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(varIndexItems)));
  1721. }
  1722. }
  1723. }
  1724. else {
  1725. std::vector<const TConstraintSet*> nonEmpty;
  1726. for (ui32 i = 1; i < input->ChildrenSize(); ++i) {
  1727. if (input->Child(i)->IsAtom()) {
  1728. ++i;
  1729. }
  1730. if (outFromChildren.Test(i)) {
  1731. nonEmpty.push_back(&input->Child(i)->GetConstraintSet());
  1732. }
  1733. }
  1734. EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint<TEmptyConstraintNode>(); });
  1735. if (nonEmpty.empty()) {
  1736. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1737. } else if (nonEmpty.size() == 1) {
  1738. input->SetConstraints(*nonEmpty.front());
  1739. }
  1740. if (auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  1741. TVarIndexConstraintNode::TMapType varIndexItems;
  1742. for (ui32 i = 1; i < input->ChildrenSize(); ++i) {
  1743. if (input->Child(i)->IsAtom()) {
  1744. const auto index = FromString<ui32>(input->Child(i++)->Content());
  1745. if (outFromChildren.Test(i) && IsDepended(input->Child(i)->Tail(), input->Child(i)->Head().Head())) { // Somehow depends on arg
  1746. const auto range = varIndex->GetIndexMapping().equal_range(index);
  1747. for (auto i = range.first; i != range.second; ++i) {
  1748. varIndexItems.push_back(std::make_pair(0, i->second));
  1749. }
  1750. }
  1751. }
  1752. else if (outFromChildren.Test(i)) {
  1753. if (inMulti) {
  1754. for (auto& inItem: inMulti->GetItems()) {
  1755. if (!usedAlts.Test(inItem.first)) {
  1756. auto range = varIndex->GetIndexMapping().equal_range(inItem.first);
  1757. for (auto i = range.first; i != range.second; ++i) {
  1758. varIndexItems.push_back(std::make_pair(0, i->second));
  1759. }
  1760. }
  1761. }
  1762. }
  1763. }
  1764. }
  1765. if (!varIndexItems.empty()) {
  1766. ::SortUnique(varIndexItems);
  1767. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(varIndexItems)));
  1768. }
  1769. }
  1770. }
  1771. return TStatus::Ok;
  1772. }
  1773. TStatus VariantItemWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1774. auto inputType = input->Head().GetTypeAnn();
  1775. if (inputType->GetKind() == ETypeAnnotationKind::Optional) {
  1776. inputType = inputType->Cast<TOptionalExprType>()->GetItemType();
  1777. }
  1778. const auto underlyingType = inputType->Cast<TVariantExprType>()->GetUnderlyingType();
  1779. if (underlyingType->GetKind() == ETypeAnnotationKind::Tuple) {
  1780. if (auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) {
  1781. std::vector<TMultiConstraintNode::TMapType::value_type> nonEmpty;
  1782. std::copy_if(multi->GetItems().begin(), multi->GetItems().end(), std::back_inserter(nonEmpty),
  1783. [] (const TMultiConstraintNode::TMapType::value_type& v) {
  1784. return !v.second.GetConstraint<TEmptyConstraintNode>();
  1785. }
  1786. );
  1787. if (nonEmpty.empty()) {
  1788. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1789. } else if (nonEmpty.size() == 1) {
  1790. input->SetConstraints(nonEmpty.front().second);
  1791. }
  1792. }
  1793. if (auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  1794. TVarIndexConstraintNode::TMapType varIndexItems;
  1795. for (auto& item: varIndex->GetIndexMapping()) {
  1796. varIndexItems.push_back(std::make_pair(0, item.second));
  1797. }
  1798. if (!varIndexItems.empty()) {
  1799. ::SortUnique(varIndexItems);
  1800. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(varIndexItems)));
  1801. }
  1802. }
  1803. }
  1804. return TStatus::Ok;
  1805. }
  1806. TStatus VariantWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1807. if (input->GetTypeAnn()->Cast<TVariantExprType>()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) {
  1808. const auto index = FromString<ui32>(input->Child(1)->Content());
  1809. TConstraintSet target;
  1810. CopyExcept(target, input->Head().GetConstraintSet(), TVarIndexConstraintNode::Name());
  1811. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(index, target));
  1812. if (auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  1813. TVarIndexConstraintNode::TMapType filteredItems;
  1814. for (auto& item: varIndex->GetIndexMapping()) {
  1815. filteredItems.push_back(std::make_pair(index, item.second));
  1816. }
  1817. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(filteredItems)));
  1818. }
  1819. }
  1820. return TStatus::Ok;
  1821. }
  1822. TStatus DynamicVariantWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1823. if (auto underlyingType = RemoveOptionalType(input->GetTypeAnn())->Cast<TVariantExprType>()->GetUnderlyingType(); underlyingType->GetKind() == ETypeAnnotationKind::Tuple) {
  1824. TConstraintSet target;
  1825. CopyExcept(target, input->Head().GetConstraintSet(), TVarIndexConstraintNode::Name());
  1826. TMultiConstraintNode::TMapType items;
  1827. for (ui32 i = 0; i < underlyingType->Cast<TTupleExprType>()->GetSize(); ++i) {
  1828. items.emplace_back(i, target);
  1829. }
  1830. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items)));
  1831. if (auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  1832. TVarIndexConstraintNode::TMapType filteredItems;
  1833. for (ui32 i = 0; i < underlyingType->Cast<TTupleExprType>()->GetSize(); ++i) {
  1834. for (auto& item: varIndex->GetIndexMapping()) {
  1835. filteredItems.push_back(std::make_pair(i, item.second));
  1836. }
  1837. }
  1838. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(filteredItems)));
  1839. }
  1840. }
  1841. return TStatus::Ok;
  1842. }
  1843. TStatus GuessWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1844. auto inputType = input->Head().GetTypeAnn();
  1845. if (inputType->GetKind() == ETypeAnnotationKind::Optional) {
  1846. inputType = inputType->Cast<TOptionalExprType>()->GetItemType();
  1847. }
  1848. const auto underlyingType = inputType->Cast<TVariantExprType>()->GetUnderlyingType();
  1849. if (underlyingType->GetKind() == ETypeAnnotationKind::Tuple) {
  1850. const auto guessIndex = FromString<ui32>(input->Child(1)->Content());
  1851. if (auto multi = input->Head().GetConstraint<TMultiConstraintNode>()) {
  1852. if (auto c = multi->GetItem(guessIndex)) {
  1853. input->SetConstraints(*c);
  1854. } else {
  1855. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1856. }
  1857. }
  1858. if (auto varIndex = input->Head().GetConstraint<TVarIndexConstraintNode>()) {
  1859. TVarIndexConstraintNode::TMapType filteredItems;
  1860. for (auto& item: varIndex->GetIndexMapping()) {
  1861. if (item.first == guessIndex) {
  1862. filteredItems.push_back(std::make_pair(0, item.second));
  1863. }
  1864. }
  1865. if (!filteredItems.empty()) {
  1866. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(filteredItems)));
  1867. }
  1868. }
  1869. }
  1870. return TStatus::Ok;
  1871. }
  1872. TStatus MuxWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1873. const auto listItemType = GetSeqItemType(input->GetTypeAnn());
  1874. if (!listItemType) {
  1875. return TStatus::Ok;
  1876. }
  1877. if (listItemType->Cast<TVariantExprType>()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) {
  1878. if (input->Head().IsList()) {
  1879. TMultiConstraintNode::TMapType items;
  1880. ui32 index = 0;
  1881. for (auto& child: input->Head().Children()) {
  1882. items.push_back(std::make_pair(index, child->GetConstraintSet()));
  1883. ++index;
  1884. }
  1885. if (!items.empty()) {
  1886. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(items)));
  1887. }
  1888. }
  1889. }
  1890. return TStatus::Ok;
  1891. }
  1892. TStatus NthWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1893. const auto& memberName = input->Tail().Content();
  1894. const auto& structNode = input->Head();
  1895. if (const auto emptyConstraint = structNode.GetConstraint<TEmptyConstraintNode>()) {
  1896. input->AddConstraint(emptyConstraint);
  1897. } else {
  1898. if (const auto part = structNode.GetConstraint<TPartOfSortedConstraintNode>()) {
  1899. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1900. input->AddConstraint(extracted);
  1901. }
  1902. }
  1903. if (const auto part = structNode.GetConstraint<TPartOfChoppedConstraintNode>()) {
  1904. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1905. input->AddConstraint(extracted);
  1906. }
  1907. }
  1908. if (const auto part = structNode.GetConstraint<TPartOfUniqueConstraintNode>()) {
  1909. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1910. input->AddConstraint(extracted);
  1911. }
  1912. }
  1913. if (const auto part = structNode.GetConstraint<TPartOfDistinctConstraintNode>()) {
  1914. if (const auto extracted = part->ExtractField(ctx, memberName)) {
  1915. input->AddConstraint(extracted);
  1916. }
  1917. }
  1918. }
  1919. if (input->Head().IsList()) {
  1920. input->CopyConstraints(*input->Head().Child(FromString<ui32>(input->Child(1)->Content())));
  1921. }
  1922. else if (input->Head().IsCallable("Demux")) {
  1923. if (auto multi = input->Head().Head().GetConstraint<TMultiConstraintNode>()) {
  1924. if (auto c = multi->GetItem(FromString<ui32>(input->Child(1)->Content()))) {
  1925. input->SetConstraints(*c);
  1926. }
  1927. }
  1928. }
  1929. return TStatus::Ok;
  1930. }
  1931. TStatus EquiJoinWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  1932. const auto numLists = input->ChildrenSize() - 2U;
  1933. std::vector<size_t> emptyInputs;
  1934. TJoinLabels labels;
  1935. for (auto i = 0U; i < numLists; ++i) {
  1936. const auto& list = input->Child(i)->Head();
  1937. if (list.GetConstraint<TEmptyConstraintNode>()) {
  1938. emptyInputs.push_back(i);
  1939. }
  1940. if (const auto err = labels.Add(ctx, input->Child(i)->Tail(),
  1941. GetSeqItemType(*list.GetTypeAnn()).Cast<TStructExprType>(),
  1942. GetDetailed(list.GetConstraint<TUniqueConstraintNode>(), *list.GetTypeAnn(), ctx),
  1943. GetDetailed(list.GetConstraint<TDistinctConstraintNode>(), *list.GetTypeAnn(), ctx))) {
  1944. ctx.AddError(*err);
  1945. return TStatus::Error;
  1946. }
  1947. }
  1948. const auto joinTree = input->Child(numLists);
  1949. for (auto i: emptyInputs) {
  1950. if (IsRequiredSide(joinTree, labels, i).first) {
  1951. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  1952. break;
  1953. }
  1954. }
  1955. TJoinOptions options;
  1956. if (const auto status = ValidateEquiJoinOptions(input->Pos(), input->Tail(), options, ctx); status != IGraphTransformer::TStatus::Ok || options.Flatten) {
  1957. return status;
  1958. }
  1959. const TUniqueConstraintNode* unique = nullptr;
  1960. const TDistinctConstraintNode* distinct = nullptr;
  1961. if (const auto status = EquiJoinConstraints(input->Pos(), unique, distinct, labels, *joinTree, ctx); status != IGraphTransformer::TStatus::Ok) {
  1962. return status;
  1963. }
  1964. if (const auto renames = LoadJoinRenameMap(input->Tail()); !renames.empty() && (unique || distinct)) {
  1965. const auto rename = [&renames](const TPartOfConstraintBase::TPathType& path) -> std::vector<TPartOfConstraintBase::TPathType> {
  1966. if (path.empty())
  1967. return {};
  1968. const auto it = renames.find(path.front());
  1969. if (renames.cend() == it)
  1970. return {path};
  1971. if (it->second.empty())
  1972. return {};
  1973. std::vector<TPartOfConstraintBase::TPathType> res(it->second.size());
  1974. std::transform(it->second.cbegin(), it->second.cend(), res.begin(), [&path](const std::string_view& newName) {
  1975. auto newPath = path;
  1976. newPath.front() = newName;
  1977. return newPath;
  1978. });
  1979. return res;
  1980. };
  1981. if (unique)
  1982. unique = unique->RenameFields(ctx, rename);
  1983. if (distinct)
  1984. distinct = distinct->RenameFields(ctx, rename);
  1985. }
  1986. if (unique)
  1987. input->AddConstraint(unique->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  1988. if (distinct)
  1989. input->AddConstraint(distinct->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  1990. return TStatus::Ok;
  1991. }
  1992. static std::vector<std::string_view> GetKeys(const TExprNode& keys) {
  1993. std::vector<std::string_view> result;
  1994. result.reserve(keys.ChildrenSize());
  1995. keys.ForEachChild([&result](const TExprNode& key) { result.emplace_back(key.Content()); });
  1996. return result;
  1997. }
  1998. template<bool ForDict = false>
  1999. static TPartOfConstraintBase::TPathReduce GetRenames(const TExprNode& renames) {
  2000. std::unordered_map<std::string_view, std::string_view> map(renames.ChildrenSize() >> 1U);
  2001. for (auto i = 0U; i < renames.ChildrenSize(); ++++i)
  2002. map.emplace(renames.Child(i)->Content(), renames.Child(i + 1U)->Content());
  2003. return [map](const TPartOfConstraintBase::TPathType& path) -> std::vector<TPartOfConstraintBase::TPathType> {
  2004. if constexpr (ForDict) {
  2005. if (path.size() > 1U && path.front() == "1"sv) {
  2006. auto out = path;
  2007. out.pop_front();
  2008. if (const auto it = map.find(out.front()); map.cend() != it) {
  2009. out.front() = it->second;
  2010. return {std::move(out)};
  2011. }
  2012. }
  2013. } else {
  2014. if (!path.empty()) {
  2015. if (const auto it = map.find(path.front()); map.cend() != it) {
  2016. auto out = path;
  2017. out.front() = it->second;
  2018. return {std::move(out)};
  2019. }
  2020. }
  2021. }
  2022. return {};
  2023. };
  2024. }
  2025. TStatus MapJoinCoreWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  2026. const TCoMapJoinCore core(input);
  2027. const auto& joinType = core.JoinKind().Ref();
  2028. if (const auto empty = core.LeftInput().Ref().GetConstraint<TEmptyConstraintNode>()) {
  2029. input->AddConstraint(empty);
  2030. } else if (const auto empty = core.RightDict().Ref().GetConstraint<TEmptyConstraintNode>()) {
  2031. if (joinType.IsAtom({"Inner", "LeftSemi"})) {
  2032. input->AddConstraint(empty);
  2033. }
  2034. }
  2035. if (joinType.IsAtom({"LeftSemi", "LeftOnly"})) {
  2036. const auto rename = GetRenames(core.LeftRenames().Ref());
  2037. if (const auto unique = core.LeftInput().Ref().GetConstraint<TUniqueConstraintNode>())
  2038. if (const auto renamed = unique->RenameFields(ctx, rename))
  2039. input->AddConstraint(renamed);
  2040. if (const auto distinct = core.LeftInput().Ref().GetConstraint<TDistinctConstraintNode>())
  2041. if (const auto renamed = distinct->RenameFields(ctx, rename))
  2042. input->AddConstraint(renamed);
  2043. } else {
  2044. if (const auto unique = core.LeftInput().Ref().GetConstraint<TUniqueConstraintNode>()) {
  2045. if (unique->ContainsCompleteSet(GetKeys(core.LeftKeysColumns().Ref())) && core.RightDict().Ref().GetTypeAnn()->Cast<TDictExprType>()->GetPayloadType()->GetKind() != ETypeAnnotationKind::List) {
  2046. const auto rename = GetRenames(core.LeftRenames().Ref());
  2047. const auto rightRename = GetRenames<true>(core.RightRenames().Ref());
  2048. auto commonUnique = unique->RenameFields(ctx, rename);
  2049. if (const auto rUnique = core.RightDict().Ref().GetConstraint<TUniqueConstraintNode>()) {
  2050. commonUnique = TUniqueConstraintNode::Merge(commonUnique, rUnique->RenameFields(ctx, rightRename), ctx);
  2051. }
  2052. const auto distinct = core.LeftInput().Ref().GetConstraint<TDistinctConstraintNode>();
  2053. auto commonDistinct = distinct ? distinct->RenameFields(ctx, rename) : nullptr;
  2054. if (joinType.IsAtom("Inner")) {
  2055. if (const auto rDistinct = core.RightDict().Ref().GetConstraint<TDistinctConstraintNode>()) {
  2056. commonDistinct = TDistinctConstraintNode::Merge(commonDistinct, rDistinct->RenameFields(ctx, rightRename), ctx);
  2057. }
  2058. }
  2059. if (commonUnique)
  2060. input->AddConstraint(commonUnique);
  2061. if (commonDistinct)
  2062. input->AddConstraint(commonDistinct);
  2063. }
  2064. }
  2065. }
  2066. if (const auto sorted = core.LeftInput().Ref().GetConstraint<TSortedConstraintNode>())
  2067. if (const auto renamed = sorted->RenameFields(ctx, GetRenames(core.LeftRenames().Ref())))
  2068. input->AddConstraint(renamed);
  2069. if (const auto chopped = core.LeftInput().Ref().GetConstraint<TChoppedConstraintNode>())
  2070. if (const auto renamed = chopped->RenameFields(ctx, GetRenames(core.LeftRenames().Ref())))
  2071. input->AddConstraint(renamed);
  2072. return TStatus::Ok;
  2073. }
  2074. TExprNode::TPtr GraceJoinRightInput(const TCoGraceJoinCore& core) const {
  2075. return core.RightInput().Ptr();
  2076. }
  2077. TExprNode::TPtr GraceJoinRightInput(const TCoGraceSelfJoinCore& core) const {
  2078. return core.Input().Ptr();
  2079. }
  2080. TExprNode::TPtr GraceJoinLeftInput(const TCoGraceJoinCore& core) const {
  2081. return core.LeftInput().Ptr();
  2082. }
  2083. TExprNode::TPtr GraceJoinLeftInput(const TCoGraceSelfJoinCore& core) const {
  2084. return core.Input().Ptr();
  2085. }
  2086. template<typename GraceJoinCoreType>
  2087. TStatus GraceJoinCoreWrapImpl(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2088. Y_UNUSED(output);
  2089. const GraceJoinCoreType core(input);
  2090. const auto& joinType = core.JoinKind().Ref();
  2091. TExprNode::TPtr leftInput = GraceJoinLeftInput(core);
  2092. TExprNode::TPtr rightInput = GraceJoinRightInput(core);
  2093. if (const auto lEmpty = leftInput->GetConstraint<TEmptyConstraintNode>(), rEmpty = rightInput->GetConstraint<TEmptyConstraintNode>(); lEmpty && rEmpty) {
  2094. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  2095. } else if (lEmpty && joinType.Content().starts_with("Left")) {
  2096. input->AddConstraint(lEmpty);
  2097. } else if (rEmpty && joinType.Content().starts_with("Right")) {
  2098. input->AddConstraint(rEmpty);
  2099. } else if ((lEmpty || rEmpty) && (joinType.IsAtom("Inner") || joinType.Content().ends_with("Semi"))) {
  2100. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  2101. }
  2102. bool leftAny = false, rigthAny = false;
  2103. core.Flags().Ref().ForEachChild([&](const TExprNode& flag) {
  2104. if (flag.IsAtom("LeftAny"))
  2105. leftAny = true;
  2106. else if (flag.IsAtom("RightAny"))
  2107. rigthAny = true;
  2108. });
  2109. const TUniqueConstraintNode* lUnique = leftInput->GetConstraint<TUniqueConstraintNode>();
  2110. const TUniqueConstraintNode* rUnique = rightInput->GetConstraint<TUniqueConstraintNode>();
  2111. const bool lOneRow = lUnique && (leftAny || lUnique->ContainsCompleteSet(GetKeys(core.LeftKeysColumns().Ref())));
  2112. const bool rOneRow = rUnique && (rigthAny || rUnique->ContainsCompleteSet(GetKeys(core.RightKeysColumns().Ref())));
  2113. const bool singleSide = joinType.Content().ends_with("Semi") || joinType.Content().ends_with("Only");
  2114. if (singleSide || lOneRow || rOneRow) {
  2115. const TUniqueConstraintNode* unique = nullptr;
  2116. const TDistinctConstraintNode* distinct = nullptr;
  2117. const bool leftSide = joinType.Content().starts_with("Left");
  2118. const bool rightSide = joinType.Content().starts_with("Right");
  2119. const auto leftRename = GetRenames(core.LeftRenames().Ref());
  2120. const auto rightRename = GetRenames(core.RightRenames().Ref());
  2121. if (singleSide) {
  2122. if (leftSide && lUnique)
  2123. unique = lUnique->RenameFields(ctx, leftRename);
  2124. else if (rightSide && rUnique)
  2125. unique = rUnique->RenameFields(ctx, rightRename);
  2126. } else {
  2127. const bool exclusion = joinType.IsAtom("Exclusion");
  2128. const bool useLeft = lUnique && (rOneRow || exclusion);
  2129. const bool useRight = rUnique && (lOneRow || exclusion);
  2130. if (useLeft && !useRight)
  2131. unique = lUnique->RenameFields(ctx, leftRename);
  2132. else if (useRight && !useLeft)
  2133. unique = rUnique->RenameFields(ctx, rightRename);
  2134. else if (useLeft && useRight)
  2135. unique = TUniqueConstraintNode::Merge(lUnique->RenameFields(ctx, leftRename), rUnique->RenameFields(ctx, rightRename), ctx);
  2136. }
  2137. const auto lDistinct = leftInput->GetConstraint<TDistinctConstraintNode>();
  2138. const auto rDistinct = rightInput->GetConstraint<TDistinctConstraintNode>();
  2139. if (singleSide) {
  2140. if (leftSide && lDistinct)
  2141. distinct = lDistinct->RenameFields(ctx, leftRename);
  2142. else if (rightSide && rDistinct)
  2143. distinct = rDistinct->RenameFields(ctx, rightRename);
  2144. } else {
  2145. const bool inner = joinType.IsAtom("Inner");
  2146. const bool useLeft = lDistinct && rOneRow && (inner || leftSide);
  2147. const bool useRight = rDistinct && lOneRow && (inner || rightSide);
  2148. if (useLeft && !useRight)
  2149. distinct = lDistinct->RenameFields(ctx, leftRename);
  2150. else if (useRight && !useLeft)
  2151. distinct = rDistinct->RenameFields(ctx, rightRename);
  2152. else if (useLeft && useRight)
  2153. distinct = TDistinctConstraintNode::Merge(lDistinct->RenameFields(ctx, leftRename), rDistinct->RenameFields(ctx, rightRename), ctx);
  2154. }
  2155. if (unique)
  2156. input->AddConstraint(unique);
  2157. if (distinct)
  2158. input->AddConstraint(distinct);
  2159. }
  2160. return TStatus::Ok;
  2161. }
  2162. TStatus GraceJoinCoreWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2163. return GraceJoinCoreWrapImpl<TCoGraceJoinCore>(input, output, ctx);
  2164. }
  2165. TStatus GraceSelfJoinCoreWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2166. return GraceJoinCoreWrapImpl<TCoGraceSelfJoinCore>(input, output, ctx);
  2167. }
  2168. template<bool Distinct>
  2169. static const TUniqueConstraintNodeBase<Distinct>* GetForPayload(const TExprNode& input, TExprContext& ctx) {
  2170. if (const auto constraint = input.GetConstraint<TUniqueConstraintNodeBase<Distinct>>()) {
  2171. return constraint->RenameFields(ctx, [&ctx](const TPartOfConstraintBase::TPathType& path) -> std::vector<TPartOfConstraintBase::TPathType> {
  2172. if (path.empty() || path.front() != ctx.GetIndexAsString(1U))
  2173. return {};
  2174. auto copy = path;
  2175. copy.pop_front();
  2176. return {copy};
  2177. });
  2178. }
  2179. return nullptr;
  2180. }
  2181. TStatus JoinDictWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  2182. const TCoJoinDict join(input);
  2183. const auto& joinType = join.JoinKind().Ref();
  2184. if (const auto lEmpty = join.LeftInput().Ref().GetConstraint<TEmptyConstraintNode>(), rEmpty = join.RightInput().Ref().GetConstraint<TEmptyConstraintNode>(); lEmpty && rEmpty) {
  2185. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  2186. } else if (lEmpty && joinType.Content().starts_with("Left")) {
  2187. input->AddConstraint(lEmpty);
  2188. } else if (rEmpty && joinType.Content().starts_with("Right")) {
  2189. input->AddConstraint(rEmpty);
  2190. } else if ((lEmpty || rEmpty) && (joinType.IsAtom("Inner") || joinType.Content().ends_with("Semi"))) {
  2191. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  2192. }
  2193. bool lOneRow = false, rOneRow = false;
  2194. if (const auto& flags = join.Flags()) {
  2195. flags.Cast().Ref().ForEachChild([&](const TExprNode& flag) {
  2196. lOneRow = lOneRow || flag.IsAtom("LeftUnique");
  2197. rOneRow = rOneRow || flag.IsAtom("RightUnique");
  2198. });
  2199. }
  2200. const auto lUnique = GetForPayload<false>(join.LeftInput().Ref(), ctx);
  2201. const auto rUnique = GetForPayload<false>(join.RightInput().Ref(), ctx);
  2202. const auto lDistinct = GetForPayload<true>(join.LeftInput().Ref(), ctx);
  2203. const auto rDistinct = GetForPayload<true>(join.RightInput().Ref(), ctx);
  2204. const bool leftSide = joinType.Content().starts_with("Left");
  2205. const bool rightSide = joinType.Content().starts_with("Right");
  2206. if (joinType.Content().ends_with("Semi") || joinType.Content().ends_with("Only")) {
  2207. if (leftSide) {
  2208. if (lUnique)
  2209. input->AddConstraint(lUnique);
  2210. if (lDistinct)
  2211. input->AddConstraint(lDistinct);
  2212. } else if (rightSide) {
  2213. if (rUnique)
  2214. input->AddConstraint(rUnique);
  2215. if (rDistinct)
  2216. input->AddConstraint(rDistinct);
  2217. }
  2218. } else if (lOneRow || rOneRow) {
  2219. const auto rename = [](const std::string_view& prefix, TPartOfConstraintBase::TPathType path) {
  2220. path.emplace_front(prefix);
  2221. return std::vector<TPartOfConstraintBase::TPathType>(1U, std::move(path));
  2222. };
  2223. const auto leftRename = std::bind(rename, ctx.GetIndexAsString(0U), std::placeholders::_1);
  2224. const auto rightRename = std::bind(rename, ctx.GetIndexAsString(1U), std::placeholders::_1);
  2225. if (lUnique || rUnique) {
  2226. const TUniqueConstraintNode* unique = nullptr;
  2227. const bool exclusion = joinType.IsAtom("Exclusion");
  2228. const bool useLeft = lUnique && (rOneRow || exclusion);
  2229. const bool useRight = rUnique && (lOneRow || exclusion);
  2230. if (useLeft && !useRight)
  2231. unique = lUnique->RenameFields(ctx, leftRename);
  2232. else if (useRight && !useLeft)
  2233. unique = rUnique->RenameFields(ctx, rightRename);
  2234. else if (useLeft && useRight)
  2235. unique = TUniqueConstraintNode::Merge(lUnique->RenameFields(ctx, leftRename), rUnique->RenameFields(ctx, rightRename), ctx);
  2236. if (unique)
  2237. input->AddConstraint(unique);
  2238. }
  2239. if (lDistinct || rDistinct) {
  2240. const TDistinctConstraintNode* distinct = nullptr;
  2241. const bool inner = joinType.IsAtom("Inner");
  2242. const bool useLeft = lDistinct && rOneRow && (inner || leftSide);
  2243. const bool useRight = rDistinct && lOneRow && (inner || rightSide);
  2244. if (useLeft && !useRight)
  2245. distinct = lDistinct->RenameFields(ctx, leftRename);
  2246. else if (useRight && !useLeft)
  2247. distinct = rDistinct->RenameFields(ctx, rightRename);
  2248. else if (useLeft && useRight)
  2249. distinct = TDistinctConstraintNode::Merge(lDistinct->RenameFields(ctx, leftRename), rDistinct->RenameFields(ctx, rightRename), ctx);
  2250. if (distinct)
  2251. input->AddConstraint(distinct);
  2252. }
  2253. }
  2254. return TStatus::Ok;
  2255. }
  2256. TStatus IsKeySwitchWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
  2257. const TCoIsKeySwitch keySwitch(input);
  2258. TSmallVec<TConstraintNode::TListType> itemConstraints, stateConstraints;
  2259. itemConstraints.emplace_back(keySwitch.Item().Ref().GetAllConstraints());
  2260. stateConstraints.emplace_back(keySwitch.State().Ref().GetAllConstraints());
  2261. return UpdateLambdaConstraints(input->ChildRef(TCoIsKeySwitch::idx_ItemKeyExtractor), ctx, itemConstraints)
  2262. .Combine(UpdateLambdaConstraints(input->ChildRef(TCoIsKeySwitch::idx_StateKeyExtractor), ctx, stateConstraints));
  2263. }
  2264. template<bool Wide>
  2265. static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& node, const TExprNode::TChildrenType& args, TExprContext& ctx) {
  2266. TPartOfConstraintBase::TSetType keys;
  2267. if (node.IsCallable("AggrNotEquals")) {
  2268. const TExprNode& body = node.Head().IsCallable("StablePickle") ? node.Head() : node;
  2269. if (body.Head().IsList() && body.Tail().IsList() && body.Head().ChildrenSize() == body.Tail().ChildrenSize()) {
  2270. keys.reserve(body.Tail().ChildrenSize());
  2271. for (auto i = 0U; i < body.Head().ChildrenSize(); ++i){
  2272. if (auto l = GetPathToKey(*body.Head().Child(i), args), r = GetPathToKey(*body.Tail().Child(i), args); l && r && *l == *r) {
  2273. if constexpr (Wide) {
  2274. auto path = r->first;
  2275. path.emplace_front(ctx.GetIndexAsString(r->second));
  2276. } else {
  2277. YQL_ENSURE(l->second == 0U, "Unexpected arg index: " << l->second);
  2278. keys.insert_unique(l->first);
  2279. }
  2280. }
  2281. }
  2282. } else if (auto l = GetPathToKey(body.Head(), args), r = GetPathToKey(body.Tail(), args); l && r && *l == *r) {
  2283. if constexpr (Wide) {
  2284. auto path = l->first;
  2285. path.emplace_front(ctx.GetIndexAsString(l->second));
  2286. } else {
  2287. YQL_ENSURE(r->second == 0U, "Unexpected arg index: " << r->second);
  2288. keys.insert_unique(r->first);
  2289. }
  2290. }
  2291. } else if (node.IsCallable("Or")) {
  2292. keys.reserve(node.ChildrenSize());
  2293. for (auto i = 0U; i < node.ChildrenSize(); ++i) {
  2294. const auto& part = GetSimpleKeys<Wide>(*node.Child(i), args, ctx);
  2295. keys.insert_unique(part.cbegin(), part.cend());
  2296. }
  2297. }
  2298. return keys;
  2299. }
  2300. template<bool Wide>
  2301. static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& selector, TExprContext& ctx) {
  2302. YQL_ENSURE(selector.IsLambda() && 2U == selector.ChildrenSize());
  2303. const auto& body = selector.Tail();
  2304. if constexpr (!Wide) {
  2305. if (TCoIsKeySwitch::Match(&body)) {
  2306. const TCoIsKeySwitch keySwitch(&body);
  2307. const auto& i = GetPathsToKeys(*ctx.ReplaceNode(keySwitch.ItemKeyExtractor().Body().Ptr(), keySwitch.ItemKeyExtractor().Args().Arg(0).Ref(), keySwitch.Item().Ptr()), keySwitch.Item().Ref());
  2308. const auto& s = GetPathsToKeys(*ctx.ReplaceNode(keySwitch.StateKeyExtractor().Body().Ptr(), keySwitch.StateKeyExtractor().Args().Arg(0).Ref(), keySwitch.State().Ptr()), keySwitch.Item().Ref());
  2309. return i == s ? i : TPartOfConstraintBase::TSetType();
  2310. }
  2311. }
  2312. return GetSimpleKeys<Wide>(selector.Tail(), selector.Head().Children(), ctx);
  2313. }
  2314. static TExprNode::TPtr FuseInitLambda(const TExprNode& inner, const TExprNode& outer, TExprContext& ctx) {
  2315. YQL_ENSURE(outer.IsLambda() && inner.IsLambda());
  2316. const auto& outerArgs = outer.Head();
  2317. const auto& innerArgs = inner.Head();
  2318. YQL_ENSURE(outerArgs.ChildrenSize() + 1U == innerArgs.ChildrenSize() + inner.ChildrenSize());
  2319. TNodeOnNodeOwnedMap outerReplaces(outerArgs.ChildrenSize());
  2320. auto i = 0U;
  2321. for (auto& item : innerArgs.ChildrenList())
  2322. YQL_ENSURE(outerReplaces.emplace(outerArgs.Child(i++), std::move(item)).second);
  2323. for (auto& item : GetLambdaBody(inner))
  2324. YQL_ENSURE(outerReplaces.emplace(outerArgs.Child(i++), std::move(item)).second);
  2325. return ctx.NewLambda(outer.Pos(), inner.HeadPtr(), ctx.ReplaceNodes(GetLambdaBody(outer), outerReplaces));
  2326. }
  2327. TStatus CondenseWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2328. auto argsConstraints = GetConstraintsForInputArgument<true, false, false>(*input, ctx);
  2329. const auto initState = input->Child(1);
  2330. argsConstraints.emplace_back(initState->GetAllConstraints());
  2331. if (const auto status = UpdateLambdaConstraints(input->ChildRef(2), ctx, argsConstraints)
  2332. .Combine(UpdateLambdaConstraints(input->TailRef(), ctx, argsConstraints)); status != TStatus::Ok) {
  2333. return status;
  2334. }
  2335. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2336. }
  2337. template<class TConstraint, bool Wide>
  2338. static void GetCommonFromBothLambdas(const TExprNode::TPtr& input, const typename TConstraint::TMainConstraint* original, TExprContext& ctx) {
  2339. if (original)
  2340. if (const auto initPart = GetConstraintFromLambda<TConstraint, Wide>(*input->Child(1), ctx))
  2341. if (const auto init = TConstraint::MakeComplete(ctx, initPart->GetColumnMapping(), original))
  2342. if (const auto updatePart = GetConstraintFromLambda<TConstraint, Wide>(input->Tail(), ctx))
  2343. if (const auto update = TConstraint::MakeComplete(ctx, updatePart->GetColumnMapping(), original))
  2344. if (const auto common = init->MakeCommon(update, ctx))
  2345. input->AddConstraint(common->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  2346. }
  2347. template<bool Wide>
  2348. TStatus Condense1Wrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2349. auto argsConstraints = GetConstraintsForInputArgument<true, Wide, false>(*input, ctx);
  2350. const auto initLambda = input->Child(1);
  2351. const auto switchLambda = input->Child(2);
  2352. const TUniqueConstraintNode* unique = nullptr;
  2353. const TDistinctConstraintNode* distinct = nullptr;
  2354. const auto sorted = input->Head().GetConstraint<TSortedConstraintNode>();
  2355. const auto chopped = input->Head().GetConstraint<TChoppedConstraintNode>();
  2356. if (sorted || chopped) {
  2357. if (const auto& keys = GetSimpleKeys<Wide>(*FuseInitLambda(*initLambda, *switchLambda, ctx), ctx); !keys.empty()) {
  2358. if (sorted && (sorted->StartsWith(keys) || sorted->GetSimplifiedForType(*input->Head().GetTypeAnn(), ctx)->StartsWith(keys)) ||
  2359. chopped && (chopped->Equals(keys) || chopped->GetSimplifiedForType(*input->Head().GetTypeAnn(), ctx)->Equals(keys))) {
  2360. TPartOfConstraintBase::TSetOfSetsType sets;
  2361. sets.reserve(keys.size());
  2362. for (const auto& key : keys)
  2363. sets.insert_unique(TPartOfConstraintBase::TSetType{key});
  2364. unique = ctx.MakeConstraint<TUniqueConstraintNode>(TUniqueConstraintNode::TContentType{sets})->GetComplicatedForType(*input->Head().GetTypeAnn(), ctx);
  2365. distinct = ctx.MakeConstraint<TDistinctConstraintNode>(TDistinctConstraintNode::TContentType{sets})->GetComplicatedForType(*input->Head().GetTypeAnn(), ctx);
  2366. if constexpr (Wide) {
  2367. if (const auto& mapping = TPartOfUniqueConstraintNode::GetCommonMapping(unique); !mapping.empty()) {
  2368. for (ui32 i = 0U; i < argsConstraints.size(); ++i) {
  2369. if (auto extracted = TPartOfUniqueConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) {
  2370. argsConstraints[i].emplace_back(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(extracted)));
  2371. }
  2372. }
  2373. }
  2374. if (const auto& mapping = TPartOfDistinctConstraintNode::GetCommonMapping(distinct); !mapping.empty()) {
  2375. for (ui32 i = 0U; i < argsConstraints.size(); ++i) {
  2376. if (auto extracted = TPartOfDistinctConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) {
  2377. argsConstraints[i].emplace_back(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(extracted)));
  2378. }
  2379. }
  2380. }
  2381. } else {
  2382. argsConstraints.front().emplace_back(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(TPartOfUniqueConstraintNode::GetCommonMapping(unique)));
  2383. argsConstraints.front().emplace_back(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(TPartOfDistinctConstraintNode::GetCommonMapping(distinct)));
  2384. }
  2385. }
  2386. }
  2387. }
  2388. if (const auto status = UpdateLambdaConstraints(input->ChildRef(1), ctx, argsConstraints); status != TStatus::Ok) {
  2389. return status;
  2390. }
  2391. argsConstraints.reserve(argsConstraints.size() + initLambda->ChildrenSize() - 1U);
  2392. for (ui32 i = 1U; i < initLambda->ChildrenSize(); ++i) {
  2393. argsConstraints.emplace_back(initLambda->Child(i)->GetAllConstraints());
  2394. }
  2395. if (const auto status = UpdateLambdaConstraints(input->ChildRef(2), ctx, argsConstraints)
  2396. .Combine(UpdateLambdaConstraints(input->TailRef(), ctx, argsConstraints)); status != TStatus::Ok) {
  2397. return status;
  2398. }
  2399. GetCommonFromBothLambdas<TPartOfUniqueConstraintNode, Wide>(input, unique, ctx);
  2400. GetCommonFromBothLambdas<TPartOfDistinctConstraintNode, Wide>(input, distinct, ctx);
  2401. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2402. }
  2403. template<class TConstraint, bool Wide>
  2404. static void GetCommonFromBothLambdas(const TExprNode::TPtr& input, TExprContext& ctx) {
  2405. if (const auto original = input->Head().GetConstraint<typename TConstraint::TMainConstraint>())
  2406. if (const auto initPart = GetConstraintFromLambda<TConstraint, Wide>(*input->Child(1), ctx))
  2407. if (const auto init = TConstraint::MakeComplete(ctx, initPart->GetColumnMapping(), original))
  2408. if (const auto updatePart = GetConstraintFromLambda<TConstraint, Wide>(input->Tail(), ctx))
  2409. if (const auto update = TConstraint::MakeComplete(ctx, updatePart->GetColumnMapping(), original))
  2410. if (const auto common = init->MakeCommon(update, ctx))
  2411. input->AddConstraint(common->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  2412. }
  2413. template<bool Wide>
  2414. TStatus Chain1MapWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2415. auto argsConstraints = GetConstraintsForInputArgument<true, Wide>(*input, ctx);
  2416. if (const auto status = UpdateLambdaConstraints(input->ChildRef(1), ctx, argsConstraints); status != TStatus::Ok) {
  2417. return status;
  2418. }
  2419. const auto initLambda = input->Child(1);
  2420. argsConstraints.reserve(argsConstraints.size() + initLambda->ChildrenSize() - 1U);
  2421. for (ui32 i = 1U; i < initLambda->ChildrenSize(); ++i) {
  2422. argsConstraints.emplace_back(initLambda->Child(i)->GetAllConstraints());
  2423. }
  2424. if (const auto status = UpdateLambdaConstraints(input->ChildRef(2), ctx, argsConstraints); status != TStatus::Ok) {
  2425. return status;
  2426. }
  2427. GetCommonFromBothLambdas<TPartOfSortedConstraintNode, Wide>(input, ctx);
  2428. GetCommonFromBothLambdas<TPartOfChoppedConstraintNode, Wide>(input, ctx);
  2429. GetCommonFromBothLambdas<TPartOfUniqueConstraintNode, Wide>(input, ctx);
  2430. GetCommonFromBothLambdas<TPartOfDistinctConstraintNode, Wide>(input, ctx);
  2431. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2432. }
  2433. template<bool Distinct>
  2434. static void GetUniquesForPayloads(const TExprNode::TPtr& input, TExprContext& ctx) {
  2435. typename TUniqueConstraintNodeBase<Distinct>::TContentType content{TPartOfConstraintBase::TSetOfSetsType{TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType{ctx.GetIndexAsString(0U)}}}};
  2436. if (const auto lambda = GetConstraintFromLambda<TPartOfConstraintNode<TUniqueConstraintNodeBase<Distinct>>, false>(*input->Child(2), ctx)) {
  2437. if (const auto original = GetDetailed(input->Head().GetConstraint<TUniqueConstraintNodeBase<Distinct>>(), *input->Head().GetTypeAnn(), ctx)) {
  2438. if (const auto complete = TPartOfConstraintNode<TUniqueConstraintNodeBase<Distinct>>::MakeComplete(ctx, lambda->GetColumnMapping(), original, ctx.GetIndexAsString(1U))) {
  2439. content.insert_unique(complete->GetContent().cbegin(), complete->GetContent().cend());
  2440. }
  2441. }
  2442. }
  2443. input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNodeBase<Distinct>>(std::move(content)));
  2444. }
  2445. TStatus ToDictWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2446. const auto argsConstraints = GetConstraintsForInputArgument<false, false>(*input, ctx);
  2447. if (const auto status = UpdateLambdaConstraints(input->ChildRef(1), ctx, argsConstraints); status != TStatus::Ok) {
  2448. return status;
  2449. }
  2450. if (const auto status = UpdateLambdaConstraints(input->ChildRef(2), ctx, argsConstraints); status != TStatus::Ok) {
  2451. return status;
  2452. }
  2453. GetUniquesForPayloads<true>(input, ctx);
  2454. GetUniquesForPayloads<false>(input, ctx);
  2455. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2456. }
  2457. TStatus DictItemsWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2458. if (const auto unique = input->Head().GetConstraint<TUniqueConstraintNode>())
  2459. input->AddConstraint(unique->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  2460. if (const auto distinct = input->Head().GetConstraint<TDistinctConstraintNode>())
  2461. input->AddConstraint(distinct->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
  2462. return FromFirst<TPartOfUniqueConstraintNode, TPartOfDistinctConstraintNode, TEmptyConstraintNode>(input, output, ctx);
  2463. }
  2464. template<bool Keys>
  2465. TStatus DictHalfWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2466. const auto& side = ctx.GetIndexAsString(Keys ? 0U : 1U);
  2467. const auto reduce = [&side](const TPartOfConstraintBase::TPathType& path) -> std::vector<TPartOfConstraintBase::TPathType> {
  2468. if (path.empty() || path.front() != side)
  2469. return {};
  2470. auto copy = path;
  2471. copy.pop_front();
  2472. return {copy};
  2473. };
  2474. ReduceFromHead<TUniqueConstraintNode, true>(input, reduce, ctx);
  2475. ReduceFromHead<TDistinctConstraintNode, true>(input, reduce, ctx);
  2476. ReduceFromHead<TPartOfUniqueConstraintNode>(input, reduce, ctx);
  2477. ReduceFromHead<TPartOfDistinctConstraintNode>(input, reduce, ctx);
  2478. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2479. }
  2480. template<bool Partitions>
  2481. TStatus ShuffleByKeysWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2482. using TCoBase = std::conditional_t<Partitions, TCoPartitionByKeyBase, TCoShuffleByKeys>;
  2483. if (const auto status = UpdateLambdaConstraints(*input->Child(TCoBase::idx_KeySelectorLambda)); status != TStatus::Ok) {
  2484. return status;
  2485. }
  2486. TPartOfConstraintBase::TSetType keys;
  2487. if constexpr (Partitions) {
  2488. keys = GetPathsToKeys(input->Child(TCoBase::idx_KeySelectorLambda)->Tail(), input->Child(TCoBase::idx_KeySelectorLambda)->Head().Head());
  2489. if (const auto sortKeySelector = input->Child(TCoBase::idx_SortKeySelectorLambda); sortKeySelector->IsLambda()) {
  2490. if (const auto status = UpdateLambdaConstraints(*sortKeySelector); status != TStatus::Ok) {
  2491. return status;
  2492. }
  2493. auto sortKeys = GetPathsToKeys(sortKeySelector->Tail(), sortKeySelector->Head().Head());
  2494. std::move(sortKeys.begin(), sortKeys.end(), std::back_inserter(keys));
  2495. std::sort(keys.begin(), keys.end());
  2496. }
  2497. }
  2498. const auto filter = [](const std::string_view& name) {
  2499. return name == TEmptyConstraintNode::Name() || name == TUniqueConstraintNode::Name() || name == TDistinctConstraintNode::Name();
  2500. };
  2501. TConstraintNode::TListType argConstraints;
  2502. const auto source = input->Child(TCoBase::idx_Input);
  2503. std::copy_if(source->GetAllConstraints().cbegin(), source->GetAllConstraints().cend(), std::back_inserter(argConstraints), std::bind(filter, std::bind(&TConstraintNode::GetName, std::placeholders::_1)));
  2504. if (const auto multi = source->template GetConstraint<TMultiConstraintNode>())
  2505. if (const auto filtered = multi->FilterConstraints(ctx, filter))
  2506. argConstraints.emplace_back(filtered);
  2507. if constexpr (Partitions) {
  2508. if (!keys.empty())
  2509. argConstraints.emplace_back(ctx.MakeConstraint<TChoppedConstraintNode>(keys)->GetComplicatedForType(*input->Head().GetTypeAnn(), ctx));
  2510. }
  2511. if (const auto status = UpdateLambdaConstraints(input->ChildRef(TCoBase::idx_ListHandlerLambda), ctx, {argConstraints}); status != TStatus::Ok) {
  2512. return status;
  2513. }
  2514. const auto handlerLambda = input->Child(TCoBase::idx_ListHandlerLambda);
  2515. if (const auto unique = handlerLambda->template GetConstraint<TUniqueConstraintNode>())
  2516. input->AddConstraint(unique);
  2517. if (const auto distinct = handlerLambda->template GetConstraint<TDistinctConstraintNode>())
  2518. input->AddConstraint(distinct);
  2519. const bool multiInput = ETypeAnnotationKind::Variant == GetSeqItemType(*input->Head().GetTypeAnn()).GetKind();
  2520. const auto lambdaVarIndex = handlerLambda->template GetConstraint<TVarIndexConstraintNode>();
  2521. const auto multi = input->Head().template GetConstraint<TMultiConstraintNode>();
  2522. const auto lambdaMulti = handlerLambda->template GetConstraint<TMultiConstraintNode>();
  2523. if (const auto varIndex = input->Head().template GetConstraint<TVarIndexConstraintNode>()) {
  2524. if (multiInput) {
  2525. if (lambdaVarIndex) {
  2526. if (auto outVarIndex = GetVarIndexOverVarIndexConstraint(*varIndex, *lambdaVarIndex, ctx)) {
  2527. input->AddConstraint(outVarIndex);
  2528. }
  2529. }
  2530. } else {
  2531. if (lambdaMulti) {
  2532. TVarIndexConstraintNode::TMapType remapItems;
  2533. for (auto& multiItem: lambdaMulti->GetItems()) {
  2534. for (auto& varItem: varIndex->GetIndexMapping()) {
  2535. remapItems.push_back(std::make_pair(multiItem.first, varItem.second));
  2536. }
  2537. }
  2538. if (!remapItems.empty()) {
  2539. ::SortUnique(remapItems);
  2540. input->AddConstraint(ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(remapItems)));
  2541. }
  2542. } else {
  2543. input->AddConstraint(varIndex);
  2544. }
  2545. }
  2546. }
  2547. if (lambdaMulti && !input->Head().GetConstraint<TEmptyConstraintNode>()) {
  2548. TMultiConstraintNode::TMapType remappedItems;
  2549. for (const auto& item: lambdaMulti->GetItems()) {
  2550. remappedItems.push_back(std::make_pair(item.first, TConstraintSet{}));
  2551. if (!multiInput) { // remapping one to many
  2552. if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>())
  2553. remappedItems.pop_back();
  2554. else {
  2555. if (const auto unique = item.second.template GetConstraint<TUniqueConstraintNode>())
  2556. remappedItems.back().second.AddConstraint(unique);
  2557. if (const auto distinct = item.second.template GetConstraint<TDistinctConstraintNode>())
  2558. remappedItems.back().second.AddConstraint(distinct);
  2559. }
  2560. }
  2561. else if (lambdaVarIndex && multi) {
  2562. const auto range = lambdaVarIndex->GetIndexMapping().equal_range(item.first);
  2563. switch (std::distance(range.first, range.second)) {
  2564. case 0: // new index
  2565. break;
  2566. case 1: // remapping 1 to 1
  2567. if (auto origConstr = multi->GetItem(range.first->second)) {
  2568. if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>())
  2569. remappedItems.pop_back();
  2570. else {
  2571. if (const auto unique = item.second.template GetConstraint<TUniqueConstraintNode>())
  2572. remappedItems.back().second.AddConstraint(unique);
  2573. if (const auto distinct = item.second.template GetConstraint<TDistinctConstraintNode>())
  2574. remappedItems.back().second.AddConstraint(distinct);
  2575. }
  2576. } else {
  2577. remappedItems.pop_back();
  2578. }
  2579. break;
  2580. default: // remapping many to one
  2581. {
  2582. std::vector<const TConstraintSet*> nonEmpty;
  2583. for (auto i = range.first; i != range.second; ++i) {
  2584. if (auto origConstr = multi->GetItem(i->second)) {
  2585. nonEmpty.push_back(origConstr);
  2586. }
  2587. }
  2588. EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint<TEmptyConstraintNode>(); });
  2589. if (nonEmpty.empty()) {
  2590. remappedItems.back().second.AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  2591. } else if (nonEmpty.size() == 1) {
  2592. remappedItems.back().second = std::move(*nonEmpty.front());
  2593. }
  2594. }
  2595. }
  2596. } else {
  2597. remappedItems.back().second = item.second;
  2598. }
  2599. }
  2600. if (remappedItems) {
  2601. input->AddConstraint(ctx.MakeConstraint<TMultiConstraintNode>(std::move(remappedItems)));
  2602. }
  2603. }
  2604. else if (multi && lambdaVarIndex) { // Many to one
  2605. const auto range = lambdaVarIndex->GetIndexMapping().equal_range(0);
  2606. std::vector<const TConstraintSet*> nonEmpty;
  2607. for (auto i = range.first; i != range.second; ++i) {
  2608. if (auto origConstr = multi->GetItem(i->second)) {
  2609. nonEmpty.push_back(origConstr);
  2610. }
  2611. }
  2612. EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint<TEmptyConstraintNode>(); });
  2613. if (nonEmpty.empty()) {
  2614. input->AddConstraint(ctx.MakeConstraint<TEmptyConstraintNode>());
  2615. } else if (nonEmpty.size() == 1) {
  2616. input->SetConstraints(*nonEmpty.front());
  2617. }
  2618. }
  2619. TApplyConstraintFromInput<TCoBase::idx_ListHandlerLambda, TEmptyConstraintNode>::Do(input);
  2620. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2621. }
  2622. template <bool Final>
  2623. TStatus AggregateWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
  2624. if (HasSetting(input->Tail(), "session")) {
  2625. // TODO: support sessions
  2626. return TStatus::Ok;
  2627. }
  2628. if (const auto size = input->Child(1)->ChildrenSize()) {
  2629. if constexpr (Final) {
  2630. bool allKeysInOutput = true;
  2631. if (auto outputColumnsSetting = GetSetting(input->Tail(), "output_columns")) {
  2632. THashSet<TStringBuf> outputColumns;
  2633. for (auto& col : outputColumnsSetting->Child(1)->Children()) {
  2634. YQL_ENSURE(col->IsAtom());
  2635. outputColumns.insert(col->Content());
  2636. }
  2637. allKeysInOutput = AllOf(input->Child(1)->Children(), [&](const auto& key) { return outputColumns.contains(key->Content()); });
  2638. }
  2639. if (allKeysInOutput) {
  2640. std::vector<std::string_view> columns;
  2641. columns.reserve(size);
  2642. for (const auto& child: input->Child(1)->Children()) {
  2643. columns.emplace_back(child->Content());
  2644. }
  2645. input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(columns));
  2646. input->AddConstraint(ctx.MakeConstraint<TDistinctConstraintNode>(columns));
  2647. }
  2648. }
  2649. return FromFirst<TEmptyConstraintNode>(input, output, ctx);
  2650. }
  2651. return TStatus::Ok;
  2652. }
  2653. TStatus FoldWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
  2654. const TStructExprType* inItemType = GetNonEmptyStructItemType(*input->Head().GetTypeAnn());
  2655. const TStructExprType* outItemType = GetNonEmptyStructItemType(*input->GetTypeAnn());
  2656. if (!inItemType || !outItemType) {
  2657. return UpdateAllChildLambdasConstraints(*input);
  2658. }
  2659. if (input->Child(1)->IsLambda()) {
  2660. TConstraintNode::TListType argConstraints;
  2661. if (const auto status = UpdateLambdaConstraints(input->ChildRef(1), ctx, {argConstraints}); status != TStatus::Ok) {
  2662. return status;
  2663. }
  2664. }
  2665. const auto initState = input->Child(1);
  2666. auto stateConstraints = initState->GetAllConstraints();
  2667. stateConstraints.erase(
  2668. std::remove_if(
  2669. stateConstraints.begin(),
  2670. stateConstraints.end(),
  2671. [](const TConstraintNode* c) { return c->GetName() == TEmptyConstraintNode::Name(); }
  2672. ),
  2673. stateConstraints.end()
  2674. );
  2675. TConstraintNode::TListType argConstraints;
  2676. if (const auto status = UpdateLambdaConstraints(input->TailRef(), ctx, {argConstraints, stateConstraints}); status != TStatus::Ok) {
  2677. return status;
  2678. }
  2679. return TStatus::Ok;
  2680. }
  2681. TStatus MultiHoppingCoreWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
  2682. if (const auto status = UpdateAllChildLambdasConstraints(*input); status != TStatus::Ok) {
  2683. return status;
  2684. }
  2685. TExprNode::TPtr keySelectorLambda = input->Child(TCoMultiHoppingCore::idx_KeyExtractor);
  2686. const auto keys = GetPathsToKeys(keySelectorLambda->Tail(), keySelectorLambda->Head().Head());
  2687. std::vector<std::string_view> columns(keys.size());
  2688. std::transform(keys.begin(), keys.end(), columns.begin(), [](const TPartOfConstraintBase::TPathType& path) -> std::string_view {
  2689. return path.front();
  2690. });
  2691. if (!columns.empty()) {
  2692. input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(columns));
  2693. input->AddConstraint(ctx.MakeConstraint<TDistinctConstraintNode>(columns));
  2694. }
  2695. return TStatus::Ok;
  2696. }
  2697. private:
  2698. template <class TConstraintContainer>
  2699. static void CopyExcept(TConstraintContainer& dst, const TConstraintContainer& from, const TSet<TStringBuf>& except) {
  2700. for (auto c: from.GetAllConstraints()) {
  2701. if (!except.contains(c->GetName())) {
  2702. dst.AddConstraint(c);
  2703. }
  2704. }
  2705. }
  2706. template <class TConstraintContainer>
  2707. static void CopyExcept(TConstraintContainer& dst, const TConstraintContainer& from, TStringBuf except) {
  2708. for (auto c: from.GetAllConstraints()) {
  2709. if (c->GetName() != except) {
  2710. dst.AddConstraint(c);
  2711. }
  2712. }
  2713. }
  2714. static void ExtractKeys(const TExprNode& keySelectorLambda, TVector<TStringBuf>& columns) {
  2715. const auto arg = keySelectorLambda.Head().Child(0);
  2716. auto body = keySelectorLambda.Child(1);
  2717. if (body->IsCallable("StablePickle")) {
  2718. body = body->Child(0);
  2719. }
  2720. ExtractSimpleKeys(body, arg, columns);
  2721. }
  2722. template<class TConstraintWithFields>
  2723. static const TConstraintWithFields* GetDetailed(const TConstraintWithFields* constraint, const TTypeAnnotationNode& type, TExprContext& ctx) {
  2724. return constraint ? constraint->GetComplicatedForType(type, ctx) : nullptr;
  2725. }
  2726. static const TStructExprType* GetNonEmptyStructItemType(const TTypeAnnotationNode& type) {
  2727. const auto itemType = GetSeqItemType(&type);
  2728. if (!itemType || itemType->GetKind() != ETypeAnnotationKind::Struct) {
  2729. return nullptr;
  2730. }
  2731. const TStructExprType* structType = itemType->Cast<TStructExprType>();
  2732. return structType->GetSize() ? structType : nullptr;
  2733. }
  2734. static const TSortedConstraintNode* DeduceSortConstraint(const TExprNode& directions, const TExprNode& keyExtractor, TExprContext& ctx) {
  2735. if (const auto& columns = ExtractSimpleSortTraits(directions, keyExtractor); !columns.empty()) {
  2736. TSortedConstraintNode::TContainerType content(columns.size());
  2737. std::transform(columns.cbegin(), columns.cend(), content.begin(), [](const std::pair<TPartOfConstraintBase::TPathType, bool>& item) {
  2738. return std::make_pair(TSortedConstraintNode::TSetType{item.first}, item.second);
  2739. });
  2740. return ctx.MakeConstraint<TSortedConstraintNode>(std::move(content));
  2741. }
  2742. return nullptr;
  2743. }
  2744. static const TVarIndexConstraintNode* GetVarIndexOverVarIndexConstraint(const TVarIndexConstraintNode& inputVarIndex,
  2745. const TVarIndexConstraintNode& varIndex, TExprContext& ctx)
  2746. {
  2747. TVarIndexConstraintNode::TMapType result;
  2748. for (auto& item: varIndex.GetIndexMapping()) {
  2749. auto range = inputVarIndex.GetIndexMapping().equal_range(item.second);
  2750. for (auto it = range.first; it != range.second; ++it) {
  2751. result.push_back(std::make_pair(item.first, it->second));
  2752. }
  2753. }
  2754. if (!result.empty()) {
  2755. return ctx.MakeConstraint<TVarIndexConstraintNode>(std::move(result));
  2756. }
  2757. return nullptr;
  2758. }
  2759. static const TExprNode* SkipModifiers(const TExprNode* valueNode) {
  2760. if (TCoJust::Match(valueNode)) {
  2761. return SkipModifiers(valueNode->Child(0));
  2762. }
  2763. if (TCoUnwrap::Match(valueNode)) {
  2764. return SkipModifiers(valueNode->Child(0));
  2765. }
  2766. return valueNode;
  2767. }
  2768. private:
  2769. const bool SubGraph;
  2770. std::unordered_map<std::string_view, THandler> Functions;
  2771. };
  2772. template<> const TPartOfSortedConstraintNode*
  2773. TCallableConstraintTransformer::GetConstraintFromWideResultLambda<TPartOfSortedConstraintNode>(const TExprNode& lambda, TExprContext& ctx) {
  2774. TPartOfSortedConstraintNode::TMapType sorted;
  2775. for (auto i = 1U; i < lambda.ChildrenSize(); ++i) {
  2776. if (const auto part = lambda.Child(i)->GetConstraint<TPartOfSortedConstraintNode>())
  2777. TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(ctx.GetIndexAsString(i - 1U)));
  2778. }
  2779. return sorted.empty() ? nullptr : ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(sorted));
  2780. }
  2781. template<> const TPartOfChoppedConstraintNode*
  2782. TCallableConstraintTransformer::GetConstraintFromWideResultLambda<TPartOfChoppedConstraintNode>(const TExprNode& lambda, TExprContext& ctx) {
  2783. TPartOfChoppedConstraintNode::TMapType chopped;
  2784. for (auto i = 1U; i < lambda.ChildrenSize(); ++i) {
  2785. if (const auto part = lambda.Child(i)->GetConstraint<TPartOfChoppedConstraintNode>())
  2786. TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(ctx.GetIndexAsString(i - 1U)));
  2787. }
  2788. return chopped.empty() ? nullptr : ctx.MakeConstraint<TPartOfChoppedConstraintNode>(std::move(chopped));
  2789. }
  2790. template<> const TPartOfUniqueConstraintNode*
  2791. TCallableConstraintTransformer::GetConstraintFromWideResultLambda<TPartOfUniqueConstraintNode>(const TExprNode& lambda, TExprContext& ctx) {
  2792. TPartOfUniqueConstraintNode::TMapType uniques;
  2793. for (auto i = 1U; i < lambda.ChildrenSize(); ++i) {
  2794. if (const auto part = lambda.Child(i)->GetConstraint<TPartOfUniqueConstraintNode>())
  2795. TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(ctx.GetIndexAsString(i - 1U)));
  2796. }
  2797. return uniques.empty() ? nullptr : ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(uniques));
  2798. }
  2799. template<> const TPartOfDistinctConstraintNode*
  2800. TCallableConstraintTransformer::GetConstraintFromWideResultLambda<TPartOfDistinctConstraintNode>(const TExprNode& lambda, TExprContext& ctx) {
  2801. TPartOfDistinctConstraintNode::TMapType uniques;
  2802. for (auto i = 1U; i < lambda.ChildrenSize(); ++i) {
  2803. if (const auto part = lambda.Child(i)->GetConstraint<TPartOfDistinctConstraintNode>())
  2804. TPartOfDistinctConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(ctx.GetIndexAsString(i - 1U)));
  2805. }
  2806. return uniques.empty() ? nullptr : ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(uniques));
  2807. }
  2808. template<> const TVarIndexConstraintNode*
  2809. TCallableConstraintTransformer::TCallableConstraintTransformer::GetConstraintFromWideResultLambda<TVarIndexConstraintNode>(const TExprNode& lambda, TExprContext& ctx) {
  2810. TVector<const TConstraintSet*> structConstraints;
  2811. structConstraints.reserve(lambda.ChildrenSize() - 1U);
  2812. for (auto i = 1U; i < lambda.ChildrenSize(); ++i) {
  2813. auto valueNode = lambda.Child(i);
  2814. if (TCoCoalesce::Match(valueNode)) {
  2815. if (valueNode->Head().GetTypeAnn()->GetKind() != ETypeAnnotationKind::Optional || valueNode->ChildrenSize() == 1) {
  2816. valueNode = valueNode->Child(0);
  2817. }
  2818. }
  2819. if (TCoJust::Match(valueNode)) {
  2820. valueNode = valueNode->Child(0);
  2821. }
  2822. if (TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) {
  2823. structConstraints.push_back(&valueNode->Head().GetConstraintSet());
  2824. } else if (valueNode->Type() == TExprNode::Argument) {
  2825. structConstraints.push_back(&valueNode->GetConstraintSet());
  2826. }
  2827. }
  2828. return TVarIndexConstraintNode::MakeCommon(structConstraints, ctx);
  2829. }
  2830. template<class TConstraint> const TConstraint*
  2831. TCallableConstraintTransformer::GetConstraintFromWideResultLambda(const TExprNode&, TExprContext&) { return nullptr; }
  2832. class TDefaultCallableConstraintTransformer : public TSyncTransformerBase {
  2833. public:
  2834. TStatus DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) override {
  2835. Y_UNUSED(output);
  2836. Y_UNUSED(ctx);
  2837. return UpdateAllChildLambdasConstraints(*input);
  2838. }
  2839. void Rewind() final {
  2840. }
  2841. };
  2842. template<bool DisableCheck>
  2843. class TConstraintTransformer : public TGraphTransformerBase {
  2844. public:
  2845. TConstraintTransformer(TAutoPtr<IGraphTransformer> callableTransformer, TTypeAnnotationContext& types)
  2846. : CallableTransformer(callableTransformer)
  2847. , Types(types)
  2848. {
  2849. }
  2850. ~TConstraintTransformer() = default;
  2851. TStatus DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final {
  2852. YQL_PROFILE_SCOPE(DEBUG, "ConstraintTransformer::DoTransform");
  2853. output = input;
  2854. auto status = TransformNode(input, output, ctx);
  2855. UpdateStatusIfChanged(status, input, output);
  2856. if (status.Level != TStatus::Error && HasRenames) {
  2857. output = ctx.ReplaceNodes(std::move(output), Processed);
  2858. }
  2859. Processed.clear();
  2860. if (status == TStatus::Ok) {
  2861. Types.ExpectedConstraints.clear();
  2862. }
  2863. HasRenames = false;
  2864. return status;
  2865. }
  2866. NThreading::TFuture<void> DoGetAsyncFuture(const TExprNode& input) final {
  2867. YQL_PROFILE_SCOPE(DEBUG, "ConstraintTransformer::DoGetAsyncFuture");
  2868. Y_UNUSED(input);
  2869. TVector<NThreading::TFuture<void>> futures;
  2870. for (const auto& callable : CallableInputs) {
  2871. futures.push_back(CallableTransformer->GetAsyncFuture(*callable));
  2872. }
  2873. return WaitExceptionOrAll(futures);
  2874. }
  2875. TStatus DoApplyAsyncChanges(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final {
  2876. YQL_PROFILE_SCOPE(DEBUG, "ConstraintTransformer::DoApplyAsyncChanges");
  2877. output = input;
  2878. TStatus combinedStatus = TStatus::Ok;
  2879. for (const auto& callable : CallableInputs) {
  2880. callable->SetState(TExprNode::EState::ConstrPending);
  2881. TExprNode::TPtr callableOutput;
  2882. auto status = CallableTransformer->ApplyAsyncChanges(callable, callableOutput, ctx);
  2883. Y_ABORT_UNLESS(callableOutput);
  2884. YQL_ENSURE(status != TStatus::Async);
  2885. YQL_ENSURE(callableOutput == callable);
  2886. combinedStatus = combinedStatus.Combine(status);
  2887. if (status.Level == TStatus::Error) {
  2888. callable->SetState(TExprNode::EState::Error);
  2889. }
  2890. }
  2891. CallableInputs.clear();
  2892. if (combinedStatus.Level == TStatus::Ok) {
  2893. Processed.clear();
  2894. }
  2895. return combinedStatus;
  2896. }
  2897. void Rewind() final {
  2898. CallableTransformer->Rewind();
  2899. CallableInputs.clear();
  2900. Processed.clear();
  2901. HasRenames = false;
  2902. CurrentFunctions = {};
  2903. }
  2904. private:
  2905. TStatus TransformNode(const TExprNode::TPtr& start, TExprNode::TPtr& output, TExprContext& ctx) {
  2906. output = start;
  2907. auto processedPair = Processed.emplace(start.Get(), nullptr); // by default node is not changed
  2908. if (!processedPair.second) {
  2909. if (processedPair.first->second) {
  2910. output = processedPair.first->second;
  2911. return TStatus::Repeat;
  2912. }
  2913. switch (start->GetState()) {
  2914. case TExprNode::EState::Initial:
  2915. case TExprNode::EState::TypeInProgress:
  2916. case TExprNode::EState::TypePending:
  2917. return TStatus(TStatus::Repeat, true);
  2918. case TExprNode::EState::TypeComplete:
  2919. break;
  2920. case TExprNode::EState::ConstrInProgress:
  2921. return IGraphTransformer::TStatus::Async;
  2922. case TExprNode::EState::ConstrPending:
  2923. if (start->Type() == TExprNode::Lambda) {
  2924. if (start->Head().GetState() != TExprNode::EState::ConstrComplete) {
  2925. return TStatus::Ok;
  2926. } else if (start->Head().ChildrenSize() == 0) {
  2927. break;
  2928. }
  2929. }
  2930. if (start->Type() == TExprNode::Arguments || start->Type() == TExprNode::Argument) {
  2931. break;
  2932. }
  2933. return TStatus(TStatus::Repeat, true);
  2934. case TExprNode::EState::ConstrComplete:
  2935. case TExprNode::EState::ExecutionInProgress:
  2936. case TExprNode::EState::ExecutionRequired:
  2937. case TExprNode::EState::ExecutionPending:
  2938. case TExprNode::EState::ExecutionComplete:
  2939. return TStatus::Ok;
  2940. case TExprNode::EState::Error:
  2941. return TStatus::Error;
  2942. default:
  2943. YQL_ENSURE(false, "Unknown state");
  2944. }
  2945. }
  2946. auto input = start;
  2947. for (;;) {
  2948. TIssueScopeGuard issueScope(ctx.IssueManager, [this, input, &ctx]() -> TIssuePtr {
  2949. TStringBuilder str;
  2950. str << "At ";
  2951. switch (input->Type()) {
  2952. case TExprNode::Callable:
  2953. if (!CurrentFunctions.empty() && CurrentFunctions.top().second) {
  2954. return nullptr;
  2955. }
  2956. if (!CurrentFunctions.empty()) {
  2957. CurrentFunctions.top().second = true;
  2958. }
  2959. str << "function: " << input->Content();
  2960. break;
  2961. case TExprNode::List:
  2962. if (CurrentFunctions.empty()) {
  2963. str << "tuple";
  2964. } else if (!CurrentFunctions.top().second) {
  2965. CurrentFunctions.top().second = true;
  2966. str << "function: " << CurrentFunctions.top().first;
  2967. } else {
  2968. return nullptr;
  2969. }
  2970. break;
  2971. case TExprNode::Lambda:
  2972. if (CurrentFunctions.empty()) {
  2973. str << "lambda";
  2974. } else if (!CurrentFunctions.top().second) {
  2975. CurrentFunctions.top().second = true;
  2976. str << "function: " << CurrentFunctions.top().first;
  2977. } else {
  2978. return nullptr;
  2979. }
  2980. break;
  2981. default:
  2982. str << "unknown";
  2983. }
  2984. return MakeIntrusive<TIssue>(ctx.GetPosition(input->Pos()), str);
  2985. });
  2986. if (input->IsCallable()) {
  2987. CurrentFunctions.emplace(input->Content(), false);
  2988. }
  2989. Y_SCOPE_EXIT(this, input) {
  2990. if (input->IsCallable()) {
  2991. CurrentFunctions.pop();
  2992. if (!CurrentFunctions.empty() && CurrentFunctions.top().first.ends_with('!')) {
  2993. CurrentFunctions.top().second = true;
  2994. }
  2995. }
  2996. };
  2997. TStatus retStatus = TStatus::Error;
  2998. switch (input->GetState()) {
  2999. case TExprNode::EState::Initial:
  3000. case TExprNode::EState::TypeInProgress:
  3001. case TExprNode::EState::TypePending:
  3002. return TStatus(TStatus::Repeat, true);
  3003. case TExprNode::EState::TypeComplete:
  3004. case TExprNode::EState::ConstrPending:
  3005. break;
  3006. case TExprNode::EState::ConstrInProgress:
  3007. return IGraphTransformer::TStatus::Async;
  3008. case TExprNode::EState::ConstrComplete:
  3009. case TExprNode::EState::ExecutionInProgress:
  3010. case TExprNode::EState::ExecutionRequired:
  3011. case TExprNode::EState::ExecutionPending:
  3012. case TExprNode::EState::ExecutionComplete:
  3013. return TStatus::Ok;
  3014. case TExprNode::EState::Error:
  3015. return TStatus::Error;
  3016. default:
  3017. YQL_ENSURE(false, "Unknown state");
  3018. }
  3019. input->SetState(TExprNode::EState::ConstrPending);
  3020. switch (input->Type()) {
  3021. case TExprNode::Atom:
  3022. case TExprNode::World:
  3023. input->SetState(TExprNode::EState::ConstrComplete);
  3024. CheckExpected(*input);
  3025. return TStatus::Ok;
  3026. case TExprNode::List:
  3027. {
  3028. retStatus = TransformChildren(input, output, ctx);
  3029. if (retStatus == TStatus::Ok) {
  3030. retStatus = CallableTransformer->Transform(input, output, ctx);
  3031. if (retStatus == TStatus::Ok) {
  3032. input->SetState(TExprNode::EState::ConstrComplete);
  3033. CheckExpected(*input);
  3034. break;
  3035. }
  3036. }
  3037. if (retStatus != TStatus::Error && input != output) {
  3038. processedPair.first->second = output;
  3039. }
  3040. break;
  3041. }
  3042. case TExprNode::Lambda:
  3043. {
  3044. YQL_ENSURE(input->ChildrenSize() > 0U);
  3045. TExprNode::TPtr out;
  3046. auto argStatus = TransformNode(input->HeadPtr(), out, ctx);
  3047. UpdateStatusIfChanged(argStatus, input->HeadPtr(), out);
  3048. if (argStatus.Level == TStatus::Error) {
  3049. input->SetState(TExprNode::EState::Error);
  3050. return argStatus;
  3051. }
  3052. if (argStatus.Level == TStatus::Repeat)
  3053. return TStatus::Ok;
  3054. TStatus bodyStatus = TStatus::Ok;
  3055. TExprNode::TListType newBody;
  3056. newBody.reserve(input->ChildrenSize() - 1U);
  3057. bool updatedChildren = false;
  3058. for (ui32 i = 1U; i < input->ChildrenSize(); ++i) {
  3059. const auto child = input->ChildPtr(i);
  3060. TExprNode::TPtr newChild;
  3061. auto childStatus = TransformNode(child, newChild, ctx);
  3062. UpdateStatusIfChanged(childStatus, child, newChild);
  3063. updatedChildren = updatedChildren || (newChild != child);
  3064. bodyStatus = bodyStatus.Combine(childStatus);
  3065. newBody.emplace_back(std::move(newChild));
  3066. }
  3067. retStatus = argStatus.Combine(bodyStatus);
  3068. if (retStatus != TStatus::Ok) {
  3069. if (retStatus.Level == TStatus::Error) {
  3070. input->SetState(TExprNode::EState::Error);
  3071. }
  3072. else if (updatedChildren) {
  3073. output = ctx.DeepCopyLambda(*input, std::move(newBody));
  3074. processedPair.first->second = output;
  3075. HasRenames = true;
  3076. }
  3077. } else {
  3078. if (input->ChildrenSize() != 2U)
  3079. input->SetState(TExprNode::EState::ConstrComplete);
  3080. else
  3081. input->CopyConstraints(input->Tail());
  3082. CheckExpected(*input);
  3083. }
  3084. break;
  3085. }
  3086. case TExprNode::Argument:
  3087. if (input->GetState() != TExprNode::EState::ConstrComplete) {
  3088. return TStatus::Repeat;
  3089. }
  3090. return TStatus::Ok;
  3091. case TExprNode::Arguments:
  3092. {
  3093. if (input->Children().empty()) {
  3094. if (TExprNode::EState::ConstrComplete == input->GetState()) {
  3095. return TStatus::Ok;
  3096. }
  3097. return TStatus::Repeat;
  3098. }
  3099. retStatus = TStatus::Ok;
  3100. for (auto& child : input->Children()) {
  3101. TExprNode::TPtr tmp;
  3102. auto childStatus = TransformNode(child, tmp, ctx);
  3103. UpdateStatusIfChanged(childStatus, child, tmp);
  3104. YQL_ENSURE(tmp == child);
  3105. retStatus = retStatus.Combine(childStatus);
  3106. }
  3107. if (retStatus != TStatus::Ok) {
  3108. if (retStatus.Level == TStatus::Error) {
  3109. input->SetState(TExprNode::EState::Error);
  3110. }
  3111. } else {
  3112. input->SetState(TExprNode::EState::ConstrComplete);
  3113. }
  3114. return retStatus;
  3115. }
  3116. case TExprNode::Callable:
  3117. {
  3118. retStatus = TransformChildren(input, output, ctx);
  3119. if (retStatus != TStatus::Ok) {
  3120. if (retStatus != TStatus::Error && input != output) {
  3121. processedPair.first->second = output;
  3122. }
  3123. break;
  3124. }
  3125. CurrentFunctions.top().second = true;
  3126. retStatus = CallableTransformer->Transform(input, output, ctx);
  3127. if (retStatus == TStatus::Error) {
  3128. input->SetState(TExprNode::EState::Error);
  3129. } else if (retStatus == TStatus::Ok) {
  3130. // Sanity check
  3131. for (size_t i = 0; i < input->ChildrenSize(); ++i) {
  3132. YQL_ENSURE(input->Child(i)->GetState() >= TExprNode::EState::ConstrComplete,
  3133. "Child with index " << i << " of callable " << TString{input->Content()}.Quote() << " has bad state after constraint transform");
  3134. }
  3135. input->SetState(TExprNode::EState::ConstrComplete);
  3136. CheckExpected(*input);
  3137. } else if (retStatus == TStatus::Async) {
  3138. CallableInputs.push_back(input);
  3139. input->SetState(TExprNode::EState::ConstrInProgress);
  3140. } else {
  3141. if (output != input.Get()) {
  3142. processedPair.first->second = output;
  3143. HasRenames = true;
  3144. }
  3145. }
  3146. break;
  3147. }
  3148. default:
  3149. YQL_ENSURE(false, "Unknown type");
  3150. }
  3151. if (retStatus.Level != TStatus::Repeat || retStatus.HasRestart) {
  3152. return retStatus;
  3153. }
  3154. input = output;
  3155. }
  3156. }
  3157. TStatus TransformChildren(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
  3158. TStatus combinedStatus = TStatus::Ok;
  3159. TExprNode::TListType newChildren;
  3160. bool updatedChildren = false;
  3161. for (ui32 i = 0; i < input->ChildrenSize(); ++i) {
  3162. const auto child = input->ChildPtr(i);
  3163. TExprNode::TPtr newChild;
  3164. auto childStatus = TransformNode(child, newChild, ctx);
  3165. UpdateStatusIfChanged(childStatus, child, newChild);
  3166. updatedChildren = updatedChildren || (newChild != child);
  3167. combinedStatus = combinedStatus.Combine(childStatus);
  3168. newChildren.emplace_back(std::move(newChild));
  3169. }
  3170. if (combinedStatus != TStatus::Ok) {
  3171. if (combinedStatus.Level == TStatus::Error) {
  3172. input->SetState(TExprNode::EState::Error);
  3173. }
  3174. else if (updatedChildren) {
  3175. output = ctx.ChangeChildren(*input, std::move(newChildren));
  3176. HasRenames = true;
  3177. }
  3178. }
  3179. return combinedStatus;
  3180. }
  3181. void UpdateStatusIfChanged(TStatus& status, const TExprNode::TPtr& input, const TExprNode::TPtr& output) {
  3182. if (status.Level == TStatus::Ok && input != output) {
  3183. status = TStatus(TStatus::Repeat, status.HasRestart);
  3184. }
  3185. }
  3186. void CheckExpected(const TExprNode& input) {
  3187. if constexpr (DisableCheck)
  3188. return;
  3189. if (const auto it = Types.ExpectedConstraints.find(input.UniqueId()); it != Types.ExpectedConstraints.cend()) {
  3190. for (const auto expectedConstr: it->second) {
  3191. if (!Types.DisableConstraintCheck.contains(expectedConstr->GetName())) {
  3192. if (auto newConstr = input.GetConstraint(expectedConstr->GetName())) {
  3193. if (expectedConstr->GetName() == TMultiConstraintNode::Name()) {
  3194. YQL_ENSURE(static_cast<const TMultiConstraintNode*>(newConstr)->FilteredIncludes(*expectedConstr, Types.DisableConstraintCheck), "Rewrite error, unequal " << *newConstr
  3195. << " constraint in node " << input.Content() << ", previous was " << *expectedConstr);
  3196. } else {
  3197. YQL_ENSURE(newConstr->Includes(*expectedConstr), "Rewrite error, unequal " << *newConstr
  3198. << " constraint in node " << input.Content() << ", previous was " << *expectedConstr);
  3199. }
  3200. } else {
  3201. if (expectedConstr->GetName() == TMultiConstraintNode::Name()) {
  3202. // Constraint Multi(0:{Empty},1:{Empty}, ..., N:{Empty}) can be reduced to Empty
  3203. newConstr = input.GetConstraint<TEmptyConstraintNode>();
  3204. }
  3205. YQL_ENSURE(newConstr, "Rewrite error, missing " << *expectedConstr << " constraint in node " << input.Content());
  3206. }
  3207. }
  3208. }
  3209. }
  3210. }
  3211. private:
  3212. TAutoPtr<IGraphTransformer> CallableTransformer;
  3213. std::deque<TExprNode::TPtr> CallableInputs;
  3214. TNodeOnNodeOwnedMap Processed;
  3215. bool HasRenames = false;
  3216. std::stack<std::pair<std::string_view, bool>> CurrentFunctions;
  3217. TTypeAnnotationContext& Types;
  3218. };
  3219. } // namespace
  3220. TAutoPtr<IGraphTransformer> CreateConstraintTransformer(TTypeAnnotationContext& types, bool instantOnly, bool subGraph, bool disableCheck) {
  3221. TAutoPtr<IGraphTransformer> callableTransformer(new TCallableConstraintTransformer(types, instantOnly, subGraph));
  3222. return disableCheck ?
  3223. static_cast<IGraphTransformer*>(new TConstraintTransformer<true>(callableTransformer, types)):
  3224. static_cast<IGraphTransformer*>(new TConstraintTransformer<false>(callableTransformer, types));
  3225. }
  3226. TAutoPtr<IGraphTransformer> CreateDefCallableConstraintTransformer() {
  3227. return new TDefaultCallableConstraintTransformer();
  3228. }
  3229. IGraphTransformer::TStatus UpdateLambdaConstraints(const TExprNode& lambda) {
  3230. const auto args = lambda.Child(0);
  3231. for (const auto& arg: args->Children()) {
  3232. if (arg->GetState() == TExprNode::EState::TypeComplete || arg->GetState() == TExprNode::EState::ConstrPending) {
  3233. arg->SetState(TExprNode::EState::ConstrComplete);
  3234. }
  3235. YQL_ENSURE(arg->GetAllConstraints().empty());
  3236. }
  3237. if (args->GetState() == TExprNode::EState::TypeComplete || args->GetState() == TExprNode::EState::ConstrPending) {
  3238. args->SetState(TExprNode::EState::ConstrComplete);
  3239. }
  3240. if (lambda.GetState() != TExprNode::EState::ConstrComplete) {
  3241. return IGraphTransformer::TStatus::Repeat;
  3242. }
  3243. return IGraphTransformer::TStatus::Ok;
  3244. }
  3245. IGraphTransformer::TStatus UpdateLambdaConstraints(TExprNode::TPtr& lambda, TExprContext& ctx, const TArrayRef<const TConstraintNode::TListType>& constraints) {
  3246. bool updateArgs = false;
  3247. const auto args = lambda->Child(0);
  3248. YQL_ENSURE(args->ChildrenSize() == constraints.size());
  3249. size_t i = 0;
  3250. for (const auto& constrList: constraints) {
  3251. const auto arg = args->Child(i++);
  3252. if (arg->GetState() == TExprNode::EState::TypeComplete || arg->GetState() == TExprNode::EState::ConstrPending) {
  3253. for (const auto c: constrList) {
  3254. arg->AddConstraint(c);
  3255. }
  3256. arg->SetState(TExprNode::EState::ConstrComplete);
  3257. } else {
  3258. if (constrList.size() != arg->GetAllConstraints().size() || !AllOf(constrList, [arg] (const TConstraintNode* c) { return arg->GetConstraint(c->GetName()) == c; })) {
  3259. updateArgs = true;
  3260. }
  3261. }
  3262. }
  3263. if (updateArgs) {
  3264. TNodeOnNodeOwnedMap replaces(constraints.size());
  3265. TExprNode::TListType argsChildren;
  3266. argsChildren.reserve(constraints.size());
  3267. i = 0;
  3268. for (const auto& constrList: constraints) {
  3269. const auto arg = args->Child(i++);
  3270. const auto newArg = ctx.ShallowCopy(*arg);
  3271. newArg->SetTypeAnn(arg->GetTypeAnn());
  3272. for (const auto c: constrList) {
  3273. newArg->AddConstraint(c);
  3274. }
  3275. newArg->SetState(TExprNode::EState::ConstrComplete);
  3276. YQL_ENSURE(replaces.emplace(arg, newArg).second);
  3277. argsChildren.emplace_back(std::move(newArg));
  3278. }
  3279. auto newArgs = ctx.NewArguments(args->Pos(), std::move(argsChildren));
  3280. newArgs->SetTypeAnn(ctx.MakeType<TUnitExprType>());
  3281. newArgs->SetState(TExprNode::EState::ConstrComplete);
  3282. const auto type = lambda->GetTypeAnn();
  3283. lambda = ctx.NewLambda(lambda->Pos(), std::move(newArgs), ctx.ReplaceNodes<true>(GetLambdaBody(*lambda), replaces));
  3284. lambda->SetTypeAnn(type);
  3285. lambda->Head().ForEachChild(std::bind(&TExprNode::SetDependencyScope, std::placeholders::_1, lambda.Get(), lambda.Get()));
  3286. return IGraphTransformer::TStatus::Repeat;
  3287. }
  3288. if (args->GetState() == TExprNode::EState::TypeComplete || args->GetState() == TExprNode::EState::ConstrPending) {
  3289. args->SetState(TExprNode::EState::ConstrComplete);
  3290. }
  3291. if (lambda->GetState() != TExprNode::EState::ConstrComplete) {
  3292. return IGraphTransformer::TStatus::Repeat;
  3293. }
  3294. return IGraphTransformer::TStatus::Ok;
  3295. }
  3296. IGraphTransformer::TStatus UpdateAllChildLambdasConstraints(const TExprNode& node) {
  3297. IGraphTransformer::TStatus status = IGraphTransformer::TStatus::Ok;
  3298. for (ui32 i = 0; i < node.ChildrenSize(); ++i) {
  3299. const auto child = node.Child(i);
  3300. if (child->Type() == TExprNode::EType::Lambda) {
  3301. status = status.Combine(UpdateLambdaConstraints(*child));
  3302. }
  3303. }
  3304. return status;
  3305. }
  3306. } // namespace NYql