Skip to content

Commit

Permalink
Introduce helper to make block or scalar type (#3797)
Browse files Browse the repository at this point in the history
  • Loading branch information
@igormunkin
igormunkin authored Apr 23, 2024
1 parent 6befca3 commit 1aceb8a
Showing 1 changed file with 28 additions and 94 deletions.
122 changes: 28 additions & 94 deletions ydb/library/yql/core/type_ann/type_ann_blocks.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,18 @@
namespace NYql {
namespace NTypeAnnImpl {

namespace {

const TTypeAnnotationNode* MakeBlockOrScalarType(const TTypeAnnotationNode* blockItemType, bool isScalar, TExprContext& ctx) {
if (isScalar) {
return ctx.MakeType<TScalarExprType>(blockItemType);
} else {
return ctx.MakeType<TBlockExprType>(blockItemType);
}
}

} // namespace

IGraphTransformer::TStatus AsScalarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
Y_UNUSED(output);
if (!EnsureArgsCount(*input, 1U, ctx.Expr)) {
Expand Down Expand Up @@ -206,12 +218,7 @@ IGraphTransformer::TStatus BlockExistsWrapper(const TExprNode::TPtr& input, TExp
}

const TTypeAnnotationNode* resultType = ctx.Expr.MakeType<TDataExprType>(EDataSlot::Bool);
if (isScalar) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}
input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -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<TScalarExprType>(outputItemType));
} else {
input->SetTypeAnn(ctx.Expr.MakeType<TBlockExprType>(outputItemType));
}
input->SetTypeAnn(MakeBlockOrScalarType(secondItemType, firstIsScalar && secondIsScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -325,12 +327,7 @@ IGraphTransformer::TStatus BlockLogicalWrapper(const TExprNode::TPtr& input, TEx
resultType = ctx.Expr.MakeType<TOptionalExprType>(resultType);
}

if (allScalars) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}
input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, allScalars, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -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<TScalarExprType>(thenItemType));
} else {
input->SetTypeAnn(ctx.Expr.MakeType<TBlockExprType>(thenItemType));
}
input->SetTypeAnn(MakeBlockOrScalarType(thenItemType, predIsScalar && thenIsScalar && elseIsScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand All @@ -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<TOptionalExprType>(blockItemType);

if (isScalar) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}
input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -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<TScalarExprType>(structType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(structType);
}
input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(structType, onlyScalars, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand All @@ -471,13 +453,7 @@ IGraphTransformer::TStatus BlockAsTupleWrapper(const TExprNode::TPtr& input, TEx
}

const TTypeAnnotationNode* resultType = ctx.Expr.MakeType<TTupleExprType>(items);
if (onlyScalars) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}

input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, onlyScalars, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -537,13 +513,7 @@ IGraphTransformer::TStatus BlockMemberWrapper(const TExprNode::TPtr& input, TExp
}
}

if (isScalar) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}

input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -611,13 +581,7 @@ IGraphTransformer::TStatus BlockNthWrapper(const TExprNode::TPtr& input, TExprNo
}
}

if (isScalar) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}

input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand All @@ -638,13 +602,7 @@ IGraphTransformer::TStatus BlockToPgWrapper(const TExprNode::TPtr& input, TExprN
return IGraphTransformer::TStatus::Error;
}

if (isScalar) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}

input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand All @@ -665,13 +623,7 @@ IGraphTransformer::TStatus BlockFromPgWrapper(const TExprNode::TPtr& input, TExp
return IGraphTransformer::TStatus::Error;
}

if (isScalar) {
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
} else {
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
}

input->SetTypeAnn(resultType);
input->SetTypeAnn(MakeBlockOrScalarType(resultType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -738,12 +690,7 @@ IGraphTransformer::TStatus BlockBitCastWrapper(const TExprNode::TPtr& input, TEx
return IGraphTransformer::TStatus::Error;
}

if (isScalar) {
input->SetTypeAnn(ctx.Expr.MakeType<TScalarExprType>(outputType));
} else {
input->SetTypeAnn(ctx.Expr.MakeType<TBlockExprType>(outputType));
}

input->SetTypeAnn(MakeBlockOrScalarType(outputType, isScalar, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -1148,12 +1095,7 @@ IGraphTransformer::TStatus BlockPgOpWrapper(const TExprNode::TPtr& input, TExprN
}

auto result = ctx.Expr.MakeType<TPgExprType>(oper.ResultType);
if (allScalars) {
input->SetTypeAnn(ctx.Expr.MakeType<TScalarExprType>(result));
} else {
input->SetTypeAnn(ctx.Expr.MakeType<TBlockExprType>(result));
}

input->SetTypeAnn(MakeBlockOrScalarType(result, allScalars, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -1240,12 +1182,7 @@ IGraphTransformer::TStatus BlockPgCallWrapper(const TExprNode::TPtr& input, TExp
return IGraphTransformer::TStatus::Error;
}

if (allScalars) {
input->SetTypeAnn(ctx.Expr.MakeType<TScalarExprType>(result));
} else {
input->SetTypeAnn(ctx.Expr.MakeType<TBlockExprType>(result));
}

input->SetTypeAnn(MakeBlockOrScalarType(result, allScalars, ctx.Expr));
return IGraphTransformer::TStatus::Ok;
}

Expand Down Expand Up @@ -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<TScalarExprType>(commonItemTypes[i]));
} else {
resultItemTypes.emplace_back(ctx.Expr.MakeType<TBlockExprType>(commonItemTypes[i]));
}
bool isScalar = (i + 1 == commonItemTypes.size());
resultItemTypes.emplace_back(MakeBlockOrScalarType(commonItemTypes[i], isScalar, ctx.Expr));
}
input->SetTypeAnn(ctx.Expr.MakeType<TFlowExprType>(ctx.Expr.MakeType<TMultiExprType>(std::move(resultItemTypes))));
return IGraphTransformer::TStatus::Ok;
Expand Down

0 comments on commit 1aceb8a

Please sign in to comment.