introduced Block/Scalar type decorators

7 files changed, 254 insertions, 1 deletions

diff --git a/ydb/library/yql/ast/yql_expr.cpp b/ydb/library/yql/ast/yql_expr.cpp
index 6c72acf1c3..8f4148c57a 100644
--- a/ydb/library/yql/ast/yql_expr.cpp
+++ b/ydb/library/yql/ast/yql_expr.cpp
@@ -559,6 +559,28 @@ namespace {
                         return nullptr;
 
                     return Expr.MakeType<TFlowExprType>(r);
+                } else if (content == TStringBuf("Block")) {
+                    if (node.GetChildrenCount() != 2) {
+                        AddError(node, "Bad block type annotation");
+                        return nullptr;
+                    }
+
+                    auto r = CompileTypeAnnotationNode(*node.GetChild(1));
+                    if (!r)
+                        return nullptr;
+
+                    return Expr.MakeType<TBlockExprType>(r);
+                } else if (content == TStringBuf("Scalar")) {
+                    if (node.GetChildrenCount() != 2) {
+                        AddError(node, "Bad scalar type annotation");
+                        return nullptr;
+                    }
+
+                    auto r = CompileTypeAnnotationNode(*node.GetChild(1));
+                    if (!r)
+                        return nullptr;
+
+                    return Expr.MakeType<TScalarExprType>(r);
                 } else {
                     AddError(node, TStringBuilder() << "Unknown type annotation");
                     return nullptr;
@@ -806,6 +828,23 @@ namespace {
         {
             return TAstNode::NewLiteralAtom(TPosition(), TStringBuf("EmptyDict"), pool);
         }
+
+        case ETypeAnnotationKind::Block:
+        {
+            auto type = annotation.Cast<TBlockExprType>();
+            auto self = TAstNode::NewLiteralAtom(TPosition(), TStringBuf("Block"), pool);
+            auto itemType = ConvertTypeAnnotationToAst(*type->GetItemType(), pool, refAtoms);
+            return TAstNode::NewList(TPosition(), pool, self, itemType);
+        }
+
+        case ETypeAnnotationKind::Scalar:
+        {
+            auto type = annotation.Cast<TScalarExprType>();
+            auto self = TAstNode::NewLiteralAtom(TPosition(), TStringBuf("Scalar"), pool);
+            auto itemType = ConvertTypeAnnotationToAst(*type->GetItemType(), pool, refAtoms);
+            return TAstNode::NewList(TPosition(), pool, self, itemType);
+        }
+
         case ETypeAnnotationKind::LastType:
             YQL_ENSURE(false, "Unknown kind: " << annotation.GetKind());
 
@@ -3269,6 +3308,24 @@ const TFlowExprType* TMakeTypeImpl<TFlowExprType>::Make(TExprContext& ctx, const
     return AddType<TFlowExprType>(ctx, hash, itemType);
 }
 
+const TBlockExprType* TMakeTypeImpl<TBlockExprType>::Make(TExprContext& ctx, const TTypeAnnotationNode* itemType) {
+    const auto hash = TBlockExprType::MakeHash(itemType);
+    TBlockExprType sample(hash, itemType);
+    if (const auto found = FindType(sample, ctx))
+        return found;
+
+    return AddType<TBlockExprType>(ctx, hash, itemType);
+}
+
+const TScalarExprType* TMakeTypeImpl<TScalarExprType>::Make(TExprContext& ctx, const TTypeAnnotationNode* itemType) {
+    const auto hash = TScalarExprType::MakeHash(itemType);
+    TScalarExprType sample(hash, itemType);
+    if (const auto found = FindType(sample, ctx))
+        return found;
+
+    return AddType<TScalarExprType>(ctx, hash, itemType);
+}
+
 bool CompareExprTrees(const TExprNode*& one, const TExprNode*& two) {
     TArgumentsMap map;
     ui32 level = 0;
diff --git a/ydb/library/yql/ast/yql_expr.h b/ydb/library/yql/ast/yql_expr.h
index 93a220df49..02790cb984 100644
--- a/ydb/library/yql/ast/yql_expr.h
+++ b/ydb/library/yql/ast/yql_expr.h
@@ -79,6 +79,8 @@ class TStreamExprType;
 class TFlowExprType;
 class TEmptyListExprType;
 class TEmptyDictExprType;
+class TBlockExprType;
+class TScalarExprType;
 
 const size_t DefaultMistypeDistance = 3;
 
@@ -111,6 +113,8 @@ struct TTypeAnnotationVisitor {
     virtual void Visit(const TVariantExprType& type) = 0;
     virtual void Visit(const TEmptyListExprType& type) = 0;
     virtual void Visit(const TEmptyDictExprType& type) = 0;
+    virtual void Visit(const TBlockExprType& type) = 0;
+    virtual void Visit(const TScalarExprType& type) = 0;
 };
 
 enum ETypeAnnotationFlags {
@@ -616,6 +620,60 @@ private:
     const TTypeAnnotationNode* ItemType;
 };
 
+class TBlockExprType : public TTypeAnnotationNode {
+public:
+    static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Block;
+
+    TBlockExprType(ui64 hash, const TTypeAnnotationNode* itemType)
+        : TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash)
+        , ItemType(itemType)
+    {
+    }
+
+    static ui64 MakeHash(const TTypeAnnotationNode* itemType) {
+        ui64 hash = TypeHashMagic | (ui64)ETypeAnnotationKind::Block;
+        return StreamHash(itemType->GetHash(), hash);
+    }
+
+    const TTypeAnnotationNode* GetItemType() const {
+        return ItemType;
+    }
+
+    bool operator==(const TBlockExprType& other) const {
+        return GetItemType() == other.GetItemType();
+    }
+
+private:
+    const TTypeAnnotationNode* ItemType;
+};
+
+class TScalarExprType : public TTypeAnnotationNode {
+public:
+    static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Scalar;
+
+    TScalarExprType(ui64 hash, const TTypeAnnotationNode* itemType)
+        : TTypeAnnotationNode(KindValue, itemType->GetFlags() | TypeNonPersistable, hash)
+        , ItemType(itemType)
+    {
+    }
+
+    static ui64 MakeHash(const TTypeAnnotationNode* itemType) {
+        ui64 hash = TypeHashMagic | (ui64)ETypeAnnotationKind::Scalar;
+        return StreamHash(itemType->GetHash(), hash);
+    }
+
+    const TTypeAnnotationNode* GetItemType() const {
+        return ItemType;
+    }
+
+    bool operator==(const TScalarExprType& other) const {
+        return GetItemType() == other.GetItemType();
+    }
+
+private:
+    const TTypeAnnotationNode* ItemType;
+};
+
 class TDataExprType : public TTypeAnnotationNode {
 public:
     static constexpr ETypeAnnotationKind KindValue = ETypeAnnotationKind::Data;
@@ -1290,6 +1348,12 @@ inline bool TTypeAnnotationNode::Equals(const TTypeAnnotationNode& node) const {
     case ETypeAnnotationKind::Multi:
         return static_cast<const TMultiExprType&>(*this) == static_cast<const TMultiExprType&>(node);
 
+    case ETypeAnnotationKind::Block:
+        return static_cast<const TBlockExprType&>(*this) == static_cast<const TBlockExprType&>(node);
+
+    case ETypeAnnotationKind::Scalar:
+        return static_cast<const TScalarExprType&>(*this) == static_cast<const TScalarExprType&>(node);
+
     case ETypeAnnotationKind::LastType:
         YQL_ENSURE(false, "Incorrect type");
 
@@ -1347,6 +1411,10 @@ inline void TTypeAnnotationNode::Accept(TTypeAnnotationVisitor& visitor) const {
         return visitor.Visit(static_cast<const TEmptyDictExprType&>(*this));
     case ETypeAnnotationKind::Multi:
         return visitor.Visit(static_cast<const TMultiExprType&>(*this));
+    case ETypeAnnotationKind::Block:
+        return visitor.Visit(static_cast<const TBlockExprType&>(*this));
+    case ETypeAnnotationKind::Scalar:
+        return visitor.Visit(static_cast<const TScalarExprType&>(*this));
     case ETypeAnnotationKind::LastType:
         YQL_ENSURE(false, "Incorrect type");
     }
@@ -2261,6 +2329,16 @@ struct TMakeTypeImpl<TFlowExprType> {
     static const TFlowExprType* Make(TExprContext& ctx, const TTypeAnnotationNode* itemType);
 };
 
+template <>
+struct TMakeTypeImpl<TBlockExprType> {
+    static const TBlockExprType* Make(TExprContext& ctx, const TTypeAnnotationNode* itemType);
+};
+
+template <>
+struct TMakeTypeImpl<TScalarExprType> {
+    static const TScalarExprType* Make(TExprContext& ctx, const TTypeAnnotationNode* itemType);
+};
+
 using TSingletonTypeCache = std::tuple<
     const TVoidExprType*,
     const TNullExprType*,
diff --git a/ydb/library/yql/ast/yql_expr_types.h b/ydb/library/yql/ast/yql_expr_types.h
index c420b45e49..bb628af272 100644
--- a/ydb/library/yql/ast/yql_expr_types.h
+++ b/ydb/library/yql/ast/yql_expr_types.h
@@ -28,7 +28,9 @@ namespace NYql {
     xx(EmptyList, 21) \
     xx(EmptyDict, 22) \
     xx(Multi, 23) \
-    xx(Pg, 24)
+    xx(Pg, 24) \
+    xx(Block, 25) \
+    xx(Scalar, 26)
 
 enum class ETypeAnnotationKind : ui64 {
     YQL_TYPE_ANN_KIND_MAP(ENUM_VALUE_GEN)
diff --git a/ydb/library/yql/ast/yql_type_string.cpp b/ydb/library/yql/ast/yql_type_string.cpp
index 928aecacdb..78ced74895 100644
--- a/ydb/library/yql/ast/yql_type_string.cpp
+++ b/ydb/library/yql/ast/yql_type_string.cpp
@@ -80,6 +80,8 @@ enum EToken
     TOKEN_TYPE_MAX = -47,
     TOKEN_JSON_DOCUMENT = -48,
     TOKEN_DYNUMBER = -49,
+    TOKEN_SCALAR = -50,
+    TOKEN_BLOCK = -51,
 
     // identifiers
     TOKEN_IDENTIFIER = -100,
@@ -142,6 +144,8 @@ EToken TokenTypeFromStr(TStringBuf str)
         { TStringBuf("EmptyDict"), TOKEN_EMPTYDICT },
         { TStringBuf("JsonDocument"), TOKEN_JSON_DOCUMENT },
         { TStringBuf("DyNumber"), TOKEN_DYNUMBER },
+        { TStringBuf("Block"), TOKEN_BLOCK},
+        { TStringBuf("Scalar"), TOKEN_SCALAR},
     };
 
     auto it = map.find(str);
@@ -303,6 +307,14 @@ private:
             type = ParseEnumType();
             break;
 
+        case TOKEN_BLOCK:
+            type = ParseBlockType();
+            break;
+
+        case TOKEN_SCALAR:
+            type = ParseScalarType();
+            break;
+
         default:
             if (Identifier.empty()) {
                 return AddError("Expected type");
@@ -645,6 +657,28 @@ private:
         return MakeFlowType(itemType);
     }
 
+    TAstNode* ParseBlockType() {
+        GetNextToken(); // eat keyword
+        EXPECT_AND_SKIP_TOKEN('<', nullptr);
+
+        auto itemType = ParseType();
+        if (!itemType) return nullptr;
+
+        EXPECT_AND_SKIP_TOKEN('>', nullptr);
+        return MakeBlockType(itemType);
+    }
+
+    TAstNode* ParseScalarType() {
+        GetNextToken(); // eat keyword
+        EXPECT_AND_SKIP_TOKEN('<', nullptr);
+
+        auto itemType = ParseType();
+        if (!itemType) return nullptr;
+
+        EXPECT_AND_SKIP_TOKEN('>', nullptr);
+        return MakeScalarType(itemType);
+    }
+
     TAstNode* ParseDecimalType() {
         GetNextToken(); // eat keyword
         EXPECT_AND_SKIP_TOKEN('(', nullptr);
@@ -890,6 +924,22 @@ private:
         return MakeList(items, Y_ARRAY_SIZE(items));
     }
 
+    TAstNode* MakeBlockType(TAstNode* itemType) {
+        TAstNode* items[] = {
+            MakeLiteralAtom(TStringBuf("BlockType")),
+            itemType,
+        };
+        return MakeList(items, Y_ARRAY_SIZE(items));
+    }
+
+    TAstNode* MakeScalarType(TAstNode* itemType) {
+        TAstNode* items[] = {
+            MakeLiteralAtom(TStringBuf("ScalarType")),
+            itemType,
+        };
+        return MakeList(items, Y_ARRAY_SIZE(items));
+    }
+
     TAstNode* MakeVariantType(TAstNode* underlyingType) {
         TAstNode* items[] = {
             MakeLiteralAtom(TStringBuf("VariantType")),
@@ -1211,6 +1261,20 @@ private:
         Out_ << '>';
     }
 
+    void Visit(const TBlockExprType& type) final {
+        TopLevel = false;
+        Out_ << TStringBuf("Block<");
+        type.GetItemType()->Accept(*this);
+        Out_ << '>';
+    }
+
+    void Visit(const TScalarExprType& type) final {
+        TopLevel = false;
+        Out_ << TStringBuf("Scalar<");
+        type.GetItemType()->Accept(*this);
+        Out_ << '>';
+    }
+
     void Visit(const TDataExprType& type) final {
         TopLevel = false;
         Out_ << type.GetName();
diff --git a/ydb/library/yql/ast/yql_type_string_ut.cpp b/ydb/library/yql/ast/yql_type_string_ut.cpp
index abeee1fec3..4d9c61a331 100644
--- a/ydb/library/yql/ast/yql_type_string_ut.cpp
+++ b/ydb/library/yql/ast/yql_type_string_ut.cpp
@@ -567,6 +567,14 @@ Y_UNIT_TEST_SUITE(TTypeString)
         TestFormat("((Flow (Data String)))", "Flow<String>");
     }
 
+    Y_UNIT_TEST(FormatBlock) {
+        TestFormat("((Block (Data String)))", "Block<String>");
+    }
+
+    Y_UNIT_TEST(FormatScalar) {
+        TestFormat("((Scalar (Data String)))", "Scalar<String>");
+    }
+
     Y_UNIT_TEST(FormatOptional) {
         TestFormat("((Optional (Data Uint32)))", "Optional<Uint32>");
         TestFormat("((List (Optional (Data Uint32))))", "List<Uint32?>");
diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 70e24c481c..09f6214977 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -11238,6 +11238,8 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["VariantType"] = &TypeWrapper<ETypeAnnotationKind::Variant>;
         Functions["StreamType"] = &TypeWrapper<ETypeAnnotationKind::Stream>;
         Functions["FlowType"] = &TypeWrapper<ETypeAnnotationKind::Flow>;
+        Functions["BlockType"] = &TypeWrapper<ETypeAnnotationKind::Block>;
+        Functions["ScalarType"] = &TypeWrapper<ETypeAnnotationKind::Scalar>;
         Functions["Nothing"] = &NothingWrapper;
         Functions["AsOptionalType"] = &AsOptionalTypeWrapper;
         Functions["List"] = &ListWrapper;
diff --git a/ydb/library/yql/core/type_ann/type_ann_types.cpp b/ydb/library/yql/core/type_ann/type_ann_types.cpp
index c27c0c433f..f596a2f6c9 100644
--- a/ydb/library/yql/core/type_ann/type_ann_types.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_types.cpp
@@ -283,6 +283,48 @@ namespace NTypeAnnImpl {
     }
 
     template <>
+    IGraphTransformer::TStatus TypeWrapper<ETypeAnnotationKind::Block>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        Y_UNUSED(output);
+        if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (auto status = EnsureTypeRewrite(input->HeadRef(), ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
+        auto itemType = input->Child(0)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
+        if (!EnsureInspectableType(input->Child(0)->Pos(), *itemType, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto type = ctx.Expr.MakeType<TBlockExprType>(itemType);
+        input->SetTypeAnn(ctx.Expr.MakeType<TTypeExprType>(type));
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    template <>
+    IGraphTransformer::TStatus TypeWrapper<ETypeAnnotationKind::Scalar>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        Y_UNUSED(output);
+        if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (auto status = EnsureTypeRewrite(input->HeadRef(), ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
+        auto itemType = input->Child(0)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
+        if (!EnsureInspectableType(input->Child(0)->Pos(), *itemType, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto type = ctx.Expr.MakeType<TScalarExprType>(itemType);
+        input->SetTypeAnn(ctx.Expr.MakeType<TTypeExprType>(type));
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    template <>
     IGraphTransformer::TStatus TypeWrapper<ETypeAnnotationKind::Optional>(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         Y_UNUSED(output);
         if (!EnsureArgsCount(*input, 1, ctx.Expr)) {