YQL-18343: Propagate EnableSpilling flag to compute nodes.

Author: Darych

ydb/core/kqp/executer_actor/kqp_literal_executer.cpp

@@ -173,6 +173,7 @@ class TKqpLiteralExecuter {
         NDqProto::TDqTask protoTask;
         protoTask.SetId(task.Id);
         protoTask.SetStageId(task.StageId.StageId);
+        protoTask.SetEnableSpilling(false); // TODO: enable spilling
         protoTask.MutableProgram()->CopyFrom(stage.GetProgram()); // it's not good...
 
         TaskId2StageId[task.Id] = task.StageId.StageId;

ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp

@@ -1141,6 +1141,7 @@ void SerializeTaskToProto(const TKqpTasksGraph& tasksGraph, const TTask& task, N
     result->SetId(task.Id);
     result->SetStageId(stageInfo.Id.StageId);
     result->SetUseLlvm(task.GetUseLlvm());
+    result->SetEnableSpilling(false); // TODO: enable spilling
     if (task.HasMetaId()) {
         result->SetMetaId(task.GetMetaIdUnsafe());
     }

ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp

@@ -49,7 +49,7 @@ class TLocalTaskRunnerActor
         , TaskId(taskId)
         , InputChannelsWithDisabledCheckpoints(std::move(inputChannelsWithDisabledCheckpoints))
         , MemoryQuota(std::move(memoryQuota))
-    { 
+    {
     }
 
     ~TLocalTaskRunnerActor()
@@ -427,16 +427,18 @@ class TLocalTaskRunnerActor
         }
 
         TaskRunner->Prepare(settings, ev->Get()->MemoryLimits, *ev->Get()->ExecCtx);
-        
+
         THashMap<ui64, std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> inputTransforms;
         for (auto i = 0; i != inputs.size(); ++i) {
             if (auto t = TaskRunner->GetInputTransform(i)) {
                 inputTransforms[i] = *t;
             }
         }
 
-        auto wakeUpCallback = ev->Get()->ExecCtx->GetWakeupCallback();
-        TaskRunner->SetSpillerFactory(std::make_shared<TDqSpillerFactory>(TxId, NActors::TActivationContext::ActorSystem(), wakeUpCallback));
+        if (settings.GetEnableSpilling()) {
+            auto wakeUpCallback = ev->Get()->ExecCtx->GetWakeupCallback();
+            TaskRunner->SetSpillerFactory(std::make_shared<TDqSpillerFactory>(TxId, NActors::TActivationContext::ActorSystem(), wakeUpCallback));
+        }
 
         auto event = MakeHolder<TEvTaskRunnerCreateFinished>(
             TaskRunner->GetSecureParams(),

ydb/library/yql/dq/proto/dq_tasks.proto

@@ -202,4 +202,5 @@ message TDqTask {
     optional bool UseLlvm = 16;
     repeated bytes ReadRanges = 17;
     map<string, string> RequestContext = 18;
+    optional bool EnableSpilling = 19;
 }

ydb/library/yql/dq/runtime/dq_tasks_runner.h

@@ -347,6 +347,10 @@ class TDqTaskSettings {
         return Task_->GetRequestContext();
     }
 
+    bool GetEnableSpilling() const {
+        return Task_->HasEnableSpilling() && Task_->GetEnableSpilling();
+    }
+
 private:
 
     // external callback to retrieve parameter value.
@@ -405,8 +409,8 @@ class IDqTaskRunner : public TSimpleRefCount<IDqTaskRunner>, private TNonCopyabl
 };
 
 TIntrusivePtr<IDqTaskRunner> MakeDqTaskRunner(
-    NKikimr::NMiniKQL::TScopedAlloc& alloc, 
-    const TDqTaskRunnerContext& ctx, 
+    NKikimr::NMiniKQL::TScopedAlloc& alloc,
+    const TDqTaskRunnerContext& ctx,
     const TDqTaskRunnerSettings& settings,
     const TLogFunc& logFunc
 );

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

@@ -443,7 +443,7 @@ using TBase = TComputationValue<TSpillingSupportState<Sort, HasCount>>;
             switch (InputStatus = Flow->FetchValues(ctx, GetFields())) {
                 case EFetchResult::One:
                     if (Put()) {
-                        if (!HasMemoryForProcessing()) {
+                        if (ctx.SpillerFactory && !HasMemoryForProcessing()) {
                             SwitchMode(EOperatingMode::Spilling, ctx);
                             return EFetchResult::Yield;
                         }
@@ -562,9 +562,7 @@ using TBase = TComputationValue<TSpillingSupportState<Sort, HasCount>>;
     EOperatingMode GetMode() const { return Mode; }
 
     bool HasMemoryForProcessing() const {
-        // TODO: Change to enable spilling
-        // return !TlsAllocState->IsMemoryYellowZoneEnabled();
-        return true;
+        return !TlsAllocState->IsMemoryYellowZoneEnabled();
     }
 
     bool IsReadFromChannelFinished() const {

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

@@ -449,6 +449,7 @@ namespace NYql::NDqs {
             taskMeta.SetStageId(publicId);
             taskDesc.MutableMeta()->PackFrom(taskMeta);
             taskDesc.SetStageId(stageId);
+            taskDesc.SetEnableSpilling(Settings->IsSpillingEnabled());
 
             if (Settings->DisableLLVMForBlockStages.Get().GetOrElse(true)) {
                 auto& stage = TasksGraph.GetStageInfo(task.StageId).Meta.Stage;