aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authoraneporada <aneporada@yandex-team.ru>2022-02-20 12:24:11 +0300
committeraneporada <aneporada@yandex-team.ru>2022-02-20 12:24:11 +0300
commitbab22eebdf894b271f661f2474068552165277b4 (patch)
tree36cd10cfbc9b6a8652265564dc3aa7a91808da8b
parent8d89aca1fc4080024ced45c8d4e85c58a8a2df47 (diff)
downloadydb-bab22eebdf894b271f661f2474068552165277b4.tar.gz
[YQL-14389] Add MapNext optimizers
ref:2c84d356de4b420dea6fd5cf2b62c2c6ab307e70
-rw-r--r--ydb/library/yql/core/common_opt/yql_co_extr_members.cpp27
-rw-r--r--ydb/library/yql/core/common_opt/yql_co_extr_members.h1
-rw-r--r--ydb/library/yql/core/common_opt/yql_co_finalizers.cpp8
-rw-r--r--ydb/library/yql/core/common_opt/yql_co_flow2.cpp38
-rw-r--r--ydb/library/yql/core/common_opt/yql_co_simple1.cpp19
-rw-r--r--ydb/library/yql/core/expr_nodes/yql_expr_nodes.json5
-rw-r--r--ydb/library/yql/core/yql_opt_utils.cpp5
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;
}