#include "yql_expr_constraint.h" #include "yql_callable_transform.h" #include "yql_opt_utils.h" #include #include #include #include #include #include #include #include #include #include namespace NYql { using namespace NNodes; namespace { template struct TApplyConstraintFromInput; template struct TApplyConstraintFromInput { static void Do(const TExprNode::TPtr&) { } }; template struct TApplyConstraintFromInput { static void Do(const TExprNode::TPtr& input) { if (auto c = input->Child(FromChild)->GetConstraint()) { input->AddConstraint(c); } TApplyConstraintFromInput::Do(input); } }; template const TConstraint* MakeCommonConstraint(const TExprNode::TPtr& input, size_t from, TExprContext& ctx) { TVector constraints; for (size_t i = from; i < input->ChildrenSize(); ++i) { constraints.push_back(&input->Child(i)->GetConstraintSet()); } return TConstraint::MakeCommon(constraints, ctx); } template struct TApplyCommonConstraint; template struct TApplyCommonConstraint { static void Do(const TExprNode::TPtr& input, const std::vector& constraints, TExprContext& ctx) { if (auto c = TConstraint::MakeCommon(constraints, ctx)) { input->AddConstraint(c); } } }; template struct TApplyCommonConstraint { static void Do(const TExprNode::TPtr& input, const std::vector& constraints, TExprContext& ctx) { if (auto c = TConstraint::MakeCommon(constraints, ctx)) { input->AddConstraint(c); } TApplyCommonConstraint::Do(input, constraints, ctx); } }; class TCallableConstraintTransformer : public TCallableTransformerBase { using THandler = TStatus(TCallableConstraintTransformer::*)(const TExprNode::TPtr&, TExprNode::TPtr&, TExprContext&) const; public: TCallableConstraintTransformer(TTypeAnnotationContext& types, bool instantOnly, bool subGraph) : TCallableTransformerBase(types, instantOnly) , SubGraph(subGraph) { Functions["Unordered"] = &TCallableConstraintTransformer::FromFirst; Functions["UnorderedSubquery"] = &TCallableConstraintTransformer::FromFirst; Functions["Sort"] = &TCallableConstraintTransformer::SortWrap; Functions["AssumeSorted"] = &TCallableConstraintTransformer::SortWrap; Functions["AssumeUnique"] = &TCallableConstraintTransformer::AssumeUniqueWrap; Functions["AssumeDistinct"] = &TCallableConstraintTransformer::AssumeUniqueWrap; Functions["AssumeUniqueHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap; Functions["AssumeDistinctHint"] = &TCallableConstraintTransformer::AssumeUniqueWrap; Functions["AssumeConstraints"] = &TCallableConstraintTransformer::AssumeConstraintsWrap; Functions["AssumeChopped"] = &TCallableConstraintTransformer::AssumeChoppedWrap; Functions["AssumeColumnOrder"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["AssumeAllMembersNullableAtOnce"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["Top"] = &TCallableConstraintTransformer::TopWrap; Functions["TopSort"] = &TCallableConstraintTransformer::TopWrap; Functions["TakeWhile"] = &TCallableConstraintTransformer::FilterWrap; Functions["SkipWhile"] = &TCallableConstraintTransformer::FilterWrap; Functions["TakeWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap; Functions["SkipWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap; Functions["WideTakeWhile"] = &TCallableConstraintTransformer::FilterWrap; Functions["WideSkipWhile"] = &TCallableConstraintTransformer::FilterWrap; Functions["WideTakeWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap; Functions["WideSkipWhileInclusive"] = &TCallableConstraintTransformer::FilterWrap; Functions["Iterator"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["ForwardList"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["LazyList"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["ToFlow"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["FromFlow"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["ToStream"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["ToSequence"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["Collect"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["FilterNullMembers"] = &TCallableConstraintTransformer::FromFirst; Functions["SkipNullMembers"] = &TCallableConstraintTransformer::FromFirst; Functions["FilterNullElements"] = &TCallableConstraintTransformer::FromFirst; Functions["SkipNullElements"] = &TCallableConstraintTransformer::FromFirst; Functions["Right!"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["Cons!"] = &TCallableConstraintTransformer::CopyAllFrom<1>; Functions["ExtractMembers"] = &TCallableConstraintTransformer::ExtractMembersWrap; Functions["RemoveSystemMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap; Functions["RemovePrefixMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap; Functions["FlattenMembers"] = &TCallableConstraintTransformer::FlattenMembersWrap; Functions["SelectMembers"] = &TCallableConstraintTransformer::SelectMembersWrap; Functions["FilterMembers"] = &TCallableConstraintTransformer::SelectMembersWrap; Functions["CastStruct"] = &TCallableConstraintTransformer::SelectMembersWrap; Functions["SafeCast"] = &TCallableConstraintTransformer::CastWrap; Functions["StrictCast"] = &TCallableConstraintTransformer::CastWrap; Functions["ToString"] = &TCallableConstraintTransformer::CastWrap; Functions["ToBytes"] = &TCallableConstraintTransformer::CastWrap; Functions["DivePrefixMembers"] = &TCallableConstraintTransformer::DivePrefixMembersWrap; Functions["OrderedFilter"] = &TCallableConstraintTransformer::FilterWrap; Functions["Filter"] = &TCallableConstraintTransformer::FilterWrap; Functions["WideFilter"] = &TCallableConstraintTransformer::FilterWrap; Functions["OrderedMap"] = &TCallableConstraintTransformer::MapWrap; Functions["Map"] = &TCallableConstraintTransformer::MapWrap; Functions["MapNext"] = &TCallableConstraintTransformer::MapWrap; Functions["OrderedFlatMap"] = &TCallableConstraintTransformer::MapWrap; Functions["FlatMap"] = &TCallableConstraintTransformer::MapWrap; Functions["OrderedMultiMap"] = &TCallableConstraintTransformer::MapWrap; Functions["MultiMap"] = &TCallableConstraintTransformer::MapWrap; Functions["ExpandMap"] = &TCallableConstraintTransformer::MapWrap; Functions["WideMap"] = &TCallableConstraintTransformer::MapWrap; Functions["NarrowMap"] = &TCallableConstraintTransformer::MapWrap; Functions["NarrowFlatMap"] = &TCallableConstraintTransformer::MapWrap; Functions["NarrowMultiMap"] = &TCallableConstraintTransformer::MapWrap; Functions["OrderedFlatMapToEquiJoin"] = &TCallableConstraintTransformer::MapWrap; Functions["FlatMapToEquiJoin"] = &TCallableConstraintTransformer::MapWrap; Functions["OrderedLMap"] = &TCallableConstraintTransformer::LMapWrap; Functions["LMap"] = &TCallableConstraintTransformer::LMapWrap; Functions["Extract"] = &TCallableConstraintTransformer::FromFirst; Functions["OrderedExtract"] = &TCallableConstraintTransformer::FromFirst; Functions["OrderedExtend"] = &TCallableConstraintTransformer::ExtendWrap; Functions["Extend"] = &TCallableConstraintTransformer::ExtendWrap; Functions["UnionAll"] = &TCallableConstraintTransformer::ExtendWrap; Functions["Merge"] = &TCallableConstraintTransformer::MergeWrap; Functions["UnionMerge"] = &TCallableConstraintTransformer::MergeWrap; Functions["Skip"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["Take"] = &TCallableConstraintTransformer::TakeWrap; Functions["Limit"] = &TCallableConstraintTransformer::TakeWrap; Functions["Member"] = &TCallableConstraintTransformer::MemberWrap; Functions["AsStruct"] = &TCallableConstraintTransformer::AsStructWrap; Functions["BlockAsStruct"] = &TCallableConstraintTransformer::AsStructWrap; Functions["Just"] = &TCallableConstraintTransformer::FromFirst; Functions["Unwrap"] = &TCallableConstraintTransformer::FromFirst; Functions["Ensure"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["ToList"] = &TCallableConstraintTransformer::FromFirst; Functions["ToOptional"] = &TCallableConstraintTransformer::FromFirst; Functions["Head"] = &TCallableConstraintTransformer::FromFirst; Functions["Last"] = &TCallableConstraintTransformer::FromFirst; Functions["Reverse"] = &TCallableConstraintTransformer::ReverseWrap; Functions["Replicate"] = &TCallableConstraintTransformer::FromFirst; Functions["AddMember"] = &TCallableConstraintTransformer::AddMemberWrap; Functions["RemoveMember"] = &TCallableConstraintTransformer::RemoveMemberWrap; Functions["ForceRemoveMember"] = &TCallableConstraintTransformer::RemoveMemberWrap; Functions["ReplaceMember"] = &TCallableConstraintTransformer::ReplaceMemberWrap; Functions["AsList"] = &TCallableConstraintTransformer::AsListWrap; Functions["OptionalIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap; Functions["FlatOptionalIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap; Functions["ListIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap; Functions["FlatListIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap; Functions["EmptyIterator"] = &TCallableConstraintTransformer::FromEmpty; Functions["EmptyFrom"] = &TCallableConstraintTransformer::EmptyFromWrap; Functions["List"] = &TCallableConstraintTransformer::ListWrap; Functions["Dict"] = &TCallableConstraintTransformer::DictWrap; Functions["EmptyList"] = &TCallableConstraintTransformer::FromEmpty; Functions["EmptyDict"] = &TCallableConstraintTransformer::FromEmpty; Functions["DictFromKeys"] = &TCallableConstraintTransformer::DictFromKeysWrap; Functions["If"] = &TCallableConstraintTransformer::IfWrap; Functions["Nothing"] = &TCallableConstraintTransformer::FromEmpty; Functions["IfPresent"] = &TCallableConstraintTransformer::IfPresentWrap; Functions["Coalesce"] = &TCallableConstraintTransformer::CommonFromChildren<0, TSortedConstraintNode, TPartOfSortedConstraintNode, TChoppedConstraintNode, TPartOfChoppedConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>; Functions["CombineByKey"] = &TCallableConstraintTransformer::FromFinalLambda; Functions["FinalizeByKey"] = &TCallableConstraintTransformer::FromFinalLambda; Functions["CombineCore"] = &TCallableConstraintTransformer::FromFinalLambda; Functions["PartitionByKey"] = &TCallableConstraintTransformer::ShuffleByKeysWrap; Functions["PartitionsByKeys"] = &TCallableConstraintTransformer::ShuffleByKeysWrap; Functions["ShuffleByKeys"] = &TCallableConstraintTransformer::ShuffleByKeysWrap; Functions["Switch"] = &TCallableConstraintTransformer::SwitchWrap; Functions["Visit"] = &TCallableConstraintTransformer::VisitWrap; Functions["VariantItem"] = &TCallableConstraintTransformer::VariantItemWrap; Functions["Variant"] = &TCallableConstraintTransformer::VariantWrap; Functions["Guess"] = &TCallableConstraintTransformer::GuessWrap; Functions["Mux"] = &TCallableConstraintTransformer::MuxWrap; Functions["Nth"] = &TCallableConstraintTransformer::NthWrap; Functions["EquiJoin"] = &TCallableConstraintTransformer::EquiJoinWrap; Functions["JoinDict"] = &TCallableConstraintTransformer::JoinDictWrap; Functions["MapJoinCore"] = &TCallableConstraintTransformer::MapJoinCoreWrap; Functions["GraceJoinCore"] = &TCallableConstraintTransformer::GraceJoinCoreWrap; Functions["GraceSelfJoinCore"] = &TCallableConstraintTransformer::GraceSelfJoinCoreWrap; Functions["CommonJoinCore"] = &TCallableConstraintTransformer::FromFirst; Functions["ToDict"] = &TCallableConstraintTransformer::ToDictWrap; Functions["DictItems"] = &TCallableConstraintTransformer::DictItemsWrap; Functions["DictKeys"] = &TCallableConstraintTransformer::DictHalfWrap; Functions["DictPayloads"] = &TCallableConstraintTransformer::DictHalfWrap; Functions["Chain1Map"] = &TCallableConstraintTransformer::Chain1MapWrap; Functions["WideChain1Map"] = &TCallableConstraintTransformer::Chain1MapWrap; Functions["IsKeySwitch"] = &TCallableConstraintTransformer::IsKeySwitchWrap; Functions["Condense"] = &TCallableConstraintTransformer::CondenseWrap; Functions["Condense1"] = &TCallableConstraintTransformer::Condense1Wrap; Functions["GroupingCore"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; Functions["Chopper"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; Functions["WideChopper"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; Functions["WideCombiner"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; Functions["WideCondense1"] = &TCallableConstraintTransformer::Condense1Wrap; Functions["Aggregate"] = &TCallableConstraintTransformer::AggregateWrap; Functions["AggregateMergeState"] = &TCallableConstraintTransformer::AggregateWrap; Functions["AggregateMergeFinalize"] = &TCallableConstraintTransformer::AggregateWrap; Functions["AggregateMergeManyFinalize"] = &TCallableConstraintTransformer::AggregateWrap; Functions["AggregateFinalize"] = &TCallableConstraintTransformer::AggregateWrap; Functions["AggregateCombine"] = &TCallableConstraintTransformer::AggregateWrap; Functions["AggregateCombineState"] = &TCallableConstraintTransformer::AggregateWrap; Functions["Fold"] = &TCallableConstraintTransformer::FoldWrap; Functions["Fold1"] = &TCallableConstraintTransformer::FoldWrap; Functions["WithContext"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["WithWorld"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["WideTop"] = &TCallableConstraintTransformer::WideTopWrap; Functions["WideTopSort"] = &TCallableConstraintTransformer::WideTopWrap; Functions["WideSort"] = &TCallableConstraintTransformer::WideTopWrap; Functions["WideTopBlocks"] = &TCallableConstraintTransformer::WideTopWrap; Functions["WideTopSortBlocks"] = &TCallableConstraintTransformer::WideTopWrap; Functions["WideSortBlocks"] = &TCallableConstraintTransformer::WideTopWrap; Functions["WideToBlocks"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["WideFromBlocks"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["BlockExpandChunked"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["ReplicateScalars"] = &TCallableConstraintTransformer::CopyAllFrom<0>; Functions["BlockMergeFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap; Functions["BlockMergeManyFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap; Functions["MultiHoppingCore"] = &TCallableConstraintTransformer::MultiHoppingCoreWrap; Functions["StablePickle"] = &TCallableConstraintTransformer::FromFirst; Functions["Unpickle"] = &TCallableConstraintTransformer::FromSecond; } std::optional ProcessCore(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) { if (const auto func = Functions.find(input->Content()); Functions.cend() != func) { return (this->*func->second)(input, output, ctx); } return std::nullopt; } std::optional ProcessList(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) { if (!input->ChildrenSize() || ETypeAnnotationKind::Tuple != input->GetTypeAnn()->GetKind()) return TStatus::Ok; return AsTupleWrap(input, output, ctx); } TStatus ProcessUnknown(const TExprNode::TPtr& input, TExprContext&) { return UpdateAllChildLambdasConstraints(*input); } TStatus ValidateProviderCommitResult(const TExprNode::TPtr&, TExprContext&) { return TStatus::Ok; } TStatus ValidateProviderReadResult(const TExprNode::TPtr&, TExprContext&) { return TStatus::Ok; } TStatus ValidateProviderWriteResult(const TExprNode::TPtr&, TExprContext&) { return TStatus::Ok; } TStatus ValidateProviderConfigureResult(const TExprNode::TPtr&, TExprContext&) { return TStatus::Ok; } IGraphTransformer& GetTransformer(IDataProvider& provider) const { return provider.GetConstraintTransformer(InstantOnly, SubGraph); } private: template TStatus CopyAllFrom(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { Y_UNUSED(output); Y_UNUSED(ctx); input->CopyConstraints(*input->Child(Ndx)); return TStatus::Ok; } template TStatus FromFirst(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { Y_UNUSED(output); Y_UNUSED(ctx); TApplyConstraintFromInput<0, TConstraints...>::Do(input); return TStatus::Ok; } template TStatus FromSecond(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { Y_UNUSED(output); Y_UNUSED(ctx); TApplyConstraintFromInput<1, TConstraints...>::Do(input); return TStatus::Ok; } template TStatus CommonFromChildren(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { Y_UNUSED(output); TVector constraints; for (size_t i = StartFromChild; i < input->ChildrenSize(); ++i) { constraints.push_back(&input->Child(i)->GetConstraintSet()); } TApplyCommonConstraint::Do(input, constraints, ctx); return TStatus::Ok; } TStatus FromEmpty(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { input->AddConstraint(ctx.MakeConstraint()); return TStatus::Ok; } TStatus EmptyFromWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { auto set = input->Head().GetConstraintSet(); set.RemoveConstraint(TEmptyConstraintNode::Name()); if (!set) { const auto type = input->GetTypeAnn(); output = ctx.NewCallable(input->Pos(), GetEmptyCollectionName(type), {ExpandType(input->Pos(), *type, ctx)}); return TStatus::Repeat; } set.AddConstraint(ctx.MakeConstraint()); input->SetConstraints(set); return TStatus::Ok; } template TStatus FromFinalLambda(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { TStatus status = UpdateAllChildLambdasConstraints(*input); if (status != TStatus::Ok) { return status; } TApplyConstraintFromInput::Do(input); return FromFirst(input, output, ctx); } TStatus InheriteEmptyFromInput(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { auto status = UpdateAllChildLambdasConstraints(*input); if (status != TStatus::Ok) { return status; } return FromFirst(input, output, ctx); } template TStatus WideTopWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if constexpr (Sort) { TSortedConstraintNode::TContainerType sorted; sorted.reserve(input->Tail().ChildrenSize()); for (const auto& item : input->Tail().Children()) { if (item->Tail().IsCallable("Bool")) sorted.emplace_back(std::make_pair(TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType(1U, item->Head().Content())}, FromString(item->Tail().Tail().Content()))); else break; } if (!sorted.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(sorted))); } } return FromFirst(input, output, ctx); } TStatus SortWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (const auto status = UpdateLambdaConstraints(input->Tail()); status != TStatus::Ok) { return status; } if (const auto sorted = DeduceSortConstraint(*input->Child(1), *input->Child(2), ctx)) { input->AddConstraint(sorted->GetSimplifiedForType(*input->GetTypeAnn(), ctx)); } return FromFirst(input, output, ctx); } TStatus AssumeConstraintsWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { TConstraintSet set; try { set = ctx.MakeConstraintSet(NYT::NodeFromYsonString(input->Tail().Content())); } catch (...) { ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << "Bad constraints yson-value: " << CurrentExceptionMessage())); return IGraphTransformer::TStatus::Error; } if (!set) { ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), "AssumeConstraints with empty constraints set")); return IGraphTransformer::TStatus::Error; } for (auto constraint: set.GetAllConstraints()) { if (!constraint->IsApplicableToType(*input->GetTypeAnn())) { ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint << " is not applicable to " << *input->GetTypeAnn())); return IGraphTransformer::TStatus::Error; } } for (auto constr: input->Head().GetAllConstraints()) { if (!constr->GetName().starts_with("PartOf") && !set.GetConstraint(constr->GetName())) { set.AddConstraint(constr); } } input->SetConstraints(set); return IGraphTransformer::TStatus::Ok; } template TStatus AssumeUniqueWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { typename TUniqueConstraintNodeBase::TContentType content; for (auto i = 1U; i < input->ChildrenSize(); ++i) { TPartOfConstraintBase::TSetOfSetsType sets; sets.reserve(input->Child(i)->ChildrenSize()); for (const auto& list : input->Child(i)->Children()) { if (list->IsAtom()) sets.insert_unique(TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType(1U, list->Content())}); else if (list->IsList()) { TPartOfConstraintBase::TSetType columns; columns.reserve(list->ChildrenSize()); for (const auto& column: list->Children()) { if (column->IsAtom()) columns.insert_unique(TPartOfConstraintBase::TPathType(1U, column->Content())); else if (column->IsList()) { TPartOfConstraintBase::TPathType path(column->ChildrenSize()); std::transform(column->Children().cbegin(), column->Children().cend(), path.begin(), [](const TExprNode::TPtr& atom) { return atom->Content(); } ); columns.insert_unique(std::move(path)); } } sets.insert_unique(std::move(columns)); } } content.insert_unique(std::move(sets)); } if (content.empty()) content.insert_unique(TPartOfConstraintBase::TSetOfSetsType{TPartOfConstraintBase::TSetType{TPartOfConstraintBase::TPathType()}}); auto constraint = ctx.MakeConstraint>(std::move(content)); if (!constraint->IsApplicableToType(*input->GetTypeAnn())) { if constexpr (Strict) { ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint << " is not applicable to " << *input->GetTypeAnn())); } else { auto issue = TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << (Distinct ? "Distinct" : "Unique") << " sql hint contains invalid column: " << Endl << *constraint << " is not applicable to " << *input->GetTypeAnn()); SetIssueCode(EYqlIssueCode::TIssuesIds_EIssueCode_YQL_HINT_INVALID_PARAMETERS, issue); if (ctx.AddWarning(issue)) { output = input->HeadPtr(); return IGraphTransformer::TStatus::Repeat; } } return IGraphTransformer::TStatus::Error; } if constexpr (!Strict) { output = ctx.RenameNode(*input, Distinct ? "AssumeDistinct" : "AssumeUnique"); return IGraphTransformer::TStatus::Repeat; } if (const auto old = input->Head().GetConstraint>()) { if (old->Includes(*constraint)) { output = input->HeadPtr(); return TStatus::Repeat; } else constraint = TUniqueConstraintNodeBase::Merge(old, constraint, ctx); } input->AddConstraint(constraint); return FromFirst, TEmptyConstraintNode, TVarIndexConstraintNode>(input, output, ctx); } TStatus AssumeChoppedWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { TPartOfConstraintBase::TSetOfSetsType sets; for (auto i = 1U; i < input->ChildrenSize(); ++i) { TPartOfConstraintBase::TSetType columns; columns.reserve(input->Child(i)->ChildrenSize()); for (const auto& column: input->Child(i)->Children()) { if (column->IsAtom()) columns.insert_unique(TPartOfConstraintBase::TPathType(1U, column->Content())); else if (column->IsList()) { TPartOfConstraintBase::TPathType path(column->ChildrenSize()); std::transform(column->Children().cbegin(), column->Children().cend(), path.begin(), [](const TExprNode::TPtr& atom) { return atom->Content(); } ); columns.insert_unique(std::move(path)); } } sets.insert_unique(std::move(columns)); } const auto constraint = ctx.MakeConstraint(std::move(sets)); if (!constraint->IsApplicableToType(*input->GetTypeAnn())) { ctx.AddError(TIssue(ctx.GetPosition(input->Pos()), TStringBuilder() << *constraint << " is not applicable to " << *input->GetTypeAnn())); return IGraphTransformer::TStatus::Error; } if (const auto old = input->Head().GetConstraint()) { if (old->Equals(*constraint)) { output = input->HeadPtr(); return TStatus::Repeat; } } input->AddConstraint(constraint); return FromFirst(input, output, ctx); } template TStatus TopWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (const auto status = UpdateLambdaConstraints(input->Tail()); status != TStatus::Ok) { return status; } if constexpr (UseSort) { if (const auto sorted = DeduceSortConstraint(*input->Child(2), *input->Child(3), ctx)) { input->AddConstraint(sorted->GetSimplifiedForType(*input->GetTypeAnn(), ctx)); } } return FromFirst(input, output, ctx); } template static void FilterFromHead(const TExprNode& input, TConstraintSet& constraints, const TPartOfConstraintBase::TPathFilter& filter, TExprContext& ctx) { if (const auto source = input.Head().GetConstraint()) { if (const auto filtered = source->FilterFields(ctx, filter)) { constraints.AddConstraint(filtered); } } } template static void ReduceFromHead(const TExprNode::TPtr& input, const TPartOfConstraintBase::TPathReduce& reduce, TExprContext& ctx) { if (const auto source = input->Head().GetConstraint()) { if (const auto filtered = source->RenameFields(ctx, reduce)) { if constexpr (Simplify) input->AddConstraint(filtered->GetSimplifiedForType(*input->GetTypeAnn(), ctx)); else input->AddConstraint(filtered); } } } template static void FilterFromHead(const TExprNode::TPtr& input, const TPartOfConstraintBase::TPathFilter& filter, TExprContext& ctx) { if (const auto source = input->Head().GetConstraint()) { if (const auto filtered = source->FilterFields(ctx, filter)) { if constexpr (Simplify) input->AddConstraint(filtered->GetSimplifiedForType(*input->GetTypeAnn(), ctx)); else input->AddConstraint(filtered); } } } template static void FilterFromHeadIfMissed(const TExprNode::TPtr& input, const TPartOfConstraintBase::TPathFilter& filter, TExprContext& ctx) { if (!input->GetConstraint()) FilterFromHead(input, filter, ctx); } TStatus SelectMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { auto outItemType = input->GetTypeAnn(); while (outItemType->GetKind() == ETypeAnnotationKind::Optional) { outItemType = outItemType->Cast()->GetItemType(); } if (outItemType->GetKind() == ETypeAnnotationKind::Variant) { if (outItemType->Cast()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) { const auto outSize = outItemType->Cast()->GetUnderlyingType()->Cast()->GetSize(); auto multi = input->Head().GetConstraint(); if (multi && multi->GetItems().back().first >= outSize) { TMultiConstraintNode::TMapType filteredItems; for (auto& item: multi->GetItems()) { if (item.first < outSize) { filteredItems.push_back(item); } } multi = filteredItems.empty() ? nullptr : ctx.MakeConstraint(std::move(filteredItems)); } if (multi) { input->AddConstraint(multi); } auto varIndex = input->Head().GetConstraint(); if (varIndex && varIndex->GetIndexMapping().back().first >= outSize) { TVarIndexConstraintNode::TMapType filteredItems; for (auto& item: varIndex->GetIndexMapping()) { if (item.first < outSize) { filteredItems.push_back(item); } } varIndex = filteredItems.empty() ? nullptr : ctx.MakeConstraint(std::move(filteredItems)); } if (varIndex) { input->AddConstraint(varIndex); } } } else if (outItemType->GetKind() == ETypeAnnotationKind::Struct) { const auto filter = [outItemType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType); }; FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); } return TStatus::Ok; } template TStatus CastWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { const auto outItemType = input->GetTypeAnn(); const auto inItemType = input->Head().GetTypeAnn(); const auto filter = [inItemType, outItemType, toString = input->IsCallable({"ToString", "ToBytes"})](const TPartOfConstraintBase::TPathType& path) { if (const auto outType = TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType)) { const auto inType = TPartOfConstraintBase::GetSubTypeByPath(path, *inItemType); return (toString && inType->GetKind() == ETypeAnnotationKind::Data && inType->Cast()->GetSlot() == EDataSlot::Utf8) || IsSameAnnotation(*outType, *inType); } return false; }; const auto filterForUnique = [inItemType, outItemType](const TPartOfConstraintBase::TPathType& path) { const auto castResult = CastResult(TPartOfConstraintBase::GetSubTypeByPath(path, *inItemType), TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType)); return NUdf::ECastOptions::Complete == castResult || NUdf::ECastOptions::MayFail == castResult; }; const auto filterForDistinct = [inItemType, outItemType](const TPartOfConstraintBase::TPathType& path) { return NUdf::ECastOptions::Complete == CastResult(TPartOfConstraintBase::GetSubTypeByPath(path, *inItemType), TPartOfConstraintBase::GetSubTypeByPath(path, *outItemType)); }; FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filterForUnique, ctx); FilterFromHead(input, filterForDistinct, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filterForUnique, ctx); FilterFromHead(input, filterForDistinct, ctx); return TStatus::Ok; } TStatus DivePrefixMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { const auto prefixes = input->Tail().Children(); const auto rename = [&prefixes](const TPartOfConstraintBase::TPathType& path) -> std::vector { if (path.empty()) return {}; for (const auto& p : prefixes) { if (const auto& prefix = p->Content(); path.front().starts_with(prefix)) { auto out = path; out.front() = out.front().substr(prefix.length()); return {std::move(out)}; } } return {}; }; ReduceFromHead(input, rename, ctx); ReduceFromHead(input, rename, ctx); ReduceFromHead(input, rename, ctx); ReduceFromHead(input, rename, ctx); return FromFirst(input, output, ctx); } TStatus ExtractMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { const auto outItemType = GetSeqItemType(*input->GetTypeAnn()).Cast(); const auto filter = [outItemType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && outItemType->FindItem(path.front()); }; FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); return FromFirst(input, output, ctx); } TStatus RemovePrefixMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { const TTypeAnnotationNode* outItemType = GetSeqItemType(input->GetTypeAnn()); if (!outItemType) { outItemType = input->GetTypeAnn(); } if (outItemType->GetKind() == ETypeAnnotationKind::Struct) { const auto outStructType = outItemType->Cast(); const auto filter = [outStructType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && outStructType->FindItem(path.front()); }; FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); } else if (outItemType->GetKind() == ETypeAnnotationKind::Variant) { if (auto multi = input->Head().GetConstraint()) { TMultiConstraintNode::TMapType multiItems; auto tupleUnderType = outItemType->Cast()->GetUnderlyingType()->Cast(); for (auto& item: multi->GetItems()) { YQL_ENSURE(item.first < tupleUnderType->GetSize()); auto& constr = multiItems[item.first]; const auto outStructType = tupleUnderType->GetItems()[item.first]->Cast(); const auto filter = [outStructType](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && outStructType->FindItem(path.front()); }; FilterFromHead(*input, constr, filter, ctx); FilterFromHead(*input, constr, filter, ctx); FilterFromHead(*input, constr, filter, ctx); FilterFromHead(*input, constr, filter, ctx); } input->AddConstraint(ctx.MakeConstraint(std::move(multiItems))); } } return FromFirst(input, output, ctx); } // TODO: Empty for false condition template TStatus FilterWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (const auto status = UpdateLambdaConstraints(*input->Child(1)); status != TStatus::Ok) { return status; } if constexpr (Ordered) { FromFirst(input, output, ctx); } return FromFirst(input, output, ctx); } template static const TConstraint* GetConstraintFromWideResultLambda(const TExprNode& lambda, TExprContext& ctx); template static const TConstraintType* GetLambdaConstraint(const TExprNode& lambda, TExprContext& ctx) { if (2U == lambda.ChildrenSize()) return lambda.Tail().GetConstraint(); TVector constraints; constraints.reserve(lambda.ChildrenSize() - 1U); for (size_t i = 1U; i < lambda.ChildrenSize(); ++i) { constraints.emplace_back(&lambda.Child(i)->GetConstraintSet()); } return TConstraintType::MakeCommon(constraints, ctx); } template static const TConstraintType* GetConstraintFromLambda(const TExprNode& lambda, TExprContext& ctx) { if constexpr (WideLambda) return GetConstraintFromWideResultLambda(lambda, ctx); else return GetLambdaConstraint(lambda, ctx); } static std::optional GetDirection(const TExprNode& dir) { if (dir.IsCallable("Bool")) return IsTrue(dir.Tail().Content()); if (dir.IsCallable("Not")) if (const auto d = GetDirection(dir.Head())) return !*d; return std::nullopt; } static std::vector> ExtractSimpleSortTraits(const TExprNode& sortDirections, const TExprNode& keySelectorLambda) { const auto& keySelectorBody = keySelectorLambda.Tail(); const auto& keySelectorArg = keySelectorLambda.Head().Head(); std::vector> columns; if (const auto dir = GetDirection(sortDirections)) columns.emplace_back(TPartOfConstraintBase::TPathType(), *dir); else if (sortDirections.IsList()) if (const auto size = keySelectorBody.ChildrenSize()) { columns.reserve(size); for (auto i = 0U; i < size; ++i) if (const auto dir = GetDirection(*sortDirections.Child(i))) columns.emplace_back(TPartOfConstraintBase::TPathType(), *dir); else return {}; } else return {}; else return {}; if (keySelectorBody.IsList()) if (const auto size = keySelectorBody.ChildrenSize()) { TPartOfConstraintBase::TSetType set; set.reserve(size); columns.resize(size, std::make_pair(TPartOfConstraintBase::TPathType(), columns.back().second)); auto it = columns.begin(); for (auto i = 0U; i < size; ++i) { if (auto path = GetPathToKey(*keySelectorBody.Child(i), keySelectorArg)) { if (set.insert(*path).second) it++->first = std::move(*path); else if (columns.cend() != it) it = columns.erase(it); } else { return {}; } } } else return {}; else if (auto path = GetPathToKey(keySelectorBody, keySelectorArg)) if (columns.size() == 1U) columns.front().first = std::move(*path); else return {}; else return {}; return columns; } template static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, bool isSingleItem, TExprContext& ctx) { constexpr bool isOrderConstraint = std::is_same() || std::is_same(); if (const auto lambda = handler.GetConstraint()) { const auto original = input.template GetConstraint(); if constexpr (OrderedMap || !isOrderConstraint) { if (original) { if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) { output.AddConstraint(complete); } } } if (const auto part = input.template GetConstraint()) { auto mapping = lambda->GetColumnMapping(); for (auto it = mapping.cbegin(); mapping.cend() != it;) { if (part->GetColumnMapping().contains(it->first)) ++it; else it = mapping.erase(it); } if (!mapping.empty()) { output.AddConstraint(ctx.MakeConstraint(std::move(mapping))); } } else if (isOrderConstraint || isSingleItem) { if (const auto filtered = lambda->RemoveOriginal(ctx, original)) output.AddConstraint(filtered); } } } template static void GetFromMapLambda(const TExprNode::TPtr& input, bool isSingleItem, TExprContext& ctx) { constexpr bool isOrderConstraint = std::is_same() || std::is_same(); if (const auto lambda = GetConstraintFromLambda(input->Tail(), ctx)) { const auto original = GetDetailed(input->Head().GetConstraint(), *input->Head().GetTypeAnn(), ctx); if constexpr (OrderedMap || !isOrderConstraint) { if (original) { if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) { input->AddConstraint(complete->GetSimplifiedForType(*input->GetTypeAnn(), ctx)); } } } if (const auto part = input->Head().GetConstraint()) { auto mapping = lambda->GetColumnMapping(); for (auto it = mapping.cbegin(); mapping.cend() != it;) { if (part->GetColumnMapping().contains(it->first)) ++it; else it = mapping.erase(it); } if (!mapping.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(mapping))); } } else if (isOrderConstraint || isSingleItem) { if (const auto filtered = lambda->RemoveOriginal(ctx, original)) input->AddConstraint(filtered); } } } template TStatus MapWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { auto argConstraints = GetConstraintsForInputArgument(*input, ctx); if constexpr (Ordered && !(Flat || WideInput || WideOutput)) { // TODO: is temporary crutch for MapNext. if (argConstraints.size() < input->Tail().Head().ChildrenSize()) argConstraints.resize(input->Tail().Head().ChildrenSize(), argConstraints.front()); } if (const auto status = UpdateLambdaConstraints(input->TailRef(), ctx, argConstraints); status != TStatus::Ok) { return status; } const bool singleItem = ETypeAnnotationKind::Optional == input->GetTypeAnn()->GetKind(); GetFromMapLambda(input, singleItem, ctx); GetFromMapLambda(input, singleItem, ctx); GetFromMapLambda(input, singleItem, ctx); GetFromMapLambda(input, singleItem, ctx); const auto lambdaVarIndex = GetConstraintFromLambda(input->Tail(), ctx); const auto lambdaMulti = GetConstraintFromLambda(input->Tail(), ctx); const auto inItemType = GetSeqItemType(input->Head().GetTypeAnn()); const bool multiInput = ETypeAnnotationKind::Variant == inItemType->GetKind(); if (const auto varIndex = input->Head().GetConstraint()) { if (multiInput) { if (lambdaVarIndex) { if (const auto outVarIndex = GetVarIndexOverVarIndexConstraint(*varIndex, *lambdaVarIndex, ctx)) { input->AddConstraint(outVarIndex); } } } else { if (lambdaMulti) { TVarIndexConstraintNode::TMapType remapItems; for (auto& multiItem: lambdaMulti->GetItems()) { for (auto& varItem: varIndex->GetIndexMapping()) { remapItems.push_back(std::make_pair(multiItem.first, varItem.second)); } } if (!remapItems.empty()) { ::SortUnique(remapItems); input->AddConstraint(ctx.MakeConstraint(std::move(remapItems))); } } else { input->AddConstraint(varIndex); } } } const auto inputMulti = input->Head().GetConstraint(); if (lambdaMulti && !input->Head().GetConstraint()) { TMultiConstraintNode::TMapType remappedItems; for (auto& item: lambdaMulti->GetItems()) { remappedItems.push_back(std::make_pair(item.first, TConstraintSet{})); if (!multiInput) { // remapping one to many GetFromMapLambda(input->Head(), item.second, remappedItems.back().second, singleItem, ctx); GetFromMapLambda(input->Head(), item.second, remappedItems.back().second, singleItem, ctx); GetFromMapLambda(input->Head(), item.second, remappedItems.back().second, singleItem, ctx); GetFromMapLambda(input->Head(), item.second, remappedItems.back().second, singleItem, ctx); if (const auto empty = item.second.template GetConstraint()) { remappedItems.pop_back(); } } else if (lambdaVarIndex && inputMulti) { const auto range = lambdaVarIndex->GetIndexMapping().equal_range(item.first); switch (std::distance(range.first, range.second)) { case 0: // new index break; case 1: // remapping 1 to 1 if (const auto origConstr = inputMulti->GetItem(range.first->second)) { GetFromMapLambda(*origConstr, item.second, remappedItems.back().second, singleItem, ctx); GetFromMapLambda(*origConstr, item.second, remappedItems.back().second, singleItem, ctx); GetFromMapLambda(*origConstr, item.second, remappedItems.back().second, singleItem, ctx); GetFromMapLambda(*origConstr, item.second, remappedItems.back().second, singleItem, ctx); if (const auto empty = item.second.template GetConstraint()) { remappedItems.pop_back(); } } else { remappedItems.pop_back(); } break; default: // remapping many to one { std::vector nonEmpty; for (auto i = range.first; i != range.second; ++i) { if (auto origConstr = inputMulti->GetItem(i->second)) { nonEmpty.push_back(origConstr); } } EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint(); }); if (nonEmpty.empty()) { remappedItems.back().second.AddConstraint(ctx.MakeConstraint()); } else if (nonEmpty.size() == 1) { remappedItems.back().second = std::move(*nonEmpty.front()); } } } } else { remappedItems.back().second = item.second; } } if (remappedItems) { input->AddConstraint(ctx.MakeConstraint(std::move(remappedItems))); } } else if (inputMulti && lambdaVarIndex) { // Many to one const auto range = lambdaVarIndex->GetIndexMapping().equal_range(0); std::vector nonEmpty; for (auto i = range.first; i != range.second; ++i) { if (auto origConstr = inputMulti->GetItem(i->second)) { nonEmpty.push_back(origConstr); } } EraseIf(nonEmpty, [] (const TConstraintSet* c) { return !!c->GetConstraint(); }); if (nonEmpty.empty()) { input->AddConstraint(ctx.MakeConstraint()); } else if (nonEmpty.size() == 1) { input->SetConstraints(*nonEmpty.front()); } } if constexpr (Flat) { if (const auto lambdaEmpty = GetConstraintFromLambda(input->Tail(), ctx)) { input->AddConstraint(lambdaEmpty); const auto& filter = std::bind(&TPartOfConstraintBase::GetSubTypeByPath, std::placeholders::_1, std::cref(GetSeqItemType(*input->GetTypeAnn()))); FilterFromHeadIfMissed(input, filter, ctx); FilterFromHeadIfMissed(input, filter, ctx); if constexpr (Ordered) { FilterFromHeadIfMissed(input, filter, ctx); FilterFromHeadIfMissed(input, filter, ctx); } } } return FromFirst(input, output, ctx); } template TStatus LMapWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { TConstraintNode::TListType argConstraints; for (const auto c: input->Head().GetAllConstraints()) { if (Ordered || (c->GetName() != TSortedConstraintNode::Name() && c->GetName() != TChoppedConstraintNode::Name())) { argConstraints.push_back(c); } } if (const auto status = UpdateLambdaConstraints(input->TailRef(), ctx, {argConstraints}); status != TStatus::Ok) { return status; } TSet except; if constexpr (!Ordered) { except.insert(TSortedConstraintNode::Name()); except.insert(TChoppedConstraintNode::Name()); } if (input->Tail().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional) { except.insert(TEmptyConstraintNode::Name()); } CopyExcept(*input, input->Tail(), except); return FromFirst(input, output, ctx); } TStatus AsListWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (input->ChildrenSize() == 1) { if (const auto unique = input->Head().GetConstraint()) { input->AddConstraint(unique); } if (const auto unique = input->Head().GetConstraint()) { input->AddConstraint(unique); } if (const auto part = input->Head().GetConstraint()) { input->AddConstraint(part); } if (const auto part = input->Head().GetConstraint()) { input->AddConstraint(part); } } return CommonFromChildren<0, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx); } template TStatus ExtendWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (input->ChildrenSize() == 1) { if (const auto unique = input->Head().GetConstraint()) { input->AddConstraint(unique); } if (const auto part = input->Head().GetConstraint()) { input->AddConstraint(part); } if (const auto unique = input->Head().GetConstraint()) { input->AddConstraint(unique); } if (const auto part = input->Head().GetConstraint()) { input->AddConstraint(part); } if constexpr (Ordered) { if (const auto sorted = input->Head().GetConstraint()) { input->AddConstraint(sorted); } if (const auto part = input->Head().GetConstraint()) { input->AddConstraint(part); } if (const auto sorted = input->Head().GetConstraint()) { input->AddConstraint(sorted); } if (const auto part = input->Head().GetConstraint()) { input->AddConstraint(part); } } } return CommonFromChildren<0, TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx); } template TStatus MergeWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (auto sort = MakeCommonConstraint(input, 0, ctx)) { if (Union && input->ChildrenSize() > 1) { // Check and exclude modified keys from final constraint const auto resultItemType = input->GetTypeAnn()->Cast()->GetItemType(); std::vector inputs; for (const auto& child: input->Children()) { inputs.emplace_back(child->GetTypeAnn()->Cast()->GetItemType()); } auto content = sort->GetContent(); for (auto i = 0U; i < content.size(); ++i) { for (auto it = content[i].first.cbegin(); content[i].first.cend() != it;) { const auto resultItemSubType = TPartOfConstraintBase::GetSubTypeByPath(*it, *resultItemType); YQL_ENSURE(resultItemSubType, "Missing " << *it << " in result type"); auto childNdx = 0U; while (childNdx < input->ChildrenSize()) { if (const auto inputItemSubType = TPartOfConstraintBase::GetSubTypeByPath(*it, *inputs[childNdx])) { if (IsSameAnnotation(*inputItemSubType, *resultItemSubType)) { ++childNdx; continue; } } else { YQL_ENSURE(input->Child(childNdx)->GetConstraint(), "Missing column " << *it << " in non empty input type"); } break; } if (childNdx < input->ChildrenSize()) it = content[i].first.erase(it); else ++it; } if (content[i].first.empty()) { content.resize(i); break; } } sort = content.empty() ? nullptr : ctx.MakeConstraint(std::move(content)); } if (sort) { input->AddConstraint(sort); } } return ExtendWrap(input, output, ctx); } TStatus TakeWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { if (input->Tail().IsCallable("Uint64") && !FromString(input->Tail().Head().Content())) { input->AddConstraint(ctx.MakeConstraint()); } return CopyAllFrom<0>(input, output, ctx); } TStatus MemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { const auto& memberName = input->Tail().Content(); const auto& structNode = input->Head(); if (const auto emptyConstraint = structNode.GetConstraint()) { input->AddConstraint(emptyConstraint); } else { if (const auto part = structNode.GetConstraint()) { if (const auto extracted = part->ExtractField(ctx, memberName)) { input->AddConstraint(extracted); } } if (const auto part = structNode.GetConstraint()) { if (const auto extracted = part->ExtractField(ctx, memberName)) { input->AddConstraint(extracted); } } if (const auto part = structNode.GetConstraint()) { if (const auto extracted = part->ExtractField(ctx, memberName)) { input->AddConstraint(extracted); } } if (const auto part = structNode.GetConstraint()) { if (const auto extracted = part->ExtractField(ctx, memberName)) { input->AddConstraint(extracted); } } } if (structNode.IsCallable("AsStruct")) { for (const auto& child: structNode.Children()) { if (child->Head().IsAtom(memberName)) { TApplyConstraintFromInput<1, TVarIndexConstraintNode>::Do(child); break; } } } else { TApplyConstraintFromInput<0, TVarIndexConstraintNode>::Do(input); } return TStatus::Ok; } TStatus AsTupleWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { TPartOfSortedConstraintNode::TMapType sorted; TPartOfChoppedConstraintNode::TMapType chopped; TPartOfUniqueConstraintNode::TMapType uniques; TPartOfDistinctConstraintNode::TMapType distincts; std::vector structConstraints; for (auto i = 0U; i < input->ChildrenSize(); ++i) { const auto child = input->Child(i); const auto& name = ctx.GetIndexAsString(i); if (const auto part = child->GetConstraint()) { TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(name)); } if (const auto part = child->GetConstraint()) { TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(name)); } if (const auto part = child->GetConstraint()) { TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(name)); } if (const auto part = child->GetConstraint()) { TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(name)); } if (const auto& valueNode = SkipModifiers(child); TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) { structConstraints.push_back(&valueNode->Head().GetConstraintSet()); } else if (valueNode->IsArgument() && ETypeAnnotationKind::Struct != valueNode->GetTypeAnn()->GetKind()) { structConstraints.push_back(&valueNode->GetConstraintSet()); } } if (!sorted.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(sorted))); } if (!chopped.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(chopped))); } if (!uniques.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(uniques))); } if (!distincts.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(distincts))); } if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) { input->AddConstraint(varIndex); } return TStatus::Ok; } TStatus AsStructWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { TPartOfSortedConstraintNode::TMapType sorted; TPartOfChoppedConstraintNode::TMapType chopped; TPartOfUniqueConstraintNode::TMapType uniques; TPartOfDistinctConstraintNode::TMapType distincts; std::vector structConstraints; for (const auto& child : input->Children()) { const auto& name = child->Head().Content(); if (const auto part = child->Tail().GetConstraint()) { TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(name)); } if (const auto part = child->Tail().GetConstraint()) { TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(name)); } if (const auto part = child->Tail().GetConstraint()) { TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(name)); } if (const auto part = child->Tail().GetConstraint()) { TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(name)); } if (const auto valueNode = SkipModifiers(&child->Tail()); TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) { structConstraints.push_back(&valueNode->Head().GetConstraintSet()); } else if (valueNode->Type() == TExprNode::Argument) { structConstraints.push_back(&valueNode->GetConstraintSet()); } } if (!sorted.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(sorted))); } if (!chopped.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(chopped))); } if (!uniques.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(uniques))); } if (!distincts.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(distincts))); } if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) { input->AddConstraint(varIndex); } return TStatus::Ok; } TStatus FlattenMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { TPartOfSortedConstraintNode::TMapType sorted; TPartOfChoppedConstraintNode::TMapType chopped; TPartOfUniqueConstraintNode::TMapType uniques; TPartOfDistinctConstraintNode::TMapType distincts; for (const auto& child : input->Children()) { const auto& prefix = child->Head().Content(); if (const auto part = child->Tail().GetConstraint()) { TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(ctx, prefix)); } if (const auto part = child->Tail().GetConstraint()) { TPartOfChoppedConstraintNode::UniqueMerge(chopped, part->GetColumnMapping(ctx, prefix)); } if (const auto part = child->Tail().GetConstraint()) { TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(ctx, prefix)); } if (const auto part = child->Tail().GetConstraint()) { TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(ctx, prefix)); } } if (!sorted.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(sorted))); } if (!chopped.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(chopped))); } if (!uniques.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(uniques))); } if (!distincts.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(distincts))); } return TStatus::Ok; } template static void AddPartOf(const TExprNode::TPtr& input, TExprContext& ctx) { typename TPartOfConstraint::TMapType map; if (const auto part = input->Head().GetConstraint()) { map = part->GetColumnMapping(); } if (const auto part = input->Tail().GetConstraint()) { TPartOfConstraint::UniqueMerge(map, part->GetColumnMapping(input->Child(1)->Content())); } if (!map.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(map))); } } template static void ReplacePartOf(const TExprNode::TPtr& input, TExprContext& ctx) { typename TPartOfConstraint::TMapType sorted; const auto& name = input->Child(1)->Content(); if (const auto part = input->Head().GetConstraint()) { if (const auto filtered = part->FilterFields(ctx, [&name](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && path.front() != name; })) { sorted = filtered->GetColumnMapping(); } } if (const auto part = input->Tail().GetConstraint()) { TPartOfConstraint::UniqueMerge(sorted, part->GetColumnMapping(name)); } if (!sorted.empty()) { input->AddConstraint(ctx.MakeConstraint(std::move(sorted))); } } TStatus AddMemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { const auto& addStructNode = input->Head(); const auto& extraFieldNode = input->Tail(); if (const auto emptyConstraint = addStructNode.GetConstraint()) { input->AddConstraint(emptyConstraint); } AddPartOf(input, ctx); AddPartOf(input, ctx); AddPartOf(input, ctx); AddPartOf(input, ctx); TVector structConstraints; structConstraints.push_back(&addStructNode.GetConstraintSet()); if (const auto& valueNode = SkipModifiers(&extraFieldNode); TCoMember::Match(valueNode) || TCoNth::Match(valueNode)) { structConstraints.push_back(&valueNode->Head().GetConstraintSet()); } else if (valueNode->Type() == TExprNode::Argument) { structConstraints.push_back(&valueNode->GetConstraintSet()); } if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) { input->AddConstraint(varIndex); } return TStatus::Ok; } TStatus RemoveMemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { const auto& name = input->Tail().Content(); const auto filter = [&name](const TPartOfConstraintBase::TPathType& path) { return !path.empty() && path.front() != name; }; FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); FilterFromHead(input, filter, ctx); return FromFirst(input, output, ctx); } TStatus ReplaceMemberWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const { TVector structConstraints; structConstraints.push_back(&input->Head().GetConstraintSet()); ReplacePartOf(input, ctx); ReplacePartOf(input, ctx); ReplacePartOf(input, ctx); ReplacePartOf(input, ctx); if (const auto varIndex = TVarIndexConstraintNode::MakeCommon(structConstraints, ctx)) { input->AddConstraint(varIndex); } return TStatus::Ok; } TStatus ListWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { switch (input->ChildrenSize()) { case 1: return FromEmpty(input, output, ctx); case 2: return FromSecond(input, output, ctx); default: break; } return CommonFromChildren<1, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx); } template static TSmallVec GetConstraintsForInputArgument(const TExprNode& node, TExprContext& ctx) { TSmallVec argsConstraints(WideInput ? node.Child(1U)->Head().ChildrenSize() : 1U); if constexpr (WideInput) { if constexpr (Ordered) { if (const auto& mapping = TPartOfSortedConstraintNode::GetCommonMapping(node.Head().GetConstraint(), node.Head().GetConstraint()); !mapping.empty()) { for (ui32 i = 0U; i < argsConstraints.size(); ++i) { if (auto extracted = TPartOfSortedConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) { argsConstraints[i].emplace_back(ctx.MakeConstraint(std::move(extracted))); } } } if (const auto& mapping = TPartOfChoppedConstraintNode::GetCommonMapping(node.Head().GetConstraint(), node.Head().GetConstraint()); !mapping.empty()) { for (ui32 i = 0U; i < argsConstraints.size(); ++i) { if (auto extracted = TPartOfChoppedConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) { argsConstraints[i].emplace_back(ctx.MakeConstraint(std::move(extracted))); } } } } if constexpr (WithUnique) { if (const auto& mapping = TPartOfUniqueConstraintNode::GetCommonMapping(node.Head().GetConstraint(), node.Head().GetConstraint()); !mapping.empty()) { for (ui32 i = 0U; i < argsConstraints.size(); ++i) { if (auto extracted = TPartOfUniqueConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) { argsConstraints[i].emplace_back(ctx.MakeConstraint(std::move(extracted))); } } } if (const auto& mapping = TPartOfDistinctConstraintNode::GetCommonMapping(node.Head().GetConstraint(), node.Head().GetConstraint()); !mapping.empty()) { for (ui32 i = 0U; i < argsConstraints.size(); ++i) { if (auto extracted = TPartOfDistinctConstraintNode::ExtractField(mapping, ctx.GetIndexAsString(i)); !extracted.empty()) { argsConstraints[i].emplace_back(ctx.MakeConstraint(std::move(extracted))); } } } } } else { if (const auto inItemType = GetSeqItemType(node.Head().GetTypeAnn())) { if (inItemType->GetKind() == ETypeAnnotationKind::Variant) { if (inItemType->Cast()->GetUnderlyingType()->GetKind() == ETypeAnnotationKind::Tuple) { const auto tupleType = inItemType->Cast()->GetUnderlyingType()->Cast(); argsConstraints.front().push_back(ctx.MakeConstraint(*inItemType->Cast())); TMultiConstraintNode::TMapType multiItems; multiItems.reserve(tupleType->GetSize()); for (size_t i = 0; i < tupleType->GetSize(); ++i) { multiItems.emplace_back(i, TConstraintSet{}); if constexpr (WithUnique) { const auto inputMulti = node.Head().GetConstraint(); if (const auto inputConstr = inputMulti ? inputMulti->GetItem(i) : nullptr) { if (auto mapping = TPartOfUniqueConstraintNode::GetCommonMapping(inputConstr->GetConstraint(), inputConstr->GetConstraint()); !mapping.empty()) { multiItems.back().second.AddConstraint(ctx.MakeConstraint(std::move(mapping))); } if (auto mapping = TPartOfDistinctConstraintNode::GetCommonMapping(inputConstr->GetConstraint(), inputConstr->GetConstraint()); !mapping.empty()) { multiItems.back().second.AddConstraint(ctx.MakeConstraint(std::move(mapping))); } } } } if (!multiItems.empty()) { argsConstraints.front().emplace_back(ctx.MakeConstraint(std::move(multiItems))); } } } else { if constexpr (Ordered) { if (auto mapping = TPartOfSortedConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint()); !mapping.empty()) { argsConstraints.front().emplace_back(ctx.MakeConstraint(std::move(mapping))); } if (auto mapping = TPartOfChoppedConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint()); !mapping.empty()) { argsConstraints.front().emplace_back(ctx.MakeConstraint(std::move(mapping))); } } if constexpr (WithUnique) { if (auto mapping = TPartOfUniqueConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint()); !mapping.empty()) { argsConstraints.front().emplace_back(ctx.MakeConstraint(std::move(mapping))); } if (auto mapping = TPartOfDistinctConstraintNode::GetCommonMapping(GetDetailed(node.Head().GetConstraint(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint()); !mapping.empty()) { argsConstraints.front().emplace_back(ctx.MakeConstraint(std::move(mapping))); } } } } } return argsConstraints; } TStatus DictWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const { const std::vector