Merge 621daf0 into 4d68f9b

lll-phill-lll · web-flow · commit 2bcae57bb1fb · 2024-07-04T12:03:08.000Z
diff --git a/ydb/core/kqp/provider/yql_kikimr_exec.cpp b/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());
diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/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);
diff --git a/ydb/library/yql/dq/tasks/dq_task_program.cpp b/ydb/library/yql/dq/tasks/dq_task_program.cpp
@@ -2,13 +2,41 @@
 
 #include <ydb/library/yql/core/yql_expr_optimize.h>
 #include <ydb/library/yql/minikql/mkql_node_serialization.h>
+#include <ydb/library/yql/minikql/mkql_runtime_version.h>
 #include <ydb/library/yql/providers/common/mkql/yql_type_mkql.h>
 
 namespace NYql::NDq {
 
 using namespace NKikimr::NMiniKQL;
 using namespace NYql::NNodes;
 
+
+class TSpillingTransformProvider {
+public:
+
+    TSpillingTransformProvider(const TSpillingSettings& spillingSettings): SpillingSettings(spillingSettings){};
+
+    TCallableVisitFunc operator()(TInternName name) {
+        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 = 0; i < callable.GetInputsCount(); ++i) {
+                    callableBuilder.Add(callable.GetInput(i));
+                }
+                return TRuntimeNode(callableBuilder.Build(), false);
+            };
+        }
+
+        return TCallableVisitFunc();
+    }
+
+private:
+
+    TSpillingSettings SpillingSettings;
+};
+
 const TStructExprType* CollectParameters(NNodes::TCoLambda program, TExprContext& ctx) {
     TVector<const TItemExprType*> memberTypes;
 
@@ -27,7 +55,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);
 
@@ -49,6 +77,13 @@ TString BuildProgram(NNodes::TCoLambda program, const TStructExprType& paramsTyp
 
     TRuntimeNode rootNode = MkqlBuildExpr(program.Body().Ref(), ctx);
 
+    TExploringNodeVisitor explorer;
+    if (spillingSettings) {
+        explorer.Walk(rootNode.GetNode(), typeEnv);
+        bool wereChanges = false;
+        rootNode = SinglePassVisitCallables(rootNode, explorer, TSpillingTransformProvider(spillingSettings), typeEnv, true, wereChanges);
+    }
+
     TStructLiteralBuilder structBuilder(typeEnv);
     structBuilder.Add("Program", rootNode);
     structBuilder.Add("Inputs", pgmBuilder.NewTuple(inputNodes));
@@ -64,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()) {
diff --git a/ydb/library/yql/dq/tasks/dq_task_program.h b/ydb/library/yql/dq/tasks/dq_task_program.h
@@ -9,11 +9,19 @@
 
 namespace NYql::NDq {
 
+struct TSpillingSettings {
+    operator bool() const {
+        return EnableSpillingInGraceJoin;
+    }
+    
+    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
diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_factory.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_factory.cpp
@@ -234,8 +234,8 @@ struct TCallableComputationNodeBuilderFuncMapFiller {
         {"JoinDict", &WrapJoinDict},
         {"GraceJoin", &WrapGraceJoin},
         {"GraceSelfJoin", &WrapGraceSelfJoin},
-        {"GraceJoinWithSpilling", &WrapGraceJoinWithSpilling},
-        {"GraceSelfJoinWithSpilling", &WrapGraceSelfJoinWithSpilling},
+        {"GraceJoinWithSpilling", &WrapGraceJoin},
+        {"GraceSelfJoinWithSpilling", &WrapGraceSelfJoin},
         {"MapJoinCore", &WrapMapJoinCore},
         {"CommonJoinCore", &WrapCommonJoinCore},
         {"CombineCore", &WrapCombineCore},
diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_grace_join.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_grace_join.cpp
@@ -1171,28 +1171,15 @@ IComputationNode* WrapGraceJoinCommon(TCallable& callable, const TComputationNod
 IComputationNode* WrapGraceJoin(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
     MKQL_ENSURE(callable.GetInputsCount() == 8, "Expected 8 args");
 
-    return WrapGraceJoinCommon(callable, ctx, false, false);
+    return WrapGraceJoinCommon(callable, ctx, false, callable.IsSpillingSupported());
 }
 
 IComputationNode* WrapGraceSelfJoin(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
     MKQL_ENSURE(callable.GetInputsCount() == 7, "Expected 7 args");
     
-    return WrapGraceJoinCommon(callable, ctx, true, false);
+    return WrapGraceJoinCommon(callable, ctx, true, callable.IsSpillingSupported());
 }
 
-IComputationNode* WrapGraceJoinWithSpilling(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
-    MKQL_ENSURE(callable.GetInputsCount() == 8, "Expected 8 args");
-
-    return WrapGraceJoinCommon(callable, ctx, false, true);
-}
-
-IComputationNode* WrapGraceSelfJoinWithSpilling(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
-    MKQL_ENSURE(callable.GetInputsCount() == 7, "Expected 7 args");
-    
-    return WrapGraceJoinCommon(callable, ctx, true, true);
-}
-
-
 }
 
 }
diff --git a/ydb/library/yql/minikql/mkql_node.h b/ydb/library/yql/minikql/mkql_node.h
@@ -1075,6 +1075,10 @@ friend class TNode;
         UniqueId = uniqueId;
     }
 
+    bool IsSpillingSupported() const {
+        return TStringBuf(GetType()->GetName()).EndsWith("WithSpilling"_sb);
+    }
+
 private:
     TCallable(ui32 inputsCount, TRuntimeNode* inputs, TCallableType* type, bool validate = true);
     TCallable(TRuntimeNode result, TCallableType* type, bool validate = true);
diff --git a/ydb/library/yql/minikql/mkql_program_builder.cpp b/ydb/library/yql/minikql/mkql_program_builder.cpp
@@ -2160,28 +2160,6 @@ TRuntimeNode TProgramBuilder::GraceSelfJoin(TRuntimeNode flowLeft,  EJoinKind jo
     return GraceJoinCommon(__func__, flowLeft, {}, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings);
 }
 
-TRuntimeNode TProgramBuilder::GraceJoinWithSpilling(TRuntimeNode flowLeft, TRuntimeNode flowRight, EJoinKind joinKind,
-        const TArrayRef<const ui32>& leftKeyColumns, const TArrayRef<const ui32>& rightKeyColumns,
-        const TArrayRef<const ui32>& leftRenames, const TArrayRef<const ui32>& rightRenames, TType* returnType, EAnyJoinSettings anyJoinSettings ) {
-
-    if constexpr (RuntimeVersion < 50U) {
-        THROW yexception() << "Runtime version (" << RuntimeVersion << ") too old for " << __func__;
-    }
-
-    return GraceJoinCommon(__func__, flowLeft, flowRight, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings);
-}
-
-TRuntimeNode TProgramBuilder::GraceSelfJoinWithSpilling(TRuntimeNode flowLeft, EJoinKind joinKind,
-        const TArrayRef<const ui32>& leftKeyColumns, const TArrayRef<const ui32>& rightKeyColumns,
-        const TArrayRef<const ui32>& leftRenames, const TArrayRef<const ui32>& rightRenames, TType* returnType, EAnyJoinSettings anyJoinSettings ) {
-
-    if constexpr (RuntimeVersion < 50U) {
-        THROW yexception() << "Runtime version (" << RuntimeVersion << ") too old for " << __func__;
-    }
-    
-    return GraceJoinCommon(__func__, flowLeft, {}, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings);
-}
-
 TRuntimeNode TProgramBuilder::ToSortedDict(TRuntimeNode list, bool all, const TUnaryLambda& keySelector,
     const TUnaryLambda& payloadSelector, bool isCompact, ui64 itemsCountHint) {
     return ToDict(list, all, keySelector, payloadSelector, __func__, isCompact, itemsCountHint);
diff --git a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
@@ -1704,13 +1704,6 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
 
         const auto returnType = BuildType(node, *node.GetTypeAnn(), ctx.ProgramBuilder);
 
-        // TODO: use PRAGMA
-        bool IsSpillingAllowed = false;
-        if (RuntimeVersion >= 50U && IsSpillingAllowed) {
-            return selfJoin
-                ? ctx.ProgramBuilder.GraceSelfJoinWithSpilling(flowLeft, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings)
-                : ctx.ProgramBuilder.GraceJoinWithSpilling(flowLeft, flowRight, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings);
-        }
         return selfJoin
             ? ctx.ProgramBuilder.GraceSelfJoin(flowLeft, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings)
             : ctx.ProgramBuilder.GraceJoin(flowLeft, flowRight, joinKind, leftKeyColumns, rightKeyColumns, leftRenames, rightRenames, returnType, anyJoinSettings);
diff --git a/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp b/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp
@@ -98,6 +98,7 @@ TDqConfiguration::TDqConfiguration() {
 
     REGISTER_SETTING(*this, _MaxAttachmentsSize);
     REGISTER_SETTING(*this, DisableCheckpoints);
+    REGISTER_SETTING(*this, EnabledSpillingNodes).Parser([](const TString& v) { return FromString<EEnabledSpillingNodes>(v); });;
 }
 
 } // namespace NYql
diff --git a/ydb/library/yql/providers/dq/common/yql_dq_settings.h b/ydb/library/yql/providers/dq/common/yql_dq_settings.h
@@ -28,6 +28,12 @@ struct TDqSettings {
         File        /* "file" */,
     };
 
+    enum class EEnabledSpillingNodes {
+        None,
+        OnlyGraceJoin,
+        All,
+    };
+
     struct TDefault {
         static constexpr ui32 MaxTasksPerStage = 20U;
         static constexpr ui32 MaxTasksPerOperation = 70U;
@@ -60,6 +66,7 @@ struct TDqSettings {
         static constexpr ui32 MaxDPccpDPTableSize = 16400U;
         static constexpr ui64 MaxAttachmentsSize = 2_GB;
         static constexpr bool SplitStageOnDqReplicate = true;
+        static constexpr EEnabledSpillingNodes EnabledSpillingNodes = EEnabledSpillingNodes::None;
     };
 
     using TPtr = std::shared_ptr<TDqSettings>;
@@ -131,6 +138,8 @@ struct TDqSettings {
     NCommon::TConfSetting<bool, false> DisableLLVMForBlockStages;
     NCommon::TConfSetting<bool, false> SplitStageOnDqReplicate;
 
+    NCommon::TConfSetting<EEnabledSpillingNodes, false> EnabledSpillingNodes;
+
     NCommon::TConfSetting<ui64, false> _MaxAttachmentsSize;
     NCommon::TConfSetting<bool, false> DisableCheckpoints;
 
@@ -215,6 +224,11 @@ struct TDqSettings {
         return SpillingEngine.Get().GetOrElse(TDqSettings::TDefault::SpillingEngine) != ESpillingEngine::Disable;
     }
 
+    bool IsSpillingInGraceJoinEnabled() const {
+        auto enabledNodes = EnabledSpillingNodes.Get().GetOrElse(TDqSettings::TDefault::EnabledSpillingNodes);
+        return IsSpillingEnabled() && (enabledNodes == EEnabledSpillingNodes::OnlyGraceJoin || enabledNodes == EEnabledSpillingNodes::All);
+    }
+
     bool IsDqReplicateEnabled(const TTypeAnnotationContext& typesCtx) const {
         return EnableDqReplicate.Get().GetOrElse(
             typesCtx.BlockEngineMode != EBlockEngineMode::Disable || TDqSettings::TDefault::EnableDqReplicate);
diff --git a/ydb/library/yql/providers/dq/planner/execution_planner.cpp b/ydb/library/yql/providers/dq/planner/execution_planner.cpp
@@ -676,10 +676,11 @@ 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);
         }
     }
diff --git a/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp b/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp
@@ -760,9 +760,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");

Original file line number	Diff line number	Diff line change
`@@ -98,6 +98,7 @@ TDqConfiguration::TDqConfiguration() {`
`98`	`98`
`99`	`99`	`REGISTER_SETTING(*this, _MaxAttachmentsSize);`
`100`	`100`	`REGISTER_SETTING(*this, DisableCheckpoints);`
	`101`	`+ REGISTER_SETTING(*this, EnabledSpillingNodes).Parser([](const TString& v) { return FromString<EEnabledSpillingNodes>(v); });;`
`101`	`102`	`}`
`102`	`103`
`103`	`104`	`} // namespace NYql`