merge to yq stable YQ-4053 support parallel s3 import (#14175)

Author: GrigoriyPA

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

@@ -143,14 +143,15 @@ Y_UNIT_TEST_SUITE(KqpS3PlanTest) {
         UNIT_ASSERT(NJson::ReadJsonTree(*queryResult.GetStats()->GetPlan(), &plan));
 
         const auto& writeStagePlan = plan["Plan"]["Plans"][0]["Plans"][0];
-        UNIT_ASSERT_VALUES_EQUAL(writeStagePlan["Node Type"].GetStringSafe(), "Limit-Sink");
-        UNIT_ASSERT(writeStagePlan["Operators"].GetArraySafe().size() >= 2);
-        const auto& sinkOp = writeStagePlan["Operators"].GetArraySafe()[1];
+        UNIT_ASSERT_VALUES_EQUAL(writeStagePlan["Node Type"].GetStringSafe(), "Stage-Sink");
+        UNIT_ASSERT(writeStagePlan["Operators"].GetArraySafe().size() >= 1);
+        const auto& sinkOp = writeStagePlan["Operators"].GetArraySafe()[0];
         UNIT_ASSERT_VALUES_EQUAL(sinkOp["ExternalDataSource"].GetStringSafe(), "write_data_source");
         UNIT_ASSERT_VALUES_EQUAL(sinkOp["Compression"].GetStringSafe(), "gzip");
 
-        const auto& readStagePlan = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0];
+        const auto& readStagePlan = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0];
         UNIT_ASSERT_VALUES_EQUAL(readStagePlan["Node Type"].GetStringSafe(), "Source");
+        UNIT_ASSERT(readStagePlan["Operators"].GetArraySafe().size() >= 1);
         const auto& sourceOp = readStagePlan["Operators"].GetArraySafe()[0];
         UNIT_ASSERT_VALUES_EQUAL(sourceOp["ExternalDataSource"].GetStringSafe(), "read_data_source");
 

ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp

@@ -573,15 +573,15 @@ class TS3WriteActor : public TActorBootstrapped<TS3WriteActor>, public IDqComput
             const auto& key = MakePartitionKey(row);
             const auto [keyIt, insertedNew] = FileWriteActors.emplace(key, std::vector<TS3FileWriteActor*>());
             if (insertedNew || keyIt->second.empty() || keyIt->second.back()->IsFinishing()) {
-            auto fileWrite = std::make_unique<TS3FileWriteActor>(
-                TxId,
-                Gateway,
-                Credentials,
-                key,
-                NS3Util::UrlEscapeRet(Url + Path + key + MakeOutputName() + Extension),
-                Compression,
-                RetryPolicy, DirtyWrite, Token);
-            keyIt->second.emplace_back(fileWrite.get());
+                auto fileWrite = std::make_unique<TS3FileWriteActor>(
+                    TxId,
+                    Gateway,
+                    Credentials,
+                    key,
+                    NS3Util::UrlEscapeRet(Url + Path + key + MakeOutputName() + Extension),
+                    Compression,
+                    RetryPolicy, DirtyWrite, Token);
+                keyIt->second.emplace_back(fileWrite.get());
                 RegisterWithSameMailbox(fileWrite.release());
             }
 
@@ -619,6 +619,10 @@ class TS3WriteActor : public TActorBootstrapped<TS3WriteActor>, public IDqComput
         NDqProto::StatusIds::StatusCode statusCode = result->Get()->StatusCode;
         if (statusCode == NDqProto::StatusIds::UNSPECIFIED) {
             statusCode = StatusFromS3ErrorCode(result->Get()->S3ErrorCode);
+            if (statusCode == NDqProto::StatusIds::UNSPECIFIED) {
+                statusCode = NDqProto::StatusIds::INTERNAL_ERROR;
+                result->Get()->Issues.AddIssue("Got upload error with unspecified error code.");
+            }
         }
 
         Callbacks->OnAsyncOutputError(OutputIndex, result->Get()->Issues, statusCode);
@@ -640,10 +644,15 @@ class TS3WriteActor : public TActorBootstrapped<TS3WriteActor>, public IDqComput
             if (const auto ft = std::find_if(it->second.cbegin(), it->second.cend(), [&](TS3FileWriteActor* actor){ return result->Get()->Url == actor->GetUrl(); }); it->second.cend() != ft) {
                 (*ft)->PassAway();
                 it->second.erase(ft);
-                if (it->second.empty())
+                if (it->second.empty()) {
                     FileWriteActors.erase(it);
+                }
             }
         }
+        if (!Finished && GetFreeSpace() > 0) {
+            LOG_D("TS3WriteActor", "Has free space, notify owner");
+            Callbacks->ResumeExecution();
+        }
         FinishIfNeeded();
     }
 

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

@@ -124,13 +124,15 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         auto keys = GetPartitionKeys(partBy);
 
         auto sinkSettingsBuilder = Build<TExprList>(ctx, target.Pos());
-        if (partBy)
+        if (partBy) {
             sinkSettingsBuilder.Add(std::move(partBy));
+        }
 
         auto compression = GetCompression(settings);
         const auto& extension = GetExtension(target.Format().Value(), compression ? compression->Tail().Content() : ""sv);
-        if (compression)
+        if (compression) {
             sinkSettingsBuilder.Add(std::move(compression));
+        }
 
         auto sinkOutputSettingsBuilder = Build<TExprList>(ctx, target.Pos());
         if (auto csvDelimiter = GetCsvDelimiter(settings)) {
@@ -199,7 +201,7 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase {
             }
         }
 
-        if (!FindNode(input.Ptr(), [] (const TExprNode::TPtr& node) { return node->IsCallable(TCoDataSource::CallableName()); })) {
+        if (IsDqPureExpr(input)) {
             YQL_CLOG(INFO, ProviderS3) << "Rewrite pure S3WriteObject `" << cluster << "`.`" << target.Path().StringValue() << "` as stage with sink.";
             return keys.empty() ?
                 Build<TDqStage>(ctx, writePos)
@@ -305,23 +307,26 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase {
                 .Build()
             .Done();
 
-        auto outputsBuilder = Build<TDqStageOutputsList>(ctx, target.Pos());
-        if (inputStage.Outputs() && keys.empty()) {
-            outputsBuilder.InitFrom(inputStage.Outputs().Cast());
-        }
-        outputsBuilder.Add(sink);
+        auto outputsBuilder = Build<TDqStageOutputsList>(ctx, target.Pos())
+            .Add(sink);
 
         if (keys.empty()) {
-            const auto outputBuilder = Build<TS3SinkOutput>(ctx, target.Pos())
-                .Input(inputStage.Program().Body().Ptr())
-                .Format(target.Format())
-                .KeyColumns().Add(std::move(keys)).Build()
-                .Settings(sinkOutputSettingsBuilder.Done())
-                .Done();
-
             return Build<TDqStage>(ctx, writePos)
-                .InitFrom(inputStage)
-                .Program(ctx.DeepCopyLambda(inputStage.Program().Ref(), outputBuilder.Ptr()))
+                .Inputs()
+                    .Add<TDqCnMap>()
+                        .Output(dqUnion.Output())
+                        .Build()
+                    .Build()
+                .Program<TCoLambda>()
+                    .Args({"in"})
+                    .Body<TS3SinkOutput>()
+                        .Input("in")
+                        .Format(target.Format())
+                        .KeyColumns().Add(std::move(keys)).Build()
+                        .Settings(sinkOutputSettingsBuilder.Done())
+                        .Build()
+                    .Build()
+                .Settings().Build()
                 .Outputs(outputsBuilder.Done())
                 .Done();
         } else {