review fixes

Author: lll-phill-lll

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

@@ -78,9 +78,11 @@ TString BuildProgram(NNodes::TCoLambda program, const TStructExprType& paramsTyp
     TRuntimeNode rootNode = MkqlBuildExpr(program.Body().Ref(), ctx);
 
     TExploringNodeVisitor explorer;
-    explorer.Walk(rootNode.GetNode(), typeEnv);
-    bool wereChanges = false;
-    rootNode = SinglePassVisitCallables(rootNode, explorer, TSpillingTransformProvider(spillingSettings), typeEnv, true, wereChanges);
+    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);

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

@@ -10,6 +10,10 @@
 namespace NYql::NDq {
 
 struct TSpillingSettings {
+    operator bool() const {
+        return EnableSpillingInGraceJoin;
+    }
+    
     bool EnableSpillingInGraceJoin = false;
 };
 

ydb/library/yql/minikql/comp_nodes/mkql_grace_join.cpp

@@ -1170,23 +1170,14 @@ IComputationNode* WrapGraceJoinCommon(TCallable& callable, const TComputationNod
 
 IComputationNode* WrapGraceJoin(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
     MKQL_ENSURE(callable.GetInputsCount() == 8, "Expected 8 args");
-    bool isSpillingAllowed = false;
-    if (callable.GetType()->GetName() == "GraceJoinWithSpilling") {
-        isSpillingAllowed = true;
-    }
 
-    return WrapGraceJoinCommon(callable, ctx, false, isSpillingAllowed);
+    return WrapGraceJoinCommon(callable, ctx, false, callable.IsSpillingSupported());
 }
 
 IComputationNode* WrapGraceSelfJoin(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
     MKQL_ENSURE(callable.GetInputsCount() == 7, "Expected 7 args");
-
-    bool isSpillingAllowed = false;
-    if (callable.GetType()->GetName() == "GraceSelfJoinWithSpilling") {
-        isSpillingAllowed = true;
-    }
 
-    return WrapGraceJoinCommon(callable, ctx, true, isSpillingAllowed);
+    return WrapGraceJoinCommon(callable, ctx, true, callable.IsSpillingSupported());
 }
 
 }

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);

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

@@ -98,7 +98,7 @@ TDqConfiguration::TDqConfiguration() {
 
     REGISTER_SETTING(*this, _MaxAttachmentsSize);
     REGISTER_SETTING(*this, DisableCheckpoints);
-    REGISTER_SETTING(*this, EnableSpillingInGraceJoin);
+    REGISTER_SETTING(*this, EnabledSpillingNodes).Parser([](const TString& v) { return FromString<EEnabledSpillingNodes>(v); });;
 }
 
 } // namespace NYql

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,7 +66,7 @@ struct TDqSettings {
         static constexpr ui32 MaxDPccpDPTableSize = 16400U;
         static constexpr ui64 MaxAttachmentsSize = 2_GB;
         static constexpr bool SplitStageOnDqReplicate = true;
-        static constexpr bool EnableSpillingInGraceJoin = false;
+        static constexpr EEnabledSpillingNodes EnabledSpillingNodes = EEnabledSpillingNodes::None;
     };
 
     using TPtr = std::shared_ptr<TDqSettings>;
@@ -132,7 +138,7 @@ struct TDqSettings {
     NCommon::TConfSetting<bool, false> DisableLLVMForBlockStages;
     NCommon::TConfSetting<bool, false> SplitStageOnDqReplicate;
 
-    NCommon::TConfSetting<bool, false> EnableSpillingInGraceJoin;
+    NCommon::TConfSetting<EEnabledSpillingNodes, false> EnabledSpillingNodes;
 
     NCommon::TConfSetting<ui64, false> _MaxAttachmentsSize;
     NCommon::TConfSetting<bool, false> DisableCheckpoints;
@@ -219,7 +225,8 @@ struct TDqSettings {
     }
 
     bool IsSpillingInGraceJoinEnabled() const {
-        return IsSpillingEnabled() && EnableSpillingInGraceJoin.Get().GetOrElse(false);
+        auto enabledNodes = EnabledSpillingNodes.Get().GetOrElse(TDqSettings::TDefault::EnabledSpillingNodes);
+        return IsSpillingEnabled() && (enabledNodes == EEnabledSpillingNodes::OnlyGraceJoin || enabledNodes == EEnabledSpillingNodes::All);
     }
 
     bool IsDqReplicateEnabled(const TTypeAnnotationContext& typesCtx) const {