add spilling nodes and pragma (ydb-platform#6769)

vporyadke · Jul 18, 2024 · 6de7340 · 6de7340
1 parent 31ea815
commit 6de7340
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 3 deletions.
diff --git a/ydb/core/kqp/provider/yql_kikimr_exec.cpp b/ydb/core/kqp/provider/yql_kikimr_exec.cpp
@@ -765,10 +765,11 @@ class TKiSourceCallableExecutionTransformer : public TAsyncCallbackTransformer<T
 
         TVector<TExprBase> fakeReads;
         auto paramsType = NDq::CollectParameters(programLambda, ctx);
+        NDq::TSpillingSettings spillingSettings{SessionCtx->Config().GetEnabledSpillingNodes()};
         lambda = NDq::BuildProgram(
             programLambda, *paramsType, compiler, SessionCtx->Query().QueryData->GetAllocState()->TypeEnv,
                 *SessionCtx->Query().QueryData->GetAllocState()->HolderFactory.GetFunctionRegistry(),
-                ctx, fakeReads, {});
+                ctx, fakeReads, spillingSettings);
 
         NKikimr::NMiniKQL::TProgramBuilder programBuilder(SessionCtx->Query().QueryData->GetAllocState()->TypeEnv,
             *SessionCtx->Query().QueryData->GetAllocState()->HolderFactory.GetFunctionRegistry());

diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.cpp b/ydb/core/kqp/provider/yql_kikimr_settings.cpp
@@ -3,6 +3,8 @@
 #include <ydb/core/protos/config.pb.h>
 #include <ydb/core/protos/table_service_config.pb.h>
 #include <util/generic/size_literals.h>
+#include <util/string/split.h>
+#include <ydb/library/yql/providers/dq/common/yql_dq_settings.h>
 
 namespace NYql {
 
@@ -69,6 +71,21 @@ TKikimrConfiguration::TKikimrConfiguration() {
 
     REGISTER_SETTING(*this, OptUseFinalizeByKey);
     REGISTER_SETTING(*this, CostBasedOptimizationLevel);
+    REGISTER_SETTING(*this, EnableSpillingNodes)
+        .Parser([](const TString& v) {
+            ui64 res = 0;
+            TVector<TString> vec;
+            StringSplitter(v).SplitBySet(",;| ").AddTo(&vec);
+            for (auto& s: vec) {
+                if (s.empty()) {
+                    throw yexception() << "Empty value item";
+                }
+                auto value = FromStringWithDefault<NYql::TDqSettings::EEnabledSpillingNodes>(
+                    s, NYql::TDqSettings::EEnabledSpillingNodes::None);
+                res |= ui64(value);
+            }
+            return res;
+        });
 
     REGISTER_SETTING(*this, MaxDPccpDPTableSize);
 
@@ -125,6 +142,10 @@ bool TKikimrSettings::HasOptUseFinalizeByKey() const {
     return GetOptionalFlagValue(OptUseFinalizeByKey.Get()) != EOptionalFlag::Disabled;
 }
 
+ui64 TKikimrSettings::GetEnabledSpillingNodes() const {
+    return EnableSpillingNodes.Get().GetOrElse(0);
+}
+
 
 EOptionalFlag TKikimrSettings::GetOptPredicateExtract() const {
     return GetOptionalFlagValue(OptEnablePredicateExtract.Get());

diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.h b/ydb/core/kqp/provider/yql_kikimr_settings.h
@@ -51,6 +51,7 @@ struct TKikimrSettings {
     NCommon::TConfSetting<bool, false> EnableLlvm;
     NCommon::TConfSetting<NDq::EHashJoinMode, false> HashJoinMode;
     NCommon::TConfSetting<TString, false> OverrideStatistics;
+    NCommon::TConfSetting<ui64, false> EnableSpillingNodes;
 
     /* Disable optimizer rules */
     NCommon::TConfSetting<bool, false> OptDisableTopSort;
@@ -82,6 +83,7 @@ struct TKikimrSettings {
     bool HasOptEnableOlapPushdown() const;
     bool HasOptEnableOlapProvideComputeSharding() const;
     bool HasOptUseFinalizeByKey() const;
+    ui64 GetEnabledSpillingNodes() const;
 
     EOptionalFlag GetOptPredicateExtract() const;
     EOptionalFlag GetUseLlvm() const;

diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
@@ -777,8 +777,9 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
         stageProto.SetIsEffectsStage(hasEffects || hasTxTableSink);
 
         auto paramsType = CollectParameters(stage, ctx);
+        NDq::TSpillingSettings spillingSettings{Config->GetEnabledSpillingNodes()};
         auto programBytecode = NDq::BuildProgram(stage.Program(), *paramsType, *KqlCompiler, TypeEnv, FuncRegistry,
-            ctx, {}, {});
+            ctx, {}, spillingSettings);
 
         auto& programProto = *stageProto.MutableProgram();
         programProto.SetRuntimeVersion(NYql::NDqProto::ERuntimeVersion::RUNTIME_VERSION_YQL_1_0);

diff --git a/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp b/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp
@@ -110,7 +110,7 @@ TDqConfiguration::TDqConfiguration() {
                 if (s.empty()) {
                     throw yexception() << "Empty value item";
                 }
-                auto value = FromString<EEnabledSpillingNodes>(s);
+                auto value = FromStringWithDefault<EEnabledSpillingNodes>(s, EEnabledSpillingNodes::None);
                 res |= ui64(value);
             }
             return res;

diff --git a/ydb/library/yql/providers/dq/common/yql_dq_settings.h b/ydb/library/yql/providers/dq/common/yql_dq_settings.h
@@ -29,6 +29,7 @@ struct TDqSettings {
     };
 
     enum class EEnabledSpillingNodes : ui64 {
+        None        = 0ULL      /* None */,
         GraceJoin   = 1ULL      /* "GraceJoin" */,
         All         = ~0ULL     /* "All" */,
     };