[yql] Precompute type annotation

YQL-12393

ref:e0441a303d4ffdb5803d5c976042ff392852c986

3 files changed, 28 insertions, 24 deletions

diff --git a/ydb/library/yql/dq/type_ann/dq_type_ann.cpp b/ydb/library/yql/dq/type_ann/dq_type_ann.cpp
index 04ff29b03f6..cefe4c7fa3c 100644
--- a/ydb/library/yql/dq/type_ann/dq_type_ann.cpp
+++ b/ydb/library/yql/dq/type_ann/dq_type_ann.cpp
@@ -440,30 +440,6 @@ const TStructExprType* GetDqJoinResultType(const TExprNode::TPtr& input, bool st
         rightTableLabel, join.JoinType(), join.JoinKeys(), ctx);
 }
 
-TStatus AnnotateDqPrecompute(const TExprNode::TPtr& node, TExprContext& ctx) {
-    if (!EnsureArgsCount(*node, 1, ctx)) {
-        return TStatus::Error;
-    }
-
-    node->SetTypeAnn(node->Child(TDqPrecompute::idx_Input)->GetTypeAnn());
-    return TStatus::Ok;
-}
-
-TStatus AnnotateDqPhyPrecompute(const TExprNode::TPtr& node, TExprContext& ctx) {
-    if (!EnsureArgsCount(*node, 1, ctx)) {
-        return TStatus::Error;
-    }
-
-    auto* cn = node->Child(TDqPhyPrecompute::idx_Connection);
-    if (!TDqConnection::Match(cn)) {
-        ctx.AddError(TIssue(ctx.GetPosition(cn->Pos()), TStringBuilder() << "Expected DqConnection, got " << cn->Content()));
-        return TStatus::Error;
-    }
-
-    node->SetTypeAnn(cn->GetTypeAnn());
-    return TStatus::Ok;
-}
-
 } // unnamed
 
 TStatus AnnotateDqStage(const TExprNode::TPtr& input, TExprContext& ctx) {
@@ -854,6 +830,30 @@ TStatus AnnotateTransformSettings(const TExprNode::TPtr& input, TExprContext& ct
     return TStatus::Ok;
 }
 
+TStatus AnnotateDqPrecompute(const TExprNode::TPtr& node, TExprContext& ctx) {
+    if (!EnsureArgsCount(*node, 1, ctx)) {
+        return TStatus::Error;
+    }
+
+    node->SetTypeAnn(node->Child(TDqPrecompute::idx_Input)->GetTypeAnn());
+    return TStatus::Ok;
+}
+
+TStatus AnnotateDqPhyPrecompute(const TExprNode::TPtr& node, TExprContext& ctx) {
+    if (!EnsureArgsCount(*node, 1, ctx)) {
+        return TStatus::Error;
+    }
+
+    auto* cn = node->Child(TDqPhyPrecompute::idx_Connection);
+    if (!TDqConnection::Match(cn)) {
+        ctx.AddError(TIssue(ctx.GetPosition(cn->Pos()), TStringBuilder() << "Expected DqConnection, got " << cn->Content()));
+        return TStatus::Error;
+    }
+
+    node->SetTypeAnn(cn->GetTypeAnn());
+    return TStatus::Ok;
+}
+
 THolder<IGraphTransformer> CreateDqTypeAnnotationTransformer(TTypeAnnotationContext& typesCtx) {
     auto coreTransformer = CreateExtCallableTypeAnnotationTransformer(typesCtx);
 
diff --git a/ydb/library/yql/dq/type_ann/dq_type_ann.h b/ydb/library/yql/dq/type_ann/dq_type_ann.h
index 84f8afe289e..c870ce7c7fc 100644
--- a/ydb/library/yql/dq/type_ann/dq_type_ann.h
+++ b/ydb/library/yql/dq/type_ann/dq_type_ann.h
@@ -23,6 +23,8 @@ IGraphTransformer::TStatus AnnotateDqSource(const TExprNode::TPtr& input, TExprC
 IGraphTransformer::TStatus AnnotateDqSink(const TExprNode::TPtr& input, TExprContext& ctx);
 IGraphTransformer::TStatus AnnotateDqQuery(const TExprNode::TPtr& input, TExprContext& ctx);
 IGraphTransformer::TStatus AnnotateTransformSettings(const TExprNode::TPtr& input, TExprContext& ctx);
+IGraphTransformer::TStatus AnnotateDqPrecompute(const TExprNode::TPtr& node, TExprContext& ctx);
+IGraphTransformer::TStatus AnnotateDqPhyPrecompute(const TExprNode::TPtr& node, TExprContext& ctx);
 
 THolder<IGraphTransformer> CreateDqTypeAnnotationTransformer(NYql::TTypeAnnotationContext& typesCtx);
 
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_datasink_type_ann.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_datasink_type_ann.cpp
index 7d7f7b8fee7..2763be33c2b 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_datasink_type_ann.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_datasink_type_ann.cpp
@@ -35,6 +35,8 @@ public:
         AddHandler({TDqWrite::CallableName()}, Hndl(&TDqsDataSinkTypeAnnotationTransformer::AnnotateDqWrite));
         AddHandler({TDqQuery::CallableName()}, Hndl(&NDq::AnnotateDqQuery));
         AddHandler({TTransformSettings::CallableName()}, Hndl(&NDq::AnnotateTransformSettings));
+        AddHandler({TDqPrecompute::CallableName()}, Hndl(&NDq::AnnotateDqPrecompute));
+        AddHandler({TDqPhyPrecompute::CallableName()}, Hndl(&NDq::AnnotateDqPhyPrecompute));
     }
 
 private: