rework solomon sink type ann (#6636)

Author: uzhastik

ydb/library/yql/providers/solomon/expr_nodes/yql_solomon_expr_nodes.json

@@ -105,7 +105,8 @@
                 {"Index": 1, "Name": "Project", "Type": "TCoAtom"},
                 {"Index": 2, "Name": "Cluster", "Type": "TCoAtom"},
                 {"Index": 3, "Name": "Service", "Type": "TCoAtom"},
-                {"Index": 4, "Name": "Token", "Type": "TCoSecureParam", "Optional": true}
+                {"Index": 4, "Name": "RowType", "Type": "TExprBase"},
+                {"Index": 5, "Name": "Token", "Type": "TCoSecureParam", "Optional": true}
             ]
         }
     ]

ydb/library/yql/providers/solomon/provider/yql_solomon_datasink_type_ann.cpp

@@ -115,7 +115,7 @@ class TSolomonDataSinkTypeAnnotationTransformer : public TVisitorTransformerBase
     TStatus HandleSoShard(TExprBase input, TExprContext& ctx) {
         YQL_ENSURE(!State_->IsRtmrMode(), "SoShard can't be used in rtmr mode");
 
-        if (!EnsureMinArgsCount(input.Ref(), 4, ctx) || !EnsureMaxArgsCount(input.Ref(), 5, ctx)) {
+        if (!EnsureMinMaxArgsCount(input.Ref(), 5, 6, ctx)) {
             return TStatus::Error;
         }
 
@@ -137,6 +137,10 @@ class TSolomonDataSinkTypeAnnotationTransformer : public TVisitorTransformerBase
             return TStatus::Error;
         }
 
+        if (!EnsureType(shard.RowType().Ref(), ctx)) {
+            return TStatus::Error;
+        }
+
         if (shard.Token() && !EnsureCallable(shard.Token().Ref(), ctx)) {
             return TStatus::Error;
         }

ydb/library/yql/providers/solomon/provider/yql_solomon_dq_integration.cpp

@@ -44,19 +44,6 @@ NSo::NProto::ESolomonClusterType MapClusterType(TSolomonClusterConfig::ESolomonC
     }
 }
 
-const TTypeAnnotationNode* GetItemType(const TExprNode& node) {
-    const TTypeAnnotationNode* typeAnn = node.GetTypeAnn();
-    switch (typeAnn->GetKind()) {
-        case ETypeAnnotationKind::Flow:
-            return typeAnn->Cast<TFlowExprType>()->GetItemType();
-        case ETypeAnnotationKind::Stream:
-            return typeAnn->Cast<TStreamExprType>()->GetItemType();
-        default: break;
-    }
-    YQL_ENSURE(false, "Invalid solomon sink type " << typeAnn->GetKind());
-    return nullptr;
-}
-
 void FillScheme(const TTypeAnnotationNode& itemType, NSo::NProto::TDqSolomonShardScheme& scheme) {
     int index = 0;
     for (const TItemExprType* structItem : itemType.Cast<TStructExprType>()->GetItems()) {
@@ -316,7 +303,7 @@ class TSolomonDqIntegration: public TDqIntegrationBase {
         shardDesc.SetClusterType(MapClusterType(clusterDesc->GetClusterType()));
         shardDesc.SetUseSsl(clusterDesc->GetUseSsl());
 
-        const TTypeAnnotationNode* itemType = GetItemType(node);
+        const TTypeAnnotationNode* itemType = shard.RowType().Ref().GetTypeAnn()->Cast<TTypeExprType>()->GetType();
         FillScheme(*itemType, *shardDesc.MutableScheme());
 
         if (auto maybeToken = shard.Token()) {

ydb/library/yql/providers/solomon/provider/yql_solomon_physical_optimize.cpp

@@ -100,7 +100,14 @@ class TSoPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         YQL_CLOG(INFO, ProviderSolomon) << "Optimize SoWriteToShard";
 
         const auto solomonCluster = TString(write.DataSink().Cluster().Value());
-        auto shard = BuildSolomonShard(write.Shard().Cast<TCoAtom>(), ctx, solomonCluster);
+        auto* typeAnn = write.Input().Ref().GetTypeAnn();
+        const TTypeAnnotationNode* inputItemType = nullptr;
+        if (!EnsureNewSeqType<false, true, false>(write.Input().Pos(), *typeAnn, ctx, &inputItemType)) {
+            return {};
+        }
+
+        auto rowTypeNode = ExpandType(write.Pos(), *inputItemType, ctx);
+        auto shard = BuildSolomonShard(write.Shard().Cast<TCoAtom>(), TExprBase(rowTypeNode), ctx, solomonCluster);
 
         auto dqSink = Build<TDqSink>(ctx, write.Pos())
             .DataSink(write.DataSink())
@@ -131,7 +138,7 @@ class TSoPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
     }
 
 private:
-    TCallable BuildSolomonShard(TCoAtom shardNode, TExprContext& ctx, TString solomonCluster) const {
+    TCallable BuildSolomonShard(TCoAtom shardNode, TExprBase rowType, TExprContext& ctx, TString solomonCluster) const {
         const auto* clusterDesc = State_->Configuration->ClusterConfigs.FindPtr(solomonCluster);
         YQL_ENSURE(clusterDesc, "Unknown cluster " << solomonCluster);
 
@@ -148,31 +155,14 @@ class TSoPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         }
         YQL_ENSURE(!cluster.empty(), "Cluster is not defined. You can define it inside connection, or inside query.");
 
-        auto solomonClusterAtom = Build<TCoAtom>(ctx, shardNode.Pos())
-            .Value(solomonCluster)
-            .Done();
-
-        auto projectAtom = Build<TCoAtom>(ctx, shardNode.Pos())
-            .Value(project)
-            .Done();
-
-        auto clusterAtom = Build<TCoAtom>(ctx, shardNode.Pos())
-            .Value(cluster)
+        return Build<TSoShard>(ctx, shardNode.Pos())
+            .SolomonCluster<TCoAtom>().Value(solomonCluster).Build()
+            .Project<TCoAtom>().Value(project).Build()
+            .Cluster<TCoAtom>().Value(cluster).Build()
+            .Service<TCoAtom>().Value(service).Build()
+            .RowType(rowType)
+            .Token<TCoSecureParam>().Name().Build("cluster:default_" + solomonCluster).Build()
             .Done();
-
-        auto serviceAtom = Build<TCoAtom>(ctx, shardNode.Pos())
-            .Value(service)
-            .Done();
-
-        auto dqSoShardBuilder = Build<TSoShard>(ctx, shardNode.Pos());
-        dqSoShardBuilder.SolomonCluster(solomonClusterAtom);
-        dqSoShardBuilder.Project(projectAtom);
-        dqSoShardBuilder.Cluster(clusterAtom);
-        dqSoShardBuilder.Service(serviceAtom);
-
-        dqSoShardBuilder.Token<TCoSecureParam>().Name().Build("cluster:default_" + solomonCluster).Build();
-
-        return dqSoShardBuilder.Done();
     }
 
 private: