From 1aceb8a16aeec092357c6c181fd6b33056ef5fb5 Mon Sep 17 00:00:00 2001 From: Igor Munkin Date: Tue, 23 Apr 2024 17:06:22 +0500 Subject: [PATCH] Introduce helper to make block or scalar type (#3797) --- .../yql/core/type_ann/type_ann_blocks.cpp | 122 ++++-------------- 1 file changed, 28 insertions(+), 94 deletions(-) diff --git a/ydb/library/yql/core/type_ann/type_ann_blocks.cpp b/ydb/library/yql/core/type_ann/type_ann_blocks.cpp index 20cbd99ef31f..5d38e63b7aea 100644 --- a/ydb/library/yql/core/type_ann/type_ann_blocks.cpp +++ b/ydb/library/yql/core/type_ann/type_ann_blocks.cpp @@ -14,6 +14,18 @@ namespace NYql { namespace NTypeAnnImpl { +namespace { + +const TTypeAnnotationNode* MakeBlockOrScalarType(const TTypeAnnotationNode* blockItemType, bool isScalar, TExprContext& ctx) { + if (isScalar) { + return ctx.MakeType(blockItemType); + } else { + return ctx.MakeType(blockItemType); + } +} + +} // namespace + IGraphTransformer::TStatus AsScalarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) { Y_UNUSED(output); if (!EnsureArgsCount(*input, 1U, ctx.Expr)) { @@ -206,12 +218,7 @@ IGraphTransformer::TStatus BlockExistsWrapper(const TExprNode::TPtr& input, TExp } const TTypeAnnotationNode* resultType = ctx.Expr.MakeType(EDataSlot::Bool); - if (isScalar) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -280,12 +287,7 @@ IGraphTransformer::TStatus BlockCoalesceWrapper(const TExprNode::TPtr& input, TE return IGraphTransformer::TStatus::Error; } - auto outputItemType = secondItemType; - if (firstIsScalar && secondIsScalar) { - input->SetTypeAnn(ctx.Expr.MakeType(outputItemType)); - } else { - input->SetTypeAnn(ctx.Expr.MakeType(outputItemType)); - } + input->SetTypeAnn(MakeBlockOrScalarType(secondItemType, firstIsScalar && secondIsScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -325,12 +327,7 @@ IGraphTransformer::TStatus BlockLogicalWrapper(const TExprNode::TPtr& input, TEx resultType = ctx.Expr.MakeType(resultType); } - if (allScalars) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, allScalars, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -369,11 +366,7 @@ IGraphTransformer::TStatus BlockIfWrapper(const TExprNode::TPtr& input, TExprNod return IGraphTransformer::TStatus::Error; } - if (predIsScalar && thenIsScalar && elseIsScalar) { - input->SetTypeAnn(ctx.Expr.MakeType(thenItemType)); - } else { - input->SetTypeAnn(ctx.Expr.MakeType(thenItemType)); - } + input->SetTypeAnn(MakeBlockOrScalarType(thenItemType, predIsScalar && thenIsScalar && elseIsScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -391,12 +384,7 @@ IGraphTransformer::TStatus BlockJustWrapper(const TExprNode::TPtr& input, TExprN const TTypeAnnotationNode* blockItemType = GetBlockItemType(*child->GetTypeAnn(), isScalar); const TTypeAnnotationNode* resultType = ctx.Expr.MakeType(blockItemType); - if (isScalar) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -441,13 +429,7 @@ IGraphTransformer::TStatus BlockAsStructWrapper(const TExprNode::TPtr& input, TE return IGraphTransformer::TStatus::Repeat; } - const TTypeAnnotationNode* resultType; - if (onlyScalars) { - resultType = ctx.Expr.MakeType(structType); - } else { - resultType = ctx.Expr.MakeType(structType); - } - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(structType, onlyScalars, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -471,13 +453,7 @@ IGraphTransformer::TStatus BlockAsTupleWrapper(const TExprNode::TPtr& input, TEx } const TTypeAnnotationNode* resultType = ctx.Expr.MakeType(items); - if (onlyScalars) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, onlyScalars, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -537,13 +513,7 @@ IGraphTransformer::TStatus BlockMemberWrapper(const TExprNode::TPtr& input, TExp } } - if (isScalar) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -611,13 +581,7 @@ IGraphTransformer::TStatus BlockNthWrapper(const TExprNode::TPtr& input, TExprNo } } - if (isScalar) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -638,13 +602,7 @@ IGraphTransformer::TStatus BlockToPgWrapper(const TExprNode::TPtr& input, TExprN return IGraphTransformer::TStatus::Error; } - if (isScalar) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -665,13 +623,7 @@ IGraphTransformer::TStatus BlockFromPgWrapper(const TExprNode::TPtr& input, TExp return IGraphTransformer::TStatus::Error; } - if (isScalar) { - resultType = ctx.Expr.MakeType(resultType); - } else { - resultType = ctx.Expr.MakeType(resultType); - } - - input->SetTypeAnn(resultType); + input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -738,12 +690,7 @@ IGraphTransformer::TStatus BlockBitCastWrapper(const TExprNode::TPtr& input, TEx return IGraphTransformer::TStatus::Error; } - if (isScalar) { - input->SetTypeAnn(ctx.Expr.MakeType(outputType)); - } else { - input->SetTypeAnn(ctx.Expr.MakeType(outputType)); - } - + input->SetTypeAnn(MakeBlockOrScalarType(outputType, isScalar, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -1148,12 +1095,7 @@ IGraphTransformer::TStatus BlockPgOpWrapper(const TExprNode::TPtr& input, TExprN } auto result = ctx.Expr.MakeType(oper.ResultType); - if (allScalars) { - input->SetTypeAnn(ctx.Expr.MakeType(result)); - } else { - input->SetTypeAnn(ctx.Expr.MakeType(result)); - } - + input->SetTypeAnn(MakeBlockOrScalarType(result, allScalars, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -1240,12 +1182,7 @@ IGraphTransformer::TStatus BlockPgCallWrapper(const TExprNode::TPtr& input, TExp return IGraphTransformer::TStatus::Error; } - if (allScalars) { - input->SetTypeAnn(ctx.Expr.MakeType(result)); - } else { - input->SetTypeAnn(ctx.Expr.MakeType(result)); - } - + input->SetTypeAnn(MakeBlockOrScalarType(result, allScalars, ctx.Expr)); return IGraphTransformer::TStatus::Ok; } @@ -1285,11 +1222,8 @@ IGraphTransformer::TStatus BlockExtendWrapper(const TExprNode::TPtr& input, TExp TTypeAnnotationNode::TListType resultItemTypes; for (size_t i = 0; i < commonItemTypes.size(); ++i) { - if (i + 1 == commonItemTypes.size()) { - resultItemTypes.emplace_back(ctx.Expr.MakeType(commonItemTypes[i])); - } else { - resultItemTypes.emplace_back(ctx.Expr.MakeType(commonItemTypes[i])); - } + bool isScalar = (i + 1 == commonItemTypes.size()); + resultItemTypes.emplace_back(MakeBlockOrScalarType(commonItemTypes[i], isScalar, ctx.Expr)); } input->SetTypeAnn(ctx.Expr.MakeType(ctx.Expr.MakeType(std::move(resultItemTypes)))); return IGraphTransformer::TStatus::Ok;