yql_expr_constraint.cpp 192 KB

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