Pragma ydb.OverridePlanner to tune runtime execution (#6904)

Hor911 · web-flow · commit 370928159494 · 2024-07-24T18:00:27.000+03:00
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -962,7 +962,12 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
 
         ui32 taskCount = externalSource.GetPartitionedTaskParams().size();
 
-        if (!resourceSnapshot.empty()) {
+        auto taskCountHint = stage.GetTaskCount();
+        if (taskCountHint) {
+            if (taskCount > taskCountHint) {
+                taskCount = taskCountHint;
+            }
+        } else if (!resourceSnapshot.empty()) {
             ui32 maxTaskcount = resourceSnapshot.size() * 2;
             if (taskCount > maxTaskcount) {
                 taskCount = maxTaskcount;
@@ -1342,7 +1347,11 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         }
 
         if (isShuffle) {
-            partitionsCount = std::max(partitionsCount, GetMaxTasksAggregation(stageInfo, inputTasks, nodesCount));
+            if (stage.GetTaskCount()) {
+                partitionsCount = stage.GetTaskCount();
+            } else {
+                partitionsCount = std::max(partitionsCount, GetMaxTasksAggregation(stageInfo, inputTasks, nodesCount));
+            }
         }
 
         for (ui32 i = 0; i < partitionsCount; ++i) {
diff --git a/ydb/core/kqp/opt/kqp_query_plan.cpp b/ydb/core/kqp/opt/kqp_query_plan.cpp
@@ -2551,6 +2551,7 @@ TString AddExecStatsToTxPlan(const TString& txPlanJson, const NYql::NDqProto::TD
                     stats["UseLlvm"] = "undefined";
                 }
 
+                stats["PhysicalStageId"] = (*stat)->GetStageId();
                 stats["Tasks"] = (*stat)->GetTotalTasksCount();
 
                 stats["StageDurationUs"] = (*stat)->GetStageDurationUs();
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.cpp b/ydb/core/kqp/provider/yql_kikimr_settings.cpp
@@ -67,6 +67,7 @@ TKikimrConfiguration::TKikimrConfiguration() {
     REGISTER_SETTING(*this, OptEnableOlapPushdown);
     REGISTER_SETTING(*this, OptEnableOlapProvideComputeSharding);
     REGISTER_SETTING(*this, OverrideStatistics);
+    REGISTER_SETTING(*this, OverridePlanner);
 
 
     REGISTER_SETTING(*this, OptUseFinalizeByKey);
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.h b/ydb/core/kqp/provider/yql_kikimr_settings.h
@@ -52,6 +52,7 @@ struct TKikimrSettings {
     NCommon::TConfSetting<NDq::EHashJoinMode, false> HashJoinMode;
     NCommon::TConfSetting<TString, false> OverrideStatistics;
     NCommon::TConfSetting<ui64, false> EnableSpillingNodes;
+    NCommon::TConfSetting<TString, false> OverridePlanner;
 
     /* Disable optimizer rules */
     NCommon::TConfSetting<bool, false> OptDisableTopSort;
diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
@@ -515,6 +515,33 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             CompileTransaction(tx, *queryProto.AddTransactions(), ctx);
         }
 
+        auto overridePlanner = Config->OverridePlanner.Get();
+        if (overridePlanner) {
+            NJson::TJsonReaderConfig jsonConfig;
+            NJson::TJsonValue jsonNode;
+            if (NJson::ReadJsonTree(*overridePlanner, &jsonConfig, &jsonNode)) {
+                for (auto& stageOverride : jsonNode.GetArray()) {
+                    ui32 txId = 0;
+                    if (auto* txNode = stageOverride.GetValueByPath("tx")) {
+                        txId = txNode->GetIntegerSafe();
+                    }
+                    if (txId < static_cast<ui32>(queryProto.GetTransactions().size())) {
+                        auto& tx = *queryProto.MutableTransactions(txId);
+                        ui32 stageId = 0;
+                        if (auto* stageNode = stageOverride.GetValueByPath("stage")) {
+                            stageId = stageNode->GetIntegerSafe();
+                        }
+                        if (stageId < static_cast<ui32>(tx.GetStages().size())) {
+                            auto& stage = *tx.MutableStages(stageId);
+                            if (auto* tasksNode = stageOverride.GetValueByPath("tasks")) {
+                                stage.SetTaskCount(tasksNode->GetIntegerSafe());
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
         for (ui32 i = 0; i < query.Results().Size(); ++i) {
             const auto& result = query.Results().Item(i);
 
diff --git a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp
@@ -1961,6 +1961,175 @@ Y_UNIT_TEST_SUITE(KqpFederatedQuery) {
 
         ValidateTables(client, oltpTable, olapTable);
     }
+
+    Y_UNIT_TEST(OverridePlannerDefaults) {
+        const TString root = "/Root/";
+        const TString source = "source";
+        const TString table1 = "table1";
+        const TString table2 = "table2";
+        const TString bucket  = "bucket";
+        const TString object1 = "object1";
+        const TString object2 = "object2";
+        const TString content1 = "foo,bar\naaa,0\nbbb,2";
+        const TString content2 = "foo,bar\naaa,1\nbbb,3";
+
+        Aws::S3::S3Client s3Client = MakeS3Client();
+        CreateBucket(bucket, s3Client);
+        UploadObject(bucket, table1 + "/" + object1, content1, s3Client);
+        UploadObject(bucket, table1 + "/" + object2, content2, s3Client);
+        UploadObject(bucket, table2 + "/" + object1, content1, s3Client);
+        UploadObject(bucket, table2 + "/" + object2, content2, s3Client);
+
+        auto kikimr = NTestUtils::MakeKikimrRunner();
+
+        auto tc = kikimr->GetTableClient();
+        auto session = tc.CreateSession().GetValueSync().GetSession();
+        const TString query = fmt::format(R"(
+            CREATE EXTERNAL DATA SOURCE `{source}` WITH (
+                SOURCE_TYPE="ObjectStorage",
+                LOCATION="{location}",
+                AUTH_METHOD="NONE"
+            );
+            CREATE EXTERNAL TABLE `{table1}` (
+                foo STRING NOT NULL,
+                bar UINT32 NOT NULL
+            ) WITH (
+                DATA_SOURCE="{source}",
+                LOCATION="/{location_table1}/",
+                FORMAT="csv_with_names"
+            );
+            CREATE EXTERNAL TABLE `{table2}` (
+                foo STRING NOT NULL,
+                bar UINT32 NOT NULL
+            ) WITH (
+                DATA_SOURCE="{source}",
+                LOCATION="/{location_table2}/",
+                FORMAT="csv_with_names"
+            );
+            )",
+            "source"_a = root + source,
+            "table1"_a = root + table1,
+            "table2"_a = root + table2,
+            "location_table1"_a = table1,
+            "location_table2"_a = table2,
+            "location"_a = TStringBuilder() << GetEnv("S3_ENDPOINT") << '/' << bucket
+            );
+
+        auto result = session.ExecuteSchemeQuery(query).GetValueSync();
+        UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+
+        ui32 source1_id = 0;
+        ui32 source2_id = 0;
+        ui32 join_id = 0;
+        ui32 limit_id = 0;
+        auto queryClient = kikimr->GetQueryClient();
+
+        {
+            // default planner values
+
+            const TString sql = fmt::format(R"(
+                    SELECT SUM(t1.bar + t2.bar) as sum FROM `{table1}` as t1 JOIN /*+grace()*/ `{table2}`as t2 ON t1.foo = t2.foo
+                )",
+                "table1"_a = root + table1,
+                "table2"_a = root + table2);
+
+            TExecuteQueryResult queryResult = queryClient.ExecuteQuery(
+                sql,
+                TTxControl::BeginTx().CommitTx(),
+                TExecuteQuerySettings().ExecMode(EExecMode::Execute).StatsMode(EStatsMode::Full)).GetValueSync();
+
+            UNIT_ASSERT_C(queryResult.IsSuccess(), queryResult.GetIssues().ToString());
+            UNIT_ASSERT(queryResult.GetStats());
+            UNIT_ASSERT(queryResult.GetStats()->GetPlan());
+            NJson::TJsonValue plan;
+            UNIT_ASSERT(NJson::ReadJsonTree(*queryResult.GetStats()->GetPlan(), &plan));
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][1]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 4);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 1);
+
+            source1_id = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["PhysicalStageId"].GetIntegerSafe();
+            source2_id = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][1]["Plans"][0]["Stats"]["PhysicalStageId"].GetIntegerSafe();
+            join_id    = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["PhysicalStageId"].GetIntegerSafe();
+            limit_id   = plan["Plan"]["Plans"][0]["Plans"][0]["Stats"]["PhysicalStageId"].GetIntegerSafe();
+        }
+
+        {
+            // scale down
+
+            const TString sql = fmt::format(R"(
+                    pragma ydb.OverridePlanner = @@ [
+                        {{ "tx": 0, "stage": {source1_id}, "tasks": 1 }},
+                        {{ "tx": 0, "stage": {source2_id}, "tasks": 1 }},
+                        {{ "tx": 0, "stage": {join_id}, "tasks": 1 }},
+                        {{ "tx": 0, "stage": {limit_id}, "tasks": 1 }}
+                    ] @@;
+
+                    SELECT SUM(t1.bar + t2.bar) as sum FROM `{table1}` as t1 JOIN /*+grace()*/ `{table2}`as t2 ON t1.foo = t2.foo
+                )",
+                "source1_id"_a = source1_id,
+                "source2_id"_a = source2_id,
+                "join_id"_a = join_id,
+                "limit_id"_a = limit_id,
+                "table1"_a = root + table1,
+                "table2"_a = root + table2);
+
+            TExecuteQueryResult queryResult = queryClient.ExecuteQuery(
+                sql,
+                TTxControl::BeginTx().CommitTx(),
+                TExecuteQuerySettings().ExecMode(EExecMode::Execute).StatsMode(EStatsMode::Full)).GetValueSync();
+
+            UNIT_ASSERT_C(queryResult.IsSuccess(), queryResult.GetIssues().ToString());
+            UNIT_ASSERT(queryResult.GetStats());
+            UNIT_ASSERT(queryResult.GetStats()->GetPlan());
+            NJson::TJsonValue plan;
+            UNIT_ASSERT(NJson::ReadJsonTree(*queryResult.GetStats()->GetPlan(), &plan));
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][1]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 1);
+        }
+
+        {
+            // scale up
+
+            const TString sql = fmt::format(R"(
+                    pragma ydb.OverridePlanner = @@ [
+                        {{ "tx": 0, "stage": {source1_id}, "tasks": 10 }},
+                        {{ "tx": 0, "stage": {source2_id}, "tasks": 10 }},
+                        {{ "tx": 0, "stage": {join_id}, "tasks": 10 }},
+                        {{ "tx": 0, "stage": {limit_id}, "tasks": 10 }}
+                    ] @@;
+
+                    SELECT SUM(t1.bar + t2.bar) as sum FROM `{table1}` as t1 JOIN /*+grace()*/ `{table2}`as t2 ON t1.foo = t2.foo
+                )",
+                "source1_id"_a = source1_id,
+                "source2_id"_a = source2_id,
+                "join_id"_a = join_id,
+                "limit_id"_a = limit_id,
+                "table1"_a = root + table1,
+                "table2"_a = root + table2);
+
+            TExecuteQueryResult queryResult = queryClient.ExecuteQuery(
+                sql,
+                TTxControl::BeginTx().CommitTx(),
+                TExecuteQuerySettings().ExecMode(EExecMode::Execute).StatsMode(EStatsMode::Full)).GetValueSync();
+
+            UNIT_ASSERT_C(queryResult.IsSuccess(), queryResult.GetIssues().ToString());
+            UNIT_ASSERT(queryResult.GetStats());
+            UNIT_ASSERT(queryResult.GetStats()->GetPlan());
+            NJson::TJsonValue plan;
+            UNIT_ASSERT(NJson::ReadJsonTree(*queryResult.GetStats()->GetPlan(), &plan));
+            // only 2 files => sources stay with 2 tasks
+            // join scales to 10 tasks
+            // limit ignores hint and keeps being in the only task
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][1]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 10);
+            UNIT_ASSERT_VALUES_EQUAL(plan["Plan"]["Plans"][0]["Plans"][0]["Stats"]["Tasks"].GetIntegerSafe(), 1);
+        }
+    }
+
 }
 
 } // namespace NKikimr::NKqp
diff --git a/ydb/core/protos/kqp_physical.proto b/ydb/core/protos/kqp_physical.proto
@@ -376,6 +376,7 @@ message TKqpPhyStage {
     repeated TKqpSink Sinks = 11;
     map<string, string> SecureParams = 12;
     bool AllowWithSpilling = 13;
+    uint32 TaskCount = 14;
 }
 
 message TKqpPhyResult {

Original file line number	Diff line number	Diff line change
`@@ -2551,6 +2551,7 @@ TString AddExecStatsToTxPlan(const TString& txPlanJson, const NYql::NDqProto::TD`
`2551`	`2551`	`stats["UseLlvm"] = "undefined";`
`2552`	`2552`	`}`
`2553`	`2553`
	`2554`	`+ stats["PhysicalStageId"] = (*stat)->GetStageId();`
`2554`	`2555`	`stats["Tasks"] = (*stat)->GetTotalTasksCount();`
`2555`	`2556`
`2556`	`2557`	`stats["StageDurationUs"] = (*stat)->GetStageDurationUs();`
Original file line number	Diff line number	Diff line change
`@@ -376,6 +376,7 @@ message TKqpPhyStage {`
`376`	`376`	`repeated TKqpSink Sinks = 11;`
`377`	`377`	`map<string, string> SecureParams = 12;`
`378`	`378`	`bool AllowWithSpilling = 13;`
	`379`	`+ uint32 TaskCount = 14;`
`379`	`380`	`}`
`380`	`381`
`381`	`382`	`message TKqpPhyResult {`