Fix type annotation for Block{Combine,MergeFinalize,MergeManyFinalize}Hashed

4 files changed, 12 insertions, 11 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 6b8df2433d..9db337ec7f 100644
--- a/ydb/library/yql/core/type_ann/type_ann_blocks.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_blocks.cpp
@@ -264,10 +264,8 @@ IGraphTransformer::TStatus BlockBitCastWrapper(const TExprNode::TPtr& input, TEx
 
     if (isScalar) {
         input->SetTypeAnn(ctx.Expr.MakeType<TScalarExprType>(outputType));
-    } else if (outputType->HasFixedSizeRepr()) {
-        input->SetTypeAnn(ctx.Expr.MakeType<TBlockExprType>(outputType));
     } else {
-        input->SetTypeAnn(ctx.Expr.MakeType<TChunkedBlockExprType>(outputType));
+        input->SetTypeAnn(MakeBlockType(*outputType, ctx.Expr));
     }
 
     return IGraphTransformer::TStatus::Ok;
@@ -431,7 +429,7 @@ IGraphTransformer::TStatus BlockCombineHashedWrapper(const TExprNode::TPtr& inpu
     }
 
     for (auto& t : retMultiType) {
-        t = ctx.Expr.MakeType<TBlockExprType>(t);
+        t = MakeBlockType(*t, ctx.Expr);
     }
 
     retMultiType.push_back(ctx.Expr.MakeType<TScalarExprType>(ctx.Expr.MakeType<TDataExprType>(EDataSlot::Uint64)));
@@ -462,7 +460,7 @@ IGraphTransformer::TStatus BlockMergeFinalizeHashedWrapper(const TExprNode::TPtr
     }
 
     for (auto& t : retMultiType) {
-        t = ctx.Expr.MakeType<TBlockExprType>(t);
+        t = MakeBlockType(*t, ctx.Expr);
     }
 
     if (many) {
diff --git a/ydb/library/yql/core/type_ann/type_ann_wide.cpp b/ydb/library/yql/core/type_ann/type_ann_wide.cpp
index a60ed79c84..ef2429c79e 100644
--- a/ydb/library/yql/core/type_ann/type_ann_wide.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_wide.cpp
@@ -647,12 +647,7 @@ IGraphTransformer::TStatus WideToBlocksWrapper(const TExprNode::TPtr& input, TEx
             return IGraphTransformer::TStatus::Error;
         }
 
-        if (type->HasFixedSizeRepr()) {
-            retMultiType.push_back(ctx.Expr.MakeType<TBlockExprType>(type));
-        } else {
-            retMultiType.push_back(ctx.Expr.MakeType<TChunkedBlockExprType>(type));
-        }
-
+        retMultiType.push_back(MakeBlockType(*type, ctx.Expr));
     }
 
     retMultiType.push_back(ctx.Expr.MakeType<TScalarExprType>(ctx.Expr.MakeType<TDataExprType>(EDataSlot::Uint64)));
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index 607615739f..0436652dfc 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -5370,6 +5370,13 @@ const TTypeAnnotationNode* GetBlockItemType(const TTypeAnnotationNode& type, boo
     }
 }
 
+const TTypeAnnotationNode* MakeBlockType(const TTypeAnnotationNode& blockItemType, TExprContext& ctx, bool withChunked) {
+    if (withChunked && !blockItemType.HasFixedSizeRepr()) {
+        return ctx.MakeType<TChunkedBlockExprType>(&blockItemType);
+    }
+    return ctx.MakeType<TBlockExprType>(&blockItemType);
+}
+
 const TTypeAnnotationNode* AggApplySerializedStateType(const TExprNode::TPtr& input, TExprContext& ctx) {
     auto name = input->Child(0)->Content();
     if (name == "count" || name == "count_all" || name == "sum" || name == "min" || name == "max") {
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index 4ac91b5ed2..7f55dd94ae 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -301,6 +301,7 @@ bool HasContextFuncs(const TExprNode& input);
 bool EnsureBlockOrScalarType(const TExprNode& node, TExprContext& ctx, bool allowChunked = false);
 bool EnsureBlockOrScalarType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx, bool allowChunked = false);
 const TTypeAnnotationNode* GetBlockItemType(const TTypeAnnotationNode& type, bool& isScalar);
+const TTypeAnnotationNode* MakeBlockType(const TTypeAnnotationNode& blockItemType, TExprContext& ctx, bool withChunked = true);
 
 const TTypeAnnotationNode* AggApplySerializedStateType(const TExprNode::TPtr& input, TExprContext& ctx);
 bool GetSumResultType(const TPositionHandle& pos, const TTypeAnnotationNode& inputType, const TTypeAnnotationNode*& retType, TExprContext& ctx);