Skip to content

Commit a7816a0

Browse files
pavelvelikhovPavel Velikhov
pavelvelikhov
and
Pavel Velikhov
authored
Added support for CTEs, alias renames and graphs into new RBO (#24467)
Co-authored-by: Pavel Velikhov <pavelvelikhov@localhost.localdomain>
1 parent 36b11b6 commit a7816a0

15 files changed

+1650
-520
lines changed

‎ydb/core/kqp/expr_nodes/kqp_expr_nodes.json‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -776,26 +776,51 @@
776776
"Match": {"Type": "Callable", "Name": "KqpOpEmptySource"}
777777
},
778778
{
779-
"Name": "TKqpOpMapElement",
779+
"Name": "TKqpOpMapElementBase",
780780
"Base": "TCallable",
781-
"Match": {"Type": "Callable", "Name": "KqpOpMapElement"},
781+
"Match": {"Type": "CallableBase"},
782782
"Children": [
783783
{"Index": 0, "Name": "Input", "Type": "TExprBase"},
784-
{"Index": 1, "Name": "Variable", "Type": "TCoAtom"},
784+
{"Index": 1, "Name": "Variable", "Type": "TCoAtom"}
785+
]
786+
},
787+
{
788+
"Name": "TKqpOpMapElementRename",
789+
"Base": "TKqpOpMapElementBase",
790+
"Match": {"Type": "Callable", "Name": "KqpOpMapElementRename"},
791+
"Children": [
792+
{"Index": 2, "Name": "From", "Type": "TCoAtom"}
793+
]
794+
},
795+
{
796+
"Name": "TKqpOpMapElementLambda",
797+
"Base": "TKqpOpMapElementBase",
798+
"Match": {"Type": "Callable", "Name": "KqpOpMapElementLambda"},
799+
"Children": [
785800
{"Index": 2, "Name": "Lambda", "Type": "TCoLambda"}
786801
]
787802
},
788803
{
789804
"Name": "TKqpOpElementList",
790-
"ListBase" : "TKqpOpMapElement"
805+
"ListBase" : "TKqpOpMapElementBase"
791806
},
792807
{
793808
"Name": "TKqpOpMap",
794809
"Base": "TCallable",
795810
"Match": {"Type": "Callable", "Name": "KqpOpMap"},
796811
"Children": [
797812
{"Index": 0, "Name": "Input", "Type": "TExprBase"},
798-
{"Index": 1, "Name": "MapElements", "Type": "TKqpOpElementList"}
813+
{"Index": 1, "Name": "MapElements", "Type": "TKqpOpElementList"},
814+
{"Index": 2, "Name": "Project", "Type": "TCoAtom", "Optional" : true}
815+
]
816+
},
817+
{
818+
"Name": "TKqpOpProject",
819+
"Base": "TCallable",
820+
"Match": {"Type": "Callable", "Name": "KqpOpProject"},
821+
"Children": [
822+
{"Index": 0, "Name": "Input", "Type": "TExprBase"},
823+
{"Index": 1, "Name": "ProjectList", "Type": "TCoAtomList"}
799824
]
800825
},
801826
{

‎ydb/core/kqp/opt/kqp_type_ann.cpp‎

Lines changed: 78 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2229,11 +2229,11 @@ TStatus AnnotateOpEmptySource(const TExprNode::TPtr& input, TExprContext& ctx) {
22292229
return TStatus::Ok;
22302230
}
22312231

2232-
TStatus AnnotateOpMapElement(const TExprNode::TPtr& input, TExprContext& ctx) {
2233-
const TTypeAnnotationNode* inputType = input->ChildPtr(TKqpOpMapElement::idx_Input)->GetTypeAnn();
2232+
TStatus AnnotateOpMapElementLambda(const TExprNode::TPtr& input, TExprContext& ctx) {
2233+
const TTypeAnnotationNode* inputType = input->ChildPtr(TKqpOpMapElementLambda::idx_Input)->GetTypeAnn();
22342234
const TTypeAnnotationNode* itemType = inputType->Cast<TListExprType>()->GetItemType();
22352235

2236-
auto& lambda = input->ChildRef(TKqpOpMapElement::idx_Lambda);
2236+
auto& lambda = input->ChildRef(TKqpOpMapElementLambda::idx_Lambda);
22372237
if (!UpdateLambdaAllArgumentsTypes(lambda, {itemType}, ctx)) {
22382238
return IGraphTransformer::TStatus::Error;
22392239
}
@@ -2243,44 +2243,92 @@ TStatus AnnotateOpMapElement(const TExprNode::TPtr& input, TExprContext& ctx) {
22432243
return IGraphTransformer::TStatus::Repeat;
22442244
}
22452245

2246-
input->SetTypeAnn(lambdaType);
2246+
auto variable = input->ChildRef(TKqpOpMapElementLambda::idx_Variable);
2247+
auto res = ctx.MakeType<TItemExprType>(variable->Content(), lambdaType);
2248+
2249+
input->SetTypeAnn(res);
22472250
return TStatus::Ok;
22482251
}
22492252

2250-
TStatus AnnotateOpMap(const TExprNode::TPtr& input, TExprContext& ctx, TTypeAnnotationContext& typesCtx) {
2253+
TStatus AnnotateOpMapElementRename(const TExprNode::TPtr& input, TExprContext& ctx) {
2254+
Y_UNUSED(ctx);
22512255

2252-
const TTypeAnnotationNode* inputType = input->ChildPtr(TKqpOpMap::idx_Input)->GetTypeAnn();
2253-
//YQL_CLOG(TRACE, CoreDq) << "Annotating OpMap, input type:" << *inputType;
2256+
const TTypeAnnotationNode* inputType = input->ChildPtr(TKqpOpMapElementLambda::idx_Input)->GetTypeAnn();
2257+
auto structType = inputType->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
2258+
auto typeItems = structType->GetItems();
22542259

2255-
const TTypeAnnotationNode* itemType = inputType->Cast<TListExprType>()->GetItemType();
2256-
//YQL_CLOG(TRACE, CoreDq) << "item type:" << *itemType;
2260+
auto from = input->ChildRef(TKqpOpMapElementRename::idx_From);
2261+
auto typeIt = std::find_if(typeItems.begin(), typeItems.end(), [&from](const TItemExprType* t){
2262+
return from->Content() == t->GetName();
2263+
});
2264+
2265+
if (typeIt==typeItems.end()) {
2266+
YQL_CLOG(TRACE, CoreDq) << "Trying to find " << from->Content() << " in " << *(TTypeAnnotationNode*)structType;
2267+
}
2268+
2269+
Y_ENSURE(typeIt!=typeItems.end());
2270+
2271+
auto variable = input->ChildRef(TKqpOpMapElementRename::idx_Variable);
2272+
auto res = ctx.MakeType<TItemExprType>(variable->Content(), (*typeIt)->GetItemType());
2273+
2274+
input->SetTypeAnn(res);
2275+
return TStatus::Ok;
2276+
}
2277+
2278+
TStatus AnnotateOpMap(const TExprNode::TPtr& input, TExprContext& ctx, TTypeAnnotationContext& typesCtx) {
2279+
Y_UNUSED(typesCtx);
22572280

22582281
TVector<const TItemExprType*> structItemTypes;
22592282

2283+
if (input->ChildrenSize() <= TKqpOpMap::idx_Project) {
2284+
const TTypeAnnotationNode* inputType = input->ChildPtr(TKqpOpMap::idx_Input)->GetTypeAnn();
2285+
auto structType = inputType->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
2286+
2287+
for (auto t : structType->GetItems()) {
2288+
structItemTypes.push_back(t);
2289+
}
2290+
}
2291+
22602292
for (size_t idx = 0; idx < input->ChildPtr(TKqpOpMap::idx_MapElements)->ChildrenSize(); idx++) {
22612293
auto& element = input->ChildPtr(TKqpOpMap::idx_MapElements)->ChildRef(idx);
2262-
auto variable = element->ChildPtr(TKqpOpMapElement::idx_Variable);
2263-
auto& lambda = element->ChildRef(TKqpOpMapElement::idx_Lambda);
2294+
auto type = element->GetTypeAnn();
2295+
structItemTypes.push_back((TItemExprType*)type);
2296+
}
22642297

2265-
if (!UpdateLambdaAllArgumentsTypes(lambda, {itemType}, ctx)) {
2266-
return IGraphTransformer::TStatus::Error;
2267-
}
2298+
auto resultItemType = ctx.MakeType<TStructExprType>(structItemTypes);
2299+
const TTypeAnnotationNode* resultAnn = ctx.MakeType<TListExprType>(resultItemType);
22682300

2269-
auto lambdaType = lambda->GetTypeAnn();
2270-
if (!lambdaType) {
2271-
return IGraphTransformer::TStatus::Repeat;
2301+
input->SetTypeAnn(resultAnn);
2302+
//typesCtx.ExpectedTypes[input->UniqueId()] = resultAnn;
2303+
2304+
YQL_CLOG(TRACE, CoreDq) << "Type annotation for OpMap done: " << *resultAnn;
2305+
2306+
return TStatus::Ok;
2307+
}
2308+
2309+
TStatus AnnotateOpProject(const TExprNode::TPtr& input, TExprContext& ctx) {
2310+
auto structType = input->ChildPtr(TKqpOpProject::idx_Input)->GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
2311+
TVector<const TItemExprType*> structItemTypes;
2312+
auto typeItems = structType->GetItems();
2313+
2314+
for (size_t i=0; i<input->ChildPtr(TKqpOpProject::idx_ProjectList)->ChildrenSize(); i++) {
2315+
auto proj = input->ChildPtr(TKqpOpProject::idx_ProjectList)->ChildRef(i);
2316+
auto typeItemIt = std::find_if(typeItems.begin(), typeItems.end(), [&proj](const TItemExprType* t){
2317+
return proj->Content() == t->GetName();
2318+
});
2319+
if (typeItemIt == typeItems.end()) {
2320+
continue;
22722321
}
22732322

2274-
structItemTypes.push_back(ctx.MakeType<TItemExprType>(variable->Content(),lambdaType));
2323+
structItemTypes.push_back(*typeItemIt);
22752324
}
22762325

22772326
auto resultItemType = ctx.MakeType<TStructExprType>(structItemTypes);
22782327
const TTypeAnnotationNode* resultAnn = ctx.MakeType<TListExprType>(resultItemType);
22792328

22802329
input->SetTypeAnn(resultAnn);
2281-
typesCtx.ExpectedTypes[input->UniqueId()] = resultAnn;
22822330

2283-
YQL_CLOG(TRACE, CoreDq) << "Type annotation for OpMap done: " << *resultAnn;
2331+
YQL_CLOG(TRACE, CoreDq) << "Type annotation for OpProject done: " << *resultAnn;
22842332

22852333
return TStatus::Ok;
22862334
}
@@ -2536,14 +2584,22 @@ TAutoPtr<IGraphTransformer> CreateKqpTypeAnnotationTransformer(const TString& cl
25362584
return AnnotateOpEmptySource(input, ctx);
25372585
}
25382586

2539-
if (TKqpOpMapElement::Match(input.Get())) {
2540-
return AnnotateOpMapElement(input, ctx);
2587+
if (TKqpOpMapElementLambda::Match(input.Get())) {
2588+
return AnnotateOpMapElementLambda(input, ctx);
2589+
}
2590+
2591+
if (TKqpOpMapElementRename::Match(input.Get())) {
2592+
return AnnotateOpMapElementRename(input, ctx);
25412593
}
25422594

25432595
if (TKqpOpMap::Match(input.Get())) {
25442596
return AnnotateOpMap(input, ctx, typesCtx);
25452597
}
25462598

2599+
if (TKqpOpProject::Match(input.Get())) {
2600+
return AnnotateOpProject(input, ctx);
2601+
}
2602+
25472603
if (TKqpOpFilter::Match(input.Get())) {
25482604
return AnnotateOpFilter(input, ctx);
25492605
}

‎ydb/core/kqp/opt/rbo/kqp_convert_to_physical.cpp‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -579,10 +579,28 @@ TExprNode::TPtr ConvertToPhysical(TOpRoot & root, TExprContext& ctx, TTypeAnnot
579579
auto arg = Build<TCoArgument>(ctx, root.Node->Pos()).Name("arg").Done().Ptr();
580580

581581
for (auto mapElement : map.MapElements()) {
582+
TMaybeNode<TCoLambda> mapLambda;
583+
584+
if (mapElement.Maybe<TKqpOpMapElementLambda>()) {
585+
mapLambda = mapElement.Cast<TKqpOpMapElementLambda>().Lambda();
586+
}
587+
else {
588+
auto var = mapElement.Cast<TKqpOpMapElementRename>().From();
589+
// clang-format off
590+
mapLambda = Build<TCoLambda>(ctx, root.Node->Pos())
591+
.Args({arg})
592+
.Body<TCoMember>()
593+
.Struct(arg)
594+
.Name(var)
595+
.Build()
596+
.Done();
597+
// clang-format on
598+
}
599+
582600
// clang-format off
583601
auto tuple = Build<TCoNameValueTuple>(ctx, root.Node->Pos())
584602
.Name().Build(mapElement.Variable())
585-
.Value(mapElement.Lambda().Body())
603+
.Value(mapLambda.Body())
586604
.Done();
587605
// clang-format on
588606

0 commit comments

Comments
 (0)