Skip to content

Commit e597d2e

Browse files
authored
[YQL-17700][refactoring] Improve EnsureStructOrOptionalStructType() (#4416)
1 parent a88c67f commit e597d2e

File tree

4 files changed

+44
-63
lines changed

4 files changed

+44
-63
lines changed

ydb/library/yql/core/type_ann/type_ann_core.cpp

+28-47
Original file line numberDiff line numberDiff line change
@@ -1024,10 +1024,11 @@ namespace NTypeAnnImpl {
10241024
sourceNameNode = input->ChildPtr(2);
10251025
}
10261026

1027-
if (!EnsureStructOrOptionalStructType(*rowNode, ctx.Expr)) {
1027+
bool isOptionalStruct = false;
1028+
const TStructExprType* structType = nullptr;
1029+
if (!EnsureStructOrOptionalStructType(*rowNode, isOptionalStruct, structType, ctx.Expr)) {
10281030
return IGraphTransformer::TStatus::Error;
10291031
}
1030-
const TStructExprType* structType = RemoveAllOptionals(rowNode->GetTypeAnn())->Cast<TStructExprType>();
10311032

10321033
if (!EnsureAtom(*columnNameNode, ctx.Expr)) {
10331034
return IGraphTransformer::TStatus::Error;
@@ -1109,7 +1110,9 @@ namespace NTypeAnnImpl {
11091110
return IGraphTransformer::TStatus::Error;
11101111
}
11111112

1112-
if (!EnsureStructOrOptionalStructType(input->Head(), ctx.Expr)) {
1113+
bool isOptional = false;
1114+
const TStructExprType* structType = nullptr;
1115+
if (!EnsureStructOrOptionalStructType(input->Head(), isOptional, structType, ctx.Expr)) {
11131116
return IGraphTransformer::TStatus::Error;
11141117
}
11151118

@@ -1123,15 +1126,12 @@ namespace NTypeAnnImpl {
11231126
return IGraphTransformer::TStatus::Error;
11241127
}
11251128

1126-
bool isOptional = input->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional;
1127-
const TStructExprType& structType = *RemoveAllOptionals(input->Head().GetTypeAnn())->Cast<TStructExprType>();
1128-
1129-
auto pos = FindOrReportMissingMember(columnNameNode->Content(), input->Pos(), structType, ctx.Expr);
1129+
auto pos = FindOrReportMissingMember(columnNameNode->Content(), input->Pos(), *structType, ctx.Expr);
11301130
if (!pos) {
11311131
return IGraphTransformer::TStatus::Error;
11321132
}
11331133

1134-
const TTypeAnnotationNode* resultType = structType.GetItems()[*pos]->GetItemType();
1134+
const TTypeAnnotationNode* resultType = structType->GetItems()[*pos]->GetItemType();
11351135
if (isOptional && resultType->GetKind() != ETypeAnnotationKind::Optional && resultType->GetKind() != ETypeAnnotationKind::Null) {
11361136
resultType = ctx.Expr.MakeType<TOptionalExprType>(input->GetTypeAnn());
11371137
}
@@ -1150,7 +1150,9 @@ namespace NTypeAnnImpl {
11501150
return IGraphTransformer::TStatus::Repeat;
11511151
}
11521152

1153-
if (!EnsureStructOrOptionalStructType(input->Head(), ctx.Expr)) {
1153+
const TStructExprType* structType = nullptr;
1154+
bool isStructOptional = false;
1155+
if (!EnsureStructOrOptionalStructType(input->Head(), isStructOptional, structType, ctx.Expr)) {
11541156
return IGraphTransformer::TStatus::Error;
11551157
}
11561158

@@ -1162,15 +1164,6 @@ namespace NTypeAnnImpl {
11621164
return IGraphTransformer::TStatus::Error;
11631165
}
11641166

1165-
const TStructExprType* structType = nullptr;
1166-
bool isStructOptional = false;
1167-
if (input->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional) {
1168-
isStructOptional = true;
1169-
structType = input->Head().GetTypeAnn()->Cast<TOptionalExprType>()->GetItemType()->Cast<TStructExprType>();
1170-
} else {
1171-
structType = input->Head().GetTypeAnn()->Cast<TStructExprType>();
1172-
}
1173-
11741167
auto otherType = input->Child(2)->GetTypeAnn();
11751168
const bool isOptional = otherType->IsOptionalOrNull();
11761169
auto memberName = input->Child(1)->Content();
@@ -1231,16 +1224,13 @@ namespace NTypeAnnImpl {
12311224
}
12321225

12331226
auto structObj = child->Child(1);
1234-
if (!EnsureStructOrOptionalStructType(*structObj, ctx.Expr)) {
1227+
bool optional = false;
1228+
const TStructExprType* structType = nullptr;
1229+
if (!EnsureStructOrOptionalStructType(*structObj, optional, structType, ctx.Expr)) {
12351230
return IGraphTransformer::TStatus::Error;
12361231
}
12371232

1238-
auto type = structObj->GetTypeAnn();
1239-
const bool optional = type->GetKind() == ETypeAnnotationKind::Optional;
1240-
if (optional) {
1241-
type = type->Cast<TOptionalExprType>()->GetItemType();
1242-
}
1243-
for (auto& field: type->Cast<TStructExprType>()->GetItems()) {
1233+
for (auto& field: structType->GetItems()) {
12441234
auto itemType = field->GetItemType();
12451235
if (optional && !itemType->IsOptionalOrNull()) {
12461236
itemType = ctx.Expr.MakeType<TOptionalExprType>(itemType);
@@ -1351,7 +1341,9 @@ namespace NTypeAnnImpl {
13511341
}
13521342

13531343
auto structObj = input->Child(0);
1354-
if (!EnsureStructOrOptionalStructType(*structObj, ctx.Expr)) {
1344+
bool optional = false;
1345+
const TStructExprType* structExprType = nullptr;
1346+
if (!EnsureStructOrOptionalStructType(*structObj, optional, structExprType, ctx.Expr)) {
13551347
return IGraphTransformer::TStatus::Error;
13561348
}
13571349

@@ -1360,13 +1352,7 @@ namespace NTypeAnnImpl {
13601352
return status;
13611353
}
13621354

1363-
auto type = structObj->GetTypeAnn();
1364-
const bool optional = type->GetKind() == ETypeAnnotationKind::Optional;
1365-
if (optional) {
1366-
type = type->Cast<TOptionalExprType>()->GetItemType();
1367-
}
13681355
TVector<const TItemExprType*> allItems;
1369-
auto structExprType = type->Cast<TStructExprType>();
13701356
for (auto& field: structExprType->GetItems()) {
13711357
const auto& fieldName = field->GetName();
13721358
auto prefixes = input->Child(1);
@@ -1397,7 +1383,9 @@ namespace NTypeAnnImpl {
13971383
}
13981384

13991385
auto structObj = input->Child(0);
1400-
if (!EnsureStructOrOptionalStructType(*structObj, ctx.Expr)) {
1386+
bool optional = false;
1387+
const TStructExprType* structExprType = nullptr;
1388+
if (!EnsureStructOrOptionalStructType(*structObj, optional, structExprType, ctx.Expr)) {
14011389
return IGraphTransformer::TStatus::Error;
14021390
}
14031391

@@ -1406,13 +1394,7 @@ namespace NTypeAnnImpl {
14061394
return status;
14071395
}
14081396

1409-
auto type = structObj->GetTypeAnn();
1410-
const bool optional = type->GetKind() == ETypeAnnotationKind::Optional;
1411-
if (optional) {
1412-
type = type->Cast<TOptionalExprType>()->GetItemType();
1413-
}
14141397
TVector<const TItemExprType*> allItems;
1415-
auto structExprType = type->Cast<TStructExprType>();
14161398
for (auto& field: structExprType->GetItems()) {
14171399
const auto& fieldName = field->GetName();
14181400
auto prefixes = input->Child(1);
@@ -1460,7 +1442,9 @@ namespace NTypeAnnImpl {
14601442
}
14611443

14621444
auto structObj = *iter;
1463-
if (!EnsureStructOrOptionalStructType(*structObj, ctx.Expr)) {
1445+
bool isOptionalStruct = false;
1446+
const TStructExprType* structExprType = nullptr;
1447+
if (!EnsureStructOrOptionalStructType(*structObj, isOptionalStruct, structExprType, ctx.Expr)) {
14641448
return IGraphTransformer::TStatus::Error;
14651449
}
14661450
TSet<TString> aliases;
@@ -11877,16 +11861,13 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1187711861
return IGraphTransformer::TStatus::Error;
1187811862
}
1187911863

11880-
if (!EnsureStructOrOptionalStructType(input->Head(), ctx.Expr)) {
11864+
bool isOptional = false;
11865+
const TStructExprType* structType = nullptr;
11866+
if (!EnsureStructOrOptionalStructType(input->Head(), isOptional, structType, ctx.Expr)) {
1188111867
return IGraphTransformer::TStatus::Error;
1188211868
}
1188311869

11884-
const TTypeAnnotationNode* itemType = input->Head().GetTypeAnn();
11885-
if (itemType->GetKind() == ETypeAnnotationKind::Optional) {
11886-
itemType = itemType->Cast<TOptionalExprType>()->GetItemType();
11887-
}
11888-
11889-
for (const auto& x : itemType->Cast<TStructExprType>()->GetItems()) {
11870+
for (const auto& x : structType->GetItems()) {
1189011871
if (!x->GetItemType()->IsOptionalOrNull()) {
1189111872
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Head().Pos()), TStringBuilder() << "Expected all columns to be optional. Non optional column: " << x->GetName()));
1189211873
return IGraphTransformer::TStatus::Error;

ydb/library/yql/core/type_ann/type_ann_types.cpp

+3-10
Original file line numberDiff line numberDiff line change
@@ -1160,17 +1160,10 @@ namespace NTypeAnnImpl {
11601160
}
11611161

11621162
auto type = child->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
1163-
if (!EnsureStructOrOptionalStructType(child->Child(1)->Pos(), *type, ctx.Expr)) {
1164-
return IGraphTransformer::TStatus::Error;
1165-
}
1166-
1167-
const TStructExprType* structType;
1163+
const TStructExprType* structType = nullptr;
11681164
bool optional = false;
1169-
if (type->GetKind() == ETypeAnnotationKind::Optional) {
1170-
optional = true;
1171-
structType = type->Cast<TOptionalExprType>()->GetItemType()->Cast<TStructExprType>();
1172-
} else {
1173-
structType = type->Cast<TStructExprType>();
1165+
if (!EnsureStructOrOptionalStructType(child->Child(1)->Pos(), *type, optional, structType, ctx.Expr)) {
1166+
return IGraphTransformer::TStatus::Error;
11741167
}
11751168

11761169
for (auto& field : structType->GetItems()) {

ydb/library/yql/core/yql_expr_type_annotation.cpp

+11-4
Original file line numberDiff line numberDiff line change
@@ -3981,7 +3981,7 @@ bool EnsureAnySeqType(TPositionHandle position, const TTypeAnnotationNode& type,
39813981
return false;
39823982
}
39833983

3984-
bool EnsureStructOrOptionalStructType(const TExprNode& node, TExprContext& ctx) {
3984+
bool EnsureStructOrOptionalStructType(const TExprNode& node, bool& isOptional, const TStructExprType*& structType, TExprContext& ctx) {
39853985
if (HasError(node.GetTypeAnn(), ctx)) {
39863986
return false;
39873987
}
@@ -3992,10 +3992,11 @@ bool EnsureStructOrOptionalStructType(const TExprNode& node, TExprContext& ctx)
39923992
return false;
39933993
}
39943994

3995-
return EnsureStructOrOptionalStructType(node.Pos(), *node.GetTypeAnn(), ctx);
3995+
return EnsureStructOrOptionalStructType(node.Pos(), *node.GetTypeAnn(), isOptional, structType, ctx);
39963996
}
3997-
3998-
bool EnsureStructOrOptionalStructType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx) {
3997+
bool EnsureStructOrOptionalStructType(TPositionHandle position, const TTypeAnnotationNode& type, bool& isOptional,
3998+
const TStructExprType*& structType, TExprContext& ctx)
3999+
{
39994000
if (HasError(&type, ctx)) {
40004001
return false;
40014002
}
@@ -4005,13 +4006,19 @@ bool EnsureStructOrOptionalStructType(TPositionHandle position, const TTypeAnnot
40054006
ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected either struct or optional of struct, but got: " << type));
40064007
return false;
40074008
}
4009+
40084010
if (kind == ETypeAnnotationKind::Optional) {
40094011
auto itemType = type.Cast<TOptionalExprType>()->GetItemType();
40104012
kind = itemType->GetKind();
40114013
if (kind != ETypeAnnotationKind::Struct) {
40124014
ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected either struct or optional of struct, but got: " << type));
40134015
return false;
40144016
}
4017+
isOptional = true;
4018+
structType = itemType->Cast<TStructExprType>();
4019+
} else {
4020+
isOptional = false;
4021+
structType = type.Cast<TStructExprType>();
40154022
}
40164023

40174024
return true;

ydb/library/yql/core/yql_expr_type_annotation.h

+2-2
Original file line numberDiff line numberDiff line change
@@ -173,8 +173,8 @@ bool EnsureInspectable(const TExprNode& node, TExprContext& ctx);
173173
bool EnsureInspectableType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
174174
bool EnsureListOrOptionalType(const TExprNode& node, TExprContext& ctx);
175175
bool EnsureListOrOptionalListType(const TExprNode& node, TExprContext& ctx);
176-
bool EnsureStructOrOptionalStructType(const TExprNode& node, TExprContext& ctx);
177-
bool EnsureStructOrOptionalStructType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
176+
bool EnsureStructOrOptionalStructType(const TExprNode& node, bool& isOptional, const TStructExprType*& structType, TExprContext& ctx);
177+
bool EnsureStructOrOptionalStructType(TPositionHandle position, const TTypeAnnotationNode& type, bool& isOptional, const TStructExprType*& structType, TExprContext& ctx);
178178
bool EnsureSeqType(const TExprNode& node, TExprContext& ctx, bool* isStream = nullptr);
179179
bool EnsureSeqType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx, bool* isStream = nullptr);
180180
bool EnsureSeqOrOptionalType(const TExprNode& node, TExprContext& ctx);

0 commit comments

Comments
 (0)