add TSpillingSettings

Author: lll-phill-lll

ydb/core/kqp/provider/yql_kikimr_exec.cpp

@@ -764,7 +764,7 @@ class TKiSourceCallableExecutionTransformer : public TAsyncCallbackTransformer<T
         lambda = NDq::BuildProgram(
             programLambda, *paramsType, compiler, SessionCtx->Query().QueryData->GetAllocState()->TypeEnv,
                 *SessionCtx->Query().QueryData->GetAllocState()->HolderFactory.GetFunctionRegistry(),
-                ctx, fakeReads);
+                ctx, fakeReads, {});
 
         NKikimr::NMiniKQL::TProgramBuilder programBuilder(SessionCtx->Query().QueryData->GetAllocState()->TypeEnv,
             *SessionCtx->Query().QueryData->GetAllocState()->HolderFactory.GetFunctionRegistry());

ydb/core/kqp/query_compiler/kqp_query_compiler.cpp

@@ -778,7 +778,7 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
 
         auto paramsType = CollectParameters(stage, ctx);
         auto programBytecode = NDq::BuildProgram(stage.Program(), *paramsType, *KqlCompiler, TypeEnv, FuncRegistry,
-            ctx, {});
+            ctx, {}, {});
 
         auto& programProto = *stageProto.MutableProgram();
         programProto.SetRuntimeVersion(NYql::NDqProto::ERuntimeVersion::RUNTIME_VERSION_YQL_1_0);

ydb/library/yql/dq/tasks/dq_task_program.cpp

@@ -5,8 +5,6 @@
 #include <ydb/library/yql/minikql/mkql_runtime_version.h>
 #include <ydb/library/yql/providers/common/mkql/yql_type_mkql.h>
 
-#include <util/generic/xrange.h>
-
 namespace NYql::NDq {
 
 using namespace NKikimr::NMiniKQL;
@@ -15,13 +13,17 @@ using namespace NYql::NNodes;
 
 class TSpillingTransformProvider {
 public:
+
+    TSpillingTransformProvider(const TSpillingSettings& spillingSettings): SpillingSettings(spillingSettings){};
+
+
     TCallableVisitFunc operator()(TInternName name) {
-        if (name == "GraceJoin" || name == "GraceSelfJoin") {
+        if (RuntimeVersion >= 50U && SpillingSettings.EnableSpillingInGraceJoin && (name == "GraceJoin" || name == "GraceSelfJoin")) {
             return [name](NKikimr::NMiniKQL::TCallable& callable, const TTypeEnvironment& env) {
                 TCallableBuilder callableBuilder(env,
                     TStringBuilder() << callable.GetType()->GetName() << "WithSpilling",
                     callable.GetType()->GetReturnType(), false);
-                for (ui32 i: xrange(callable.GetInputsCount())) {
+                for (ui32 i = 0; i < callable.GetInputsCount(); ++i) {
                     callableBuilder.Add(callable.GetInput(i));
                 }
                 return TRuntimeNode(callableBuilder.Build(), false);
@@ -31,6 +33,10 @@ class TSpillingTransformProvider {
 
         return TCallableVisitFunc();
     }
+
+private:
+
+    TSpillingSettings SpillingSettings;
 };
 
 const TStructExprType* CollectParameters(NNodes::TCoLambda program, TExprContext& ctx) {
@@ -51,7 +57,7 @@ const TStructExprType* CollectParameters(NNodes::TCoLambda program, TExprContext
 
 TString BuildProgram(NNodes::TCoLambda program, const TStructExprType& paramsType,
                      const NCommon::IMkqlCallableCompiler& compiler, const TTypeEnvironment& typeEnv,
-                     const IFunctionRegistry& funcRegistry, TExprContext& exprCtx, const TVector<TExprBase>& reads)
+                     const IFunctionRegistry& funcRegistry, TExprContext& exprCtx, const TVector<TExprBase>& reads, const TSpillingSettings& spillingSettings)
 {
     TProgramBuilder pgmBuilder(typeEnv, funcRegistry);
 
@@ -73,14 +79,10 @@ TString BuildProgram(NNodes::TCoLambda program, const TStructExprType& paramsTyp
 
     TRuntimeNode rootNode = MkqlBuildExpr(program.Body().Ref(), ctx);
 
-    if (RuntimeVersion >= 50U) {
-        TExploringNodeVisitor explorer;
-        explorer.Walk(rootNode.GetNode(), typeEnv);
-        bool wereChanges = false;
-        rootNode = SinglePassVisitCallables(rootNode, explorer, TSpillingTransformProvider(), typeEnv, true, wereChanges);
-
-        std::cerr << "MISHA " << wereChanges << std::endl; 
-    }
+    TExploringNodeVisitor explorer;
+    explorer.Walk(rootNode.GetNode(), typeEnv);
+    bool wereChanges = false;
+    rootNode = SinglePassVisitCallables(rootNode, explorer, TSpillingTransformProvider(spillingSettings), typeEnv, true, wereChanges);
 
     TStructLiteralBuilder structBuilder(typeEnv);
     structBuilder.Add("Program", rootNode);
@@ -97,7 +99,6 @@ TString BuildProgram(NNodes::TCoLambda program, const TStructExprType& paramsTyp
 
     auto programNode = structBuilder.Build();
 
-    TExploringNodeVisitor explorer;
     explorer.Walk(programNode, typeEnv);
     ui32 uniqueId = 0;
     for (auto& node : explorer.GetNodes()) {

ydb/library/yql/dq/tasks/dq_task_program.h

@@ -9,11 +9,15 @@
 
 namespace NYql::NDq {
 
+struct TSpillingSettings {
+    bool EnableSpillingInGraceJoin = false;
+};
+
 const TStructExprType* CollectParameters(NNodes::TCoLambda program, TExprContext& ctx);
 
 TString BuildProgram(NNodes::TCoLambda program, const TStructExprType& paramsType,
                      const NCommon::IMkqlCallableCompiler& compiler, const NKikimr::NMiniKQL::TTypeEnvironment& typeEnv,
                      const NKikimr::NMiniKQL::IFunctionRegistry& funcRegistry, TExprContext& exprCtx,
-                     const TVector<NNodes::TExprBase>& reads);
+                     const TVector<NNodes::TExprBase>& reads, const TSpillingSettings& spillingSettings);
 
 } // namespace NYql::NDq

ydb/library/yql/providers/dq/common/yql_dq_settings.h

@@ -218,6 +218,10 @@ struct TDqSettings {
         return SpillingEngine.Get().GetOrElse(TDqSettings::TDefault::SpillingEngine) != ESpillingEngine::Disable;
     }
 
+    bool IsSpillingInGraceJoinEnabled() const {
+        return IsSpillingEnabled() && EnableSpillingInGraceJoin.Get().GetOrElse(false);
+    }
+
     bool IsDqReplicateEnabled(const TTypeAnnotationContext& typesCtx) const {
         return EnableDqReplicate.Get().GetOrElse(
             typesCtx.BlockEngineMode != EBlockEngineMode::Disable || TDqSettings::TDefault::EnableDqReplicate);

ydb/library/yql/providers/dq/planner/execution_planner.cpp

@@ -1,6 +1,7 @@
 #include "execution_planner.h"
 
 #include <ydb/library/yql/dq/integration/yql_dq_integration.h>
+#include <ydb/library/yql/minikql/mkql_runtime_version.h>
 #include <ydb/library/yql/providers/dq/expr_nodes/dqs_expr_nodes.h>
 #include <ydb/library/yql/providers/dq/opt/dqs_opt.h>
 #include <ydb/library/yql/providers/dq/opt/logical_optimize.h>
@@ -676,10 +677,12 @@ namespace NYql::NDqs {
                 Y_ABORT_UNLESS(false);
             }
 */
+
+            TSpillingSettings spillingSettings{Settings->IsSpillingInGraceJoinEnabled()};
             StagePrograms[stageInfo.first] = std::make_tuple(
                 NDq::BuildProgram(
                     stage.Program(), *paramsType, compiler, typeEnv, *FunctionRegistry,
-                    ExprContext, fakeReads),
+                    ExprContext, fakeReads, spillingSettings),
                 stageId, publicId);
         }
     }

ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp

@@ -761,9 +761,10 @@ class TDqExecTransformer: public TExecTransformerBase, TCounters
 
             TVector<TExprBase> fakeReads;
             auto paramsType = NDq::CollectParameters(programLambda, ctx);
+            NDq::TSpillingSettings spillingSettings{State->Settings->IsSpillingInGraceJoinEnabled()};
             *lambda = NDq::BuildProgram(
                 programLambda, *paramsType, compiler, typeEnv, *State->FunctionRegistry,
-                ctx, fakeReads);
+                ctx, fakeReads, spillingSettings);
         }
 
         auto block = MeasureBlock("RuntimeNodeVisitor");