Fix convertion of TChunkedBlockExpr to MKQL type

2 files changed, 10 insertions, 3 deletions

diff --git a/ydb/library/yql/providers/common/arrow_resolve/yql_simple_arrow_resolver.cpp b/ydb/library/yql/providers/common/arrow_resolve/yql_simple_arrow_resolver.cpp
index c8b10f48c4..418acc486d 100644
--- a/ydb/library/yql/providers/common/arrow_resolve/yql_simple_arrow_resolver.cpp
+++ b/ydb/library/yql/providers/common/arrow_resolve/yql_simple_arrow_resolver.cpp
@@ -30,6 +30,7 @@ private:
             TVector<TType*> mkqlInputTypes;
             for (const auto& type : argTypes) {
                 auto mkqlType = NCommon::BuildType(*type, pgmBuilder, null);
+                YQL_ENSURE(mkqlType, "Failed to convert type " << *type << " to MKQL type");
                 mkqlInputTypes.emplace_back(mkqlType);
             }
             TType* mkqlOutputType = NCommon::BuildType(*returnType, pgmBuilder, null);
diff --git a/ydb/library/yql/providers/common/mkql/yql_type_mkql.cpp b/ydb/library/yql/providers/common/mkql/yql_type_mkql.cpp
index cbd56a87b2..7e31ec164f 100644
--- a/ydb/library/yql/providers/common/mkql/yql_type_mkql.cpp
+++ b/ydb/library/yql/providers/common/mkql/yql_type_mkql.cpp
@@ -203,9 +203,12 @@ NKikimr::NMiniKQL::TType* BuildType(const TTypeAnnotationNode& annotation, NKiki
         return pgmBuilder.GetTypeEnvironment().GetTypeOfEmptyDict();
     }
 
+    case ETypeAnnotationKind::ChunkedBlock:
     case ETypeAnnotationKind::Block: {
-        auto block = annotation.Cast<TBlockExprType>();
-        auto itemType = BuildType(*block->GetItemType(), pgmBuilder, err, withTagged);
+        auto blockItem = annotation.GetKind() == ETypeAnnotationKind::Block ?
+            annotation.Cast<TBlockExprType>()->GetItemType() :
+            annotation.Cast<TChunkedBlockExprType>()->GetItemType();
+        auto itemType = BuildType(*blockItem, pgmBuilder, err, withTagged);
         if (!itemType) {
             return nullptr;
         }
@@ -221,7 +224,10 @@ NKikimr::NMiniKQL::TType* BuildType(const TTypeAnnotationNode& annotation, NKiki
         return pgmBuilder.NewBlockType(itemType, NKikimr::NMiniKQL::TBlockType::EShape::Scalar);
     }
 
-    default:
+    case ETypeAnnotationKind::Item:
+    case ETypeAnnotationKind::World:
+    case ETypeAnnotationKind::Error:
+    case ETypeAnnotationKind::LastType:
         return nullptr;
     }
 }