Better scheduler for S3 Sources (#9620)

Hor911 · web-flow · commit 54b47ab3d49a · 2024-09-25T15:21:36.000+03:00
diff --git a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
@@ -1858,6 +1858,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
         size_t sourceScanPartitionsCount = 0;
         for (ui32 txIdx = 0; txIdx < Request.Transactions.size(); ++txIdx) {
             auto& tx = Request.Transactions[txIdx];
+            auto scheduledTaskCount = ScheduleByCost(tx, ResourceSnapshot);
             for (ui32 stageIdx = 0; stageIdx < tx.Body->StagesSize(); ++stageIdx) {
                 auto& stage = tx.Body->GetStages(stageIdx);
                 auto& stageInfo = TasksGraph.GetStageInfo(TStageId(txIdx, stageIdx));
@@ -1903,8 +1904,10 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                                 UnknownAffectedShardCount = true;
                             }
                             break;
-                        case NKqpProto::TKqpSource::kExternalSource:
-                            BuildReadTasksFromSource(stageInfo, ResourceSnapshot);
+                        case NKqpProto::TKqpSource::kExternalSource: {
+                            auto it = scheduledTaskCount.find(stageIdx);
+                            BuildReadTasksFromSource(stageInfo, ResourceSnapshot, it != scheduledTaskCount.end() ? it->second.TaskCount : 0);
+                        }
                             break;
                         default:
                             YQL_ENSURE(false, "unknown source type");
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -81,6 +81,10 @@ struct TShardRangesWithShardId {
     const TShardKeyRanges* Ranges;
 };
 
+struct TStageScheduleInfo {
+    double StageCost = 0.0;
+    ui32 TaskCount = 0;
+};
 
 TActorId ReportToRl(ui64 ru, const TString& database, const TString& userToken,
     const NKikimrKqp::TRlPath& path);
@@ -819,6 +823,40 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         }
     }
 
+    std::map<ui32, TStageScheduleInfo> ScheduleByCost(const IKqpGateway::TPhysicalTxData& tx, const TVector<NKikimrKqp::TKqpNodeResources>& resourceSnapshot) {
+        std::map<ui32, TStageScheduleInfo> result;
+        if (!resourceSnapshot.empty()) // can't schedule w/o node count
+        {
+            // collect costs and schedule stages with external sources only 
+            double totalCost = 0.0;
+            for (ui32 stageIdx = 0; stageIdx < tx.Body->StagesSize(); ++stageIdx) {
+                auto& stage = tx.Body->GetStages(stageIdx);
+                if (stage.SourcesSize() > 0 && stage.GetSources(0).GetTypeCase() == NKqpProto::TKqpSource::kExternalSource) {
+                    if (stage.GetStageCost() > 0.0 && stage.GetTaskCount() == 0) {
+                        totalCost += stage.GetStageCost();
+                        result.emplace(stageIdx, TStageScheduleInfo{.StageCost = stage.GetStageCost()});
+                    }
+                }
+            }
+            // assign task counts
+            if (!result.empty()) {
+                // allow use 2/3 of threads in single stage
+                ui32 maxStageTaskCount = (TStagePredictor::GetUsableThreads() * 2 + 2) / 3;
+                // total limit per mode is x2 
+                ui32 maxTotalTaskCount = maxStageTaskCount * 2;
+                for (auto& [_, stageInfo] : result) {
+                    // schedule tasks evenly between nodes
+                    stageInfo.TaskCount = 
+                        std::max<ui32>(
+                            std::min(static_cast<ui32>(maxTotalTaskCount * stageInfo.StageCost / totalCost), maxStageTaskCount)
+                            , 1
+                        ) * resourceSnapshot.size();
+                }
+            }
+        }
+        return result;
+    }
+
     void BuildSysViewScanTasks(TStageInfo& stageInfo) {
         Y_DEBUG_ABORT_UNLESS(stageInfo.Meta.IsSysView());
 
@@ -924,7 +962,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         }
     }
 
-    void BuildReadTasksFromSource(TStageInfo& stageInfo, const TVector<NKikimrKqp::TKqpNodeResources>& resourceSnapshot) {
+    void BuildReadTasksFromSource(TStageInfo& stageInfo, const TVector<NKikimrKqp::TKqpNodeResources>& resourceSnapshot, ui32 scheduledTaskCount) {
         const auto& stage = stageInfo.Meta.GetStage(stageInfo.Id);
 
         YQL_ENSURE(stage.GetSources(0).HasExternalSource());
@@ -936,6 +974,10 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         ui32 taskCount = externalSource.GetPartitionedTaskParams().size();
 
         auto taskCountHint = stage.GetTaskCount();
+        if (taskCountHint == 0) {
+            taskCountHint = scheduledTaskCount;
+        }
+
         if (taskCountHint) {
             if (taskCount > taskCountHint) {
                 taskCount = taskCountHint;
diff --git a/ydb/core/kqp/opt/kqp_statistics_transformer.cpp b/ydb/core/kqp/opt/kqp_statistics_transformer.cpp
@@ -615,6 +615,8 @@ void InferStatisticsForDqSourceWrap(const TExprNode::TPtr& input, TTypeAnnotatio
                     }
                 }
 
+                auto dataSourceStats = stats;
+
                 auto rowType = wrapBase.Cast().RowType().Ref().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
                 if (specific->FullRawRowAvgSize == 0.0) {
                     auto newSpecific = std::make_shared<TS3ProviderStatistics>(*specific);
@@ -659,6 +661,10 @@ void InferStatisticsForDqSourceWrap(const TExprNode::TPtr& input, TTypeAnnotatio
                         if (newSpecific->Compression) {
                             stats->Cost *= 1.5;
                         }
+                        {
+                            auto specific = const_cast<TS3ProviderStatistics*>(dynamic_cast<const TS3ProviderStatistics*>((dataSourceStats->Specific.get())));
+                            specific->Costs[TStructExprType::MakeHash(rowType->GetItems())] = stats->Cost;
+                        }
                     }
                     typeCtx->SetStats(input.Get(), stats);
                 }
diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
@@ -21,8 +21,10 @@
 #include <ydb/library/yql/providers/common/provider/yql_provider_names.h>
 #include <ydb/library/yql/providers/common/structured_token/yql_token_builder.h>
 #include <ydb/library/yql/providers/dq/common/yql_dq_settings.h>
+#include <ydb/library/yql/providers/s3/statistics/yql_s3_statistics.h>
 #include <ydb/library/yql/core/yql_opt_utils.h>
 
+
 namespace NKikimr {
 namespace NKqp {
 
@@ -672,8 +674,11 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             }
         }
 
+        double stageCost = 0.0;
         VisitExpr(stage.Program().Ptr(), [&](const TExprNode::TPtr& exprNode) {
+
             TExprBase node(exprNode);
+
             if (auto maybeReadTable = node.Maybe<TKqpWideReadTable>()) {
                 auto readTable = maybeReadTable.Cast();
                 auto tableMeta = TablesData->ExistingTable(Cluster, readTable.Table().Path()).Metadata;
@@ -754,12 +759,15 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
                 FillOlapProgram(readTableRanges, miniKqlResultType, *tableMeta, *tableOp.MutableReadOlapRange(), ctx);
                 FillResultType(miniKqlResultType, *tableOp.MutableReadOlapRange());
                 tableOp.MutableReadOlapRange()->SetReadType(NKqpProto::TKqpPhyOpReadOlapRanges::BLOCKS);
+            } else if (auto maybeDqSourceWrapBase = node.Maybe<TDqSourceWrapBase>()) {
+                stageCost += GetDqSourceWrapBaseCost(maybeDqSourceWrapBase.Cast(), TypesCtx);
             } else {
                 YQL_ENSURE(!node.Maybe<TKqpReadTable>());
             }
             return true;
         });
 
+        stageProto.SetStageCost(stageCost);
         const auto& secureParams = FindSecureParams(stage.Program().Ptr(), TypesCtx, SecretNames);
         stageProto.MutableSecureParams()->insert(secureParams.begin(), secureParams.end());
 
diff --git a/ydb/core/kqp/query_compiler/ya.make b/ydb/core/kqp/query_compiler/ya.make
@@ -19,6 +19,7 @@ PEERDIR(
     ydb/library/yql/minikql
     ydb/library/yql/providers/common/mkql
     ydb/library/yql/providers/dq/common
+    ydb/library/yql/providers/s3/expr_nodes
 )
 
 YQL_LAST_ABI_VERSION()
diff --git a/ydb/core/protos/kqp_physical.proto b/ydb/core/protos/kqp_physical.proto
@@ -377,6 +377,7 @@ message TKqpPhyStage {
     map<string, string> SecureParams = 12;
     bool AllowWithSpilling = 13;
     uint32 TaskCount = 14;
+    double StageCost = 15;
 }
 
 message TKqpPhyResult {
diff --git a/ydb/library/yql/providers/s3/statistics/ya.make b/ydb/library/yql/providers/s3/statistics/ya.make
@@ -1,10 +1,13 @@
 LIBRARY()
 
 SRCS(
+    yql_s3_statistics.cpp
 )
 
 PEERDIR(
     ydb/library/yql/core
+    ydb/library/yql/providers/dq/expr_nodes
+    ydb/library/yql/providers/s3/expr_nodes
 )
 
 YQL_LAST_ABI_VERSION()
diff --git a/ydb/library/yql/providers/s3/statistics/yql_s3_statistics.cpp b/ydb/library/yql/providers/s3/statistics/yql_s3_statistics.cpp
@@ -0,0 +1,24 @@
+#include <ydb/library/yql/providers/s3/expr_nodes/yql_s3_expr_nodes.h>
+
+#include "yql_s3_statistics.h"
+
+
+namespace NYql {
+
+double GetDqSourceWrapBaseCost(const NNodes::TDqSourceWrapBase& wrapBase, TTypeAnnotationContext& typeCtx) {
+    if (auto maybeS3DataSource = wrapBase.DataSource().Maybe<NNodes::TS3DataSource>()) {
+        auto s3DataSource = maybeS3DataSource.Cast();
+        auto stats = typeCtx.GetStats(s3DataSource.Raw());
+        if (stats && stats->Specific) {
+            const TS3ProviderStatistics* specific = dynamic_cast<const TS3ProviderStatistics*>((stats->Specific.get()));
+            auto rowType = wrapBase.RowType().Ref().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
+            auto it = specific->Costs.find(TStructExprType::MakeHash(rowType->GetItems()));
+            if (it != specific->Costs.end()) {
+                return it->second;
+            }
+        }
+    }
+    return 0.0;
+}
+
+} // namespace NYql
diff --git a/ydb/library/yql/providers/s3/statistics/yql_s3_statistics.h b/ydb/library/yql/providers/s3/statistics/yql_s3_statistics.h
@@ -1,6 +1,9 @@
 #pragma once
 
 #include <ydb/library/yql/core/yql_statistics.h>
+#include <ydb/library/yql/core/yql_type_annotation.h>
+#include <ydb/library/yql/providers/dq/expr_nodes/dqs_expr_nodes.h>
+#include <map>
 
 namespace NYql {
 
@@ -13,6 +16,9 @@ struct TS3ProviderStatistics : public IProviderStatistics {
     double PrunedDecodedRowAvgSize = 0.0;
     TString Format;
     TString Compression;
+    std::unordered_map<ui64, double> Costs;
 };
 
+double GetDqSourceWrapBaseCost(const NNodes::TDqSourceWrapBase& wrapBase, TTypeAnnotationContext& typeCtx);
+
 } // namespace NYql

Original file line number	Diff line number	Diff line change
`@@ -615,6 +615,8 @@ void InferStatisticsForDqSourceWrap(const TExprNode::TPtr& input, TTypeAnnotatio`
`615`	`615`	`}`
`616`	`616`	`}`
`617`	`617`
	`618`	`+ auto dataSourceStats = stats;`
	`619`	`+`
`618`	`620`	`auto rowType = wrapBase.Cast().RowType().Ref().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();`
`619`	`621`	`if (specific->FullRawRowAvgSize == 0.0) {`
`620`	`622`	`auto newSpecific = std::make_shared<TS3ProviderStatistics>(*specific);`
`@@ -659,6 +661,10 @@ void InferStatisticsForDqSourceWrap(const TExprNode::TPtr& input, TTypeAnnotatio`
`659`	`661`	`if (newSpecific->Compression) {`
`660`	`662`	`stats->Cost *= 1.5;`
`661`	`663`	`}`
	`664`	`+ {`
	`665`	`+ auto specific = const_cast<TS3ProviderStatistics>(dynamic_cast<const TS3ProviderStatistics>((dataSourceStats->Specific.get())));`
	`666`	`+ specific->Costs[TStructExprType::MakeHash(rowType->GetItems())] = stats->Cost;`
	`667`	`+ }`
`662`	`668`	`}`
`663`	`669`	`typeCtx->SetStats(input.Get(), stats);`
`664`	`670`	`}`
Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@ PEERDIR(`
`19`	`19`	`ydb/library/yql/minikql`
`20`	`20`	`ydb/library/yql/providers/common/mkql`
`21`	`21`	`ydb/library/yql/providers/dq/common`
	`22`	`+ ydb/library/yql/providers/s3/expr_nodes`
`22`	`23`	`)`
`23`	`24`
`24`	`25`	`YQL_LAST_ABI_VERSION()`
Original file line number	Diff line number	Diff line change
`@@ -377,6 +377,7 @@ message TKqpPhyStage {`
`377`	`377`	`map<string, string> SecureParams = 12;`
`378`	`378`	`bool AllowWithSpilling = 13;`
`379`	`379`	`uint32 TaskCount = 14;`
	`380`	`+ double StageCost = 15;`
`380`	`381`	`}`
`381`	`382`
`382`	`383`	`message TKqpPhyResult {`
Original file line number	Diff line number	Diff line change
`@@ -1,10 +1,13 @@`
`1`	`1`	`LIBRARY()`
`2`	`2`
`3`	`3`	`SRCS(`
	`4`	`+ yql_s3_statistics.cpp`
`4`	`5`	`)`
`5`	`6`
`6`	`7`	`PEERDIR(`
`7`	`8`	`ydb/library/yql/core`
	`9`	`+ ydb/library/yql/providers/dq/expr_nodes`
	`10`	`+ ydb/library/yql/providers/s3/expr_nodes`
`8`	`11`	`)`
`9`	`12`
`10`	`13`	`YQL_LAST_ABI_VERSION()`