support ydb pragma MaxTasksPerStage

Added MaxTasksPerStage setting.

ydb/core/kqp/host/kqp_runner.cpp

@@ -307,7 +307,7 @@ private:
         auto& preparedQuery = *TransformCtx->QueryCtx->PreparingQuery;
         TKqpPhysicalQuery physicalQuery(transformedQuery);
 
-        auto compiler = CreateKqpQueryCompiler(Cluster, OptimizeCtx->Tables, FuncRegistry, TypesCtx);
+        auto compiler = CreateKqpQueryCompiler(Cluster, OptimizeCtx->Tables, FuncRegistry, TypesCtx, Config);
         auto ret = compiler->CompilePhysicalQuery(physicalQuery, dataQueryBlocks, *preparedQuery.MutablePhysicalQuery(), ctx);
         if (!ret) {
             ctx.AddError(TIssue(ctx.GetPosition(query->Pos()), "Failed to compile physical query."));

ydb/core/kqp/provider/yql_kikimr_settings.cpp

@@ -65,6 +65,8 @@ TKikimrConfiguration::TKikimrConfiguration() {
     REGISTER_SETTING(*this, OptUseFinalizeByKey);
     REGISTER_SETTING(*this, OptEnableCostBasedOptimization);
 
+    REGISTER_SETTING(*this, MaxTasksPerStage);
+
     /* Runtime */
     REGISTER_SETTING(*this, ScanQuery);
 }

ydb/core/kqp/provider/yql_kikimr_settings.h

@@ -58,6 +58,8 @@ struct TKikimrSettings {
     NCommon::TConfSetting<bool, false> OptUseFinalizeByKey;
     NCommon::TConfSetting<bool, false> OptEnableCostBasedOptimization;
 
+    NCommon::TConfSetting<ui32, false> MaxTasksPerStage;
+
     /* Runtime */
     NCommon::TConfSetting<bool, true> ScanQuery;
 

ydb/core/kqp/query_compiler/kqp_query_compiler.cpp

@@ -411,7 +411,7 @@ void FillOlapProgram(const TCoLambda& process, const NKikimr::NMiniKQL::TType* m
 class TKqpQueryCompiler : public IKqpQueryCompiler {
 public:
     TKqpQueryCompiler(const TString& cluster, const TIntrusivePtr<TKikimrTablesData> tablesData,
-        const NMiniKQL::IFunctionRegistry& funcRegistry, TTypeAnnotationContext& typesCtx)
+        const NMiniKQL::IFunctionRegistry& funcRegistry, TTypeAnnotationContext& typesCtx, NYql::TKikimrConfiguration::TPtr config)
         : Cluster(cluster)
         , TablesData(tablesData)
         , FuncRegistry(funcRegistry)
@@ -420,6 +420,7 @@ public:
         , KqlCtx(cluster, tablesData, TypeEnv, FuncRegistry)
         , KqlCompiler(CreateKqlCompiler(KqlCtx, typesCtx))
         , TypesCtx(typesCtx)
+        , Config(config)
     {
         Alloc.Release();
     }
@@ -928,11 +929,8 @@ private:
             // In runtime, number of tasks with Sources is limited by 2x of node count
             // We prepare a lot of partitions and distribute them between these tasks
             // Constraint of 1 task per partition is NOT valid anymore
-            // We choose 120 as number with a lot of divisors for even final distribution
-            //
-            // TODO: Replace with ydb.MaxTasksPerStage when implemented
-            //
-            dqIntegration->Partition(NYql::TDqSettings(), 120, source.Ref(), partitionParams, &clusterName, ctx, false);
+            auto maxTasksPerStage = Config->MaxTasksPerStage.Get().GetOrElse(TDqSettings::TDefault::MaxTasksPerStage);
+            dqIntegration->Partition(NYql::TDqSettings(), maxTasksPerStage, source.Ref(), partitionParams, &clusterName, ctx, false);
             externalSource.SetTaskParamKey(TString(dataSourceCategory));
             for (const TString& partitionParam : partitionParams) {
                 externalSource.AddPartitionedTaskParams(partitionParam);
@@ -1146,6 +1144,7 @@ private:
     TKqlCompileContext KqlCtx;
     TIntrusivePtr<NCommon::IMkqlCallableCompiler> KqlCompiler;
     TTypeAnnotationContext& TypesCtx;
+    TKikimrConfiguration::TPtr Config;
     TSet<TString> SecretNames;
 };
 
@@ -1153,9 +1152,9 @@ private:
 
 TIntrusivePtr<IKqpQueryCompiler> CreateKqpQueryCompiler(const TString& cluster,
     const TIntrusivePtr<TKikimrTablesData> tablesData, const IFunctionRegistry& funcRegistry,
-    TTypeAnnotationContext& typesCtx)
+    TTypeAnnotationContext& typesCtx, NYql::TKikimrConfiguration::TPtr config)
 {
-    return MakeIntrusive<TKqpQueryCompiler>(cluster, tablesData, funcRegistry, typesCtx);
+    return MakeIntrusive<TKqpQueryCompiler>(cluster, tablesData, funcRegistry, typesCtx, config);
 }
 
 } // namespace NKqp

ydb/core/kqp/query_compiler/kqp_query_compiler.h

@@ -2,6 +2,7 @@
 
 #include <ydb/core/kqp/expr_nodes/kqp_expr_nodes.h>
 #include <ydb/core/kqp/provider/yql_kikimr_expr_nodes.h>
+#include <ydb/core/kqp/provider/yql_kikimr_settings.h>
 #include <ydb/core/protos/kqp_physical.pb.h>
 
 #include <ydb/core/kqp/provider/yql_kikimr_provider.h>
@@ -18,7 +19,7 @@ public:
 
 TIntrusivePtr<IKqpQueryCompiler> CreateKqpQueryCompiler(const TString& cluster,
     const TIntrusivePtr<NYql::TKikimrTablesData> tablesData, const NMiniKQL::IFunctionRegistry& funcRegistry,
-    NYql::TTypeAnnotationContext& typesCtx);
+    NYql::TTypeAnnotationContext& typesCtx, NYql::TKikimrConfiguration::TPtr config);
 
 } // namespace NKqp
 } // namespace NKikimr