fixed bug in S3 Physical Optimizer 2 (#13391)

Author: evanevanevanevannnn

ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp

@@ -2366,6 +2366,21 @@ Y_UNIT_TEST_SUITE(KqpFederatedQuery) {
 
         auto db = kikimr->GetQueryClient();
 
+        {
+            const TString sql = fmt::format(R"(
+                    INSERT INTO {destination}
+                        SELECT key, value FROM {source};)",
+                    "destination"_a = table1,
+                    "source"_a = olapTable);
+
+            auto scriptExecutionOperation = db.ExecuteScript(sql).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(scriptExecutionOperation.Status().GetStatus(), EStatus::SUCCESS, scriptExecutionOperation.Status().GetIssues().ToString());
+            UNIT_ASSERT(scriptExecutionOperation.Metadata().ExecutionId);
+
+            NYdb::NQuery::TScriptExecutionOperation readyOp = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), kikimr->GetDriver());
+            UNIT_ASSERT_EQUAL_C(readyOp.Metadata().ExecStatus, EExecStatus::Completed, readyOp.Status().GetIssues().ToString());
+        }
+
         {
             const TString sql = fmt::format(R"(
                     INSERT INTO {destination}
@@ -2381,6 +2396,21 @@ Y_UNIT_TEST_SUITE(KqpFederatedQuery) {
             UNIT_ASSERT_EQUAL_C(readyOp.Metadata().ExecStatus, EExecStatus::Completed, readyOp.Status().GetIssues().ToString());
         }
 
+        {
+            const TString sql = fmt::format(R"(
+                    INSERT INTO {destination}
+                        SELECT key, value, "2024" AS year FROM {source};)",
+                    "destination"_a = table2,
+                    "source"_a = olapTable);
+
+            auto scriptExecutionOperation = db.ExecuteScript(sql).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(scriptExecutionOperation.Status().GetStatus(), EStatus::SUCCESS, scriptExecutionOperation.Status().GetIssues().ToString());
+            UNIT_ASSERT(scriptExecutionOperation.Metadata().ExecutionId);
+
+            NYdb::NQuery::TScriptExecutionOperation readyOp = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), kikimr->GetDriver());
+            UNIT_ASSERT_EQUAL_C(readyOp.Metadata().ExecStatus, EExecStatus::Completed, readyOp.Status().GetIssues().ToString());
+        }
+
         {
             const TString sql = fmt::format(R"(
                     INSERT INTO {destination}

ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp

@@ -199,17 +199,31 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase {
             }
         }
 
-        if (!FindNode(input.Ptr(), [] (const TExprNode::TPtr& node) { return node->IsCallable(TDqCnUnionAll::CallableName()); })) {
+        if (!FindNode(input.Ptr(), [] (const TExprNode::TPtr& node) { return node->IsCallable(TCoDataSource::CallableName()); })) {
             YQL_CLOG(INFO, ProviderS3) << "Rewrite pure S3WriteObject `" << cluster << "`.`" << target.Path().StringValue() << "` as stage with sink.";
+            auto shouldBePassedAsInput = FindNode(input.Ptr(), [] (const TExprNode::TPtr& node) { return node->IsCallable(TDqStage::CallableName()); });
+
+            auto stageInputs = Build<TExprList>(ctx, writePos);
+            auto toFlow = Build<TCoToFlow>(ctx, writePos);
+            TVector<TCoArgument> args;
+
+            if (shouldBePassedAsInput) {
+                auto arg = Build<TCoArgument>(ctx, writePos).Name("in").Done();
+                stageInputs.Add(input);
+                args.push_back(arg);
+                toFlow.Input(arg);
+            }
+            else {
+                toFlow.Input(input);
+            }
+
             return keys.empty() ?
                 Build<TDqStage>(ctx, writePos)
-                    .Inputs().Build()
+                    .Inputs(stageInputs.Done())
                     .Program<TCoLambda>()
-                        .Args({})
+                        .Args(args)
                         .Body<TS3SinkOutput>()
-                            .Input<TCoToFlow>()
-                                .Input(input)
-                                .Build()
+                            .Input(toFlow.Done())
                             .Format(target.Format())
                             .KeyColumns().Build()
                             .Settings(sinkOutputSettingsBuilder.Done())
@@ -237,12 +251,10 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase {
                             .Add<TDqCnHashShuffle>()
                                 .Output<TDqOutput>()
                                     .Stage<TDqStage>()
-                                        .Inputs().Build()
+                                        .Inputs(stageInputs.Done())
                                         .Program<TCoLambda>()
-                                            .Args({})
-                                            .Body<TCoToFlow>()
-                                                .Input(input)
-                                                .Build()
+                                            .Args(args)
+                                            .Body(toFlow.Done())
                                             .Build()
                                         .Settings().Build()
                                         .Build()