[yt provider] Add YtMaterialize callable (#10057)

3 files changed, 48 insertions, 0 deletions

diff --git a/ydb/library/yql/providers/yt/expr_nodes/yql_yt_expr_nodes.json b/ydb/library/yql/providers/yt/expr_nodes/yql_yt_expr_nodes.json
index 984bdf7a3d..d6508aef6c 100644
--- a/ydb/library/yql/providers/yt/expr_nodes/yql_yt_expr_nodes.json
+++ b/ydb/library/yql/providers/yt/expr_nodes/yql_yt_expr_nodes.json
@@ -348,6 +348,15 @@
             ]
         },
         {
+            "Name": "TYtMaterialize",
+            "Base": "TYtOpBase",
+            "Match": {"Type": "Callable", "Name": "YtMaterialize!"},
+            "Children": [
+                {"Index": 2, "Name": "Input", "Type": "TExprBase"},
+                {"Index": 3, "Name": "Settings", "Type": "TCoNameValueTupleList"}
+            ]
+        },
+        {
             "Name": "TYtTouch",
             "Base": "TYtOutputOpBase",
             "Match": {"Type": "Callable", "Name": "YtTouch!"}
diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_datasink_constraints.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_datasink_constraints.cpp
index 5481d3700d..4d80264abd 100644
--- a/ydb/library/yql/providers/yt/provider/yql_yt_datasink_constraints.cpp
+++ b/ydb/library/yql/providers/yt/provider/yql_yt_datasink_constraints.cpp
@@ -42,6 +42,7 @@ public:
         AddHandler({TYtStatOut::CallableName()}, Hndl(&TYtDataSinkConstraintTransformer::HandleDefault));
         AddHandler({TYtDqProcessWrite ::CallableName()}, Hndl(&TYtDataSinkConstraintTransformer::HandleDqProcessWrite));
         AddHandler({TYtTryFirst ::CallableName()}, Hndl(&TYtDataSinkConstraintTransformer::HandleTryFirst));
+        AddHandler({TYtMaterialize ::CallableName()}, Hndl(&TYtDataSinkConstraintTransformer::HandleMaterialize));
     }
 private:
     static void CopyExcept(TExprNode* dst, const TExprNode& from, const TStringBuf& except) {
@@ -447,6 +448,12 @@ private:
         input.Ptr()->CopyConstraints(input.Ref().Tail());
         return TStatus::Ok;
     }
+
+    TStatus HandleMaterialize(TExprBase input, TExprContext&) {
+        input.Ptr()->CopyConstraints(*input.Ref().Child(TYtMaterialize::idx_Input));
+        return TStatus::Ok;
+    }
+
 private:
     const TYtState::TPtr State_;
     const bool SubGraph;
diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
index d5df0549c5..284ff67e0d 100644
--- a/ydb/library/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
+++ b/ydb/library/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
@@ -86,6 +86,7 @@ public:
         AddHandler({TYtDqWrite::CallableName()}, Hndl(&TYtDataSinkTypeAnnotationTransformer::HandleDqWrite<false>));
         AddHandler({TYtDqWideWrite::CallableName()}, Hndl(&TYtDataSinkTypeAnnotationTransformer::HandleDqWrite<true>));
         AddHandler({TYtTryFirst::CallableName()}, Hndl(&TYtDataSinkTypeAnnotationTransformer::HandleTryFirst));
+        AddHandler({TYtMaterialize::CallableName()}, Hndl(&TYtDataSinkTypeAnnotationTransformer::HandleMaterialize));
     }
 
 private:
@@ -2066,6 +2067,37 @@ private:
         input.Ptr()->SetTypeAnn(input.Ref().Head().GetTypeAnn());
         return TStatus::Ok;
     }
+
+    TStatus HandleMaterialize(TExprBase input, TExprContext& ctx) {
+        if (!EnsureArgsCount(input.Ref(), 4, ctx)) {
+            return TStatus::Error;
+        }
+
+        if (!ValidateOpBase(input.Ptr(), ctx)) {
+            return TStatus::Error;
+        }
+
+        if (!EnsureSeqOrOptionalType(*input.Ref().Child(TYtMaterialize::idx_Input), ctx)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+        const auto& itemType = GetSeqItemType(*input.Ref().Child(TYtMaterialize::idx_Input)->GetTypeAnn());
+        if (!EnsurePersistableType(input.Ref().Head().Pos(), itemType, ctx)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        // Basic Settings validation
+        if (!EnsureTuple(*input.Ref().Child(TYtMaterialize::idx_Settings), ctx)) {
+            return TStatus::Error;
+        }
+
+        if (!ValidateSettings(*input.Ref().Child(TYtMaterialize::idx_Settings), EYtSettingTypes{}, ctx)) {
+            return TStatus::Error;
+        }
+
+        input.Ptr()->SetTypeAnn(ctx.MakeType<TListExprType>(&itemType));
+        return TStatus::Ok;
+    }
+
 private:
     const TYtState::TPtr State_;
 };