diff options
author | aneporada <aneporada@yandex-team.ru> | 2022-02-20 12:24:11 +0300 |
---|---|---|
committer | aneporada <aneporada@yandex-team.ru> | 2022-02-20 12:24:11 +0300 |
commit | bab22eebdf894b271f661f2474068552165277b4 (patch) | |
tree | 36cd10cfbc9b6a8652265564dc3aa7a91808da8b | |
parent | 8d89aca1fc4080024ced45c8d4e85c58a8a2df47 (diff) | |
download | ydb-bab22eebdf894b271f661f2474068552165277b4.tar.gz |
[YQL-14389] Add MapNext optimizers
ref:2c84d356de4b420dea6fd5cf2b62c2c6ab307e70
7 files changed, 101 insertions, 2 deletions
diff --git a/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp b/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp index 215a1e9f69..35b24e9439 100644 --- a/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp +++ b/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp @@ -4,6 +4,7 @@ #include <ydb/library/yql/core/yql_join.h> #include <ydb/library/yql/core/yql_opt_utils.h> #include <ydb/library/yql/core/yql_opt_window.h> +#include <ydb/library/yql/core/yql_type_helpers.h> #include <ydb/library/yql/utils/log/log.h> @@ -794,4 +795,30 @@ TExprNode::TPtr ApplyExtractMembersToCollect(const TExprNode::TPtr& node, const .Done().Ptr(); } +TExprNode::TPtr ApplyExtractMembersToMapNext(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) { + TCoMapNext mapNext(node); + YQL_CLOG(DEBUG, Core) << "Apply ExtractMembers to " << node->Content() << logSuffix; + + const bool singleValue = GetItemType(*mapNext.Lambda().Ref().GetTypeAnn()) == nullptr; + TSet<TStringBuf> memberNames; + for (auto& member : members->ChildrenList()) { + YQL_ENSURE(member->IsAtom()); + memberNames.insert(member->Content()); + } + TExprBase newBody{FilterByFields(mapNext.Pos(), mapNext.Lambda().Body().Ptr(), memberNames, ctx, singleValue)}; + + return Build<TCoMapNext>(ctx, mapNext.Pos()) + .Input(mapNext.Input()) + .Lambda() + .Args({"current", "next"}) + .Body<TExprApplier>() + .Apply(newBody) + .With(mapNext.Lambda().Args().Arg(0), "current") + .With(mapNext.Lambda().Args().Arg(1), "next") + .Build() + .Build() + .Done() + .Ptr(); +} + } // NYql diff --git a/ydb/library/yql/core/common_opt/yql_co_extr_members.h b/ydb/library/yql/core/common_opt/yql_co_extr_members.h index 8244636784..974775b49d 100644 --- a/ydb/library/yql/core/common_opt/yql_co_extr_members.h +++ b/ydb/library/yql/core/common_opt/yql_co_extr_members.h @@ -25,5 +25,6 @@ TExprNode::TPtr ApplyExtractMembersToAggregate(const TExprNode::TPtr& node, cons TExprNode::TPtr ApplyExtractMembersToChopper(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix); TExprNode::TPtr ApplyExtractMembersToCollect(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix); TExprNode::TPtr ApplyExtractMembersToMapJoinCore(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix); +TExprNode::TPtr ApplyExtractMembersToMapNext(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix); } // NYql diff --git a/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp b/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp index b77766bc6f..9a72235ea8 100644 --- a/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp +++ b/ydb/library/yql/core/common_opt/yql_co_finalizers.cpp @@ -203,6 +203,14 @@ void RegisterCoFinalizers(TFinalizingOptimizerMap& map) { } ); }; + + map[TCoMapNext::CallableName()] = [](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) { + SubsetFieldsForNodeWithMultiUsage(node, *optCtx.ParentsMap, toOptimize, ctx, + [] (const TExprNode::TPtr& input, const TExprNode::TPtr& members, const TParentsMap&, TExprContext& ctx) { + return ApplyExtractMembersToMapNext(input, members, ctx, " with multi-usage"); + } + ); + }; } } // NYql diff --git a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp index 88f9140031..690393bbfe 100644 --- a/ydb/library/yql/core/common_opt/yql_co_flow2.cpp +++ b/ydb/library/yql/core/common_opt/yql_co_flow2.cpp @@ -1937,6 +1937,13 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) { return node; } + if (self.Input().Maybe<TCoMapNext>()) { + if (auto res = ApplyExtractMembersToMapNext(self.Input().Ptr(), self.Members().Ptr(), ctx, {})) { + return res; + } + return node; + } + return node; }; @@ -2302,6 +2309,37 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) { } return node; }; + + map[TCoMapNext::CallableName()] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) { + TCoMapNext self(node); + if (!optCtx.IsSingleUsage(self.Input().Ref())) { + return node; + } + + std::map<std::string_view, TExprNode::TPtr> usedFields; + if (( + HaveFieldsSubset(self.Lambda().Body().Ptr(), self.Lambda().Args().Arg(0).Ref(), usedFields, *optCtx.ParentsMap, false) && + HaveFieldsSubset(self.Lambda().Body().Ptr(), self.Lambda().Args().Arg(1).Ref(), usedFields, *optCtx.ParentsMap, false) + ) && usedFields.size() < GetSeqItemType(self.Input().Ref().GetTypeAnn())->Cast<TStructExprType>()->GetSize()) + { + TExprNode::TListType fields; + fields.reserve(usedFields.size()); + std::transform(usedFields.begin(), usedFields.end(), std::back_inserter(fields), + [](std::pair<const std::string_view, TExprNode::TPtr>& item){ return std::move(item.second); }); + + YQL_CLOG(DEBUG, Core) << node->Content() << "SubsetFields"; + return Build<TCoMapNext>(ctx, node->Pos()) + .Input<TCoExtractMembers>() + .Input(self.Input()) + .Members() + .Add(std::move(fields)) + .Build() + .Build() + .Lambda(ctx.DeepCopyLambda(self.Lambda().Ref())) + .Done().Ptr(); + } + return node; + }; } } diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp index 8a66d1bd60..11899138d8 100644 --- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp +++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp @@ -7194,6 +7194,25 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) { return ctx.NewCallable(node->Pos(), "Member", { node->HeadPtr(), node->ChildPtr(1) }); }; + map["MapNext"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& /*optCtx*/) { + TCoMapNext self(node); + if (!IsDepended(self.Lambda().Body().Ref(), self.Lambda().Args().Arg(1).Ref())) { + YQL_CLOG(DEBUG, Core) << node->Content() << " with unused next arg"; + return Build<TCoOrderedMap>(ctx, self.Pos()) + .Input(self.Input()) + .Lambda() + .Args({"row"}) + .Body<TExprApplier>() + .Apply(self.Lambda().Body()) + .With(self.Lambda().Args().Arg(0), "row") + .Build() + .Build() + .Done() + .Ptr(); + } + return node; + }; + // will be applied to any callable after all above map[""] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) { YQL_ENSURE(node->IsCallable()); diff --git a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json index edfcaa1530..96ed88efbb 100644 --- a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json +++ b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json @@ -183,6 +183,11 @@ "Match": {"Type": "Callable", "Name": "OrderedMap"} }, { + "Name": "TCoMapNext", + "Base": "TCoMapBase", + "Match": {"Type": "Callable", "Name": "MapNext"} + }, + { "Name": "TCoMultiMap", "Base": "TCoMapBase", "Match": {"Type": "Callable", "Name": "MultiMap"} diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp index 69aa84a500..8385c239ae 100644 --- a/ydb/library/yql/core/yql_opt_utils.cpp +++ b/ydb/library/yql/core/yql_opt_utils.cpp @@ -207,7 +207,8 @@ TExprNode::TPtr KeepColumnOrder(const TExprNode::TPtr& node, const TExprNode& sr template<class TFieldsSet> bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TFieldsSet& usedFields, const TParentsMap& parentsMap, bool allowDependsOn) { - if (arg.GetTypeAnn()->GetKind() != ETypeAnnotationKind::Struct) { + const TTypeAnnotationNode* argType = RemoveOptionalType(arg.GetTypeAnn()); + if (argType->GetKind() != ETypeAnnotationKind::Struct) { return false; } @@ -215,7 +216,7 @@ bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TField return false; } - const auto inputStructType = arg.GetTypeAnn()->Cast<TStructExprType>(); + const auto inputStructType = argType->Cast<TStructExprType>(); if (!IsDepended(*start, arg)) { return inputStructType->GetSize() > 0; } |