Merge 0ba4810 into fb4d6bb

Author: GrigoriyPA

ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp

@@ -721,14 +721,15 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         const auto& readRanges = ev->Get()->ReadRanges;
         const auto& typeEnv = ev->Get()->TypeEnv;
         const auto& holderFactory = ev->Get()->HolderFactory;
+        const auto& memoryLimits = ev->Get()->MemoryLimits;
         if (Stat) {
             Stat->AddCounters2(ev->Get()->Sensors);
         }
         TypeEnv = const_cast<NKikimr::NMiniKQL::TTypeEnvironment*>(&typeEnv);
         for (auto& [inputIndex, transform] : this->InputTransformsMap) {
             std::tie(transform.Input, transform.Buffer) = ev->Get()->InputTransforms.at(inputIndex);
         }
-        FillIoMaps(holderFactory, typeEnv, secureParams, taskParams, readRanges, nullptr);
+        FillIoMaps(holderFactory, typeEnv, secureParams, taskParams, readRanges, memoryLimits, nullptr);
 
         {
             // say "Hello" to executer

ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io.h

@@ -22,6 +22,7 @@ class TTaskOutput;
 namespace NYql::NDq {
 struct TSourceState;
 struct TSinkState;
+struct TDqTaskRunnerMemoryLimits;
 } // namespace NYql::NDq
 
 namespace NActors {
@@ -262,6 +263,7 @@ struct IDqAsyncIoFactory : public TThrRefBase {
         const google::protobuf::Message* SourceSettings = nullptr;  // used only in case if we execute compute actor locally
         TIntrusivePtr<NActors::TProtoArenaHolder> Arena;  // Arena for SourceSettings
         NWilson::TTraceId TraceId;
+        const TDqTaskRunnerMemoryLimits& MemoryLimits;
     };
 
     struct TLookupSourceArguments {

ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h

@@ -1338,6 +1338,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         const THashMap<TString, TString>& secureParams,
         const THashMap<TString, TString>& taskParams,
         const TVector<TString>& readRanges,
+        const TDqTaskRunnerMemoryLimits& memoryLimits,
         IRandomProvider* randomProvider
         )
     {
@@ -1371,7 +1372,8 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
                         .MemoryQuotaManager = MemoryLimits.MemoryQuotaManager,
                         .SourceSettings = (!settings.empty() ? settings.at(inputIndex) : nullptr),
                         .Arena = Task.GetArena(),
-                        .TraceId = ComputeActorSpan.GetTraceId()
+                        .TraceId = ComputeActorSpan.GetTraceId(),
+                        .MemoryLimits = memoryLimits
                     });
             } catch (const std::exception& ex) {
                 throw yexception() << "Failed to create source " << inputDesc.GetSource().GetType() << ": " << ex.what();

ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h

@@ -249,6 +249,7 @@ class TDqSyncComputeActorBase: public TDqComputeActorBase<TDerived, TComputeActo
             TaskRunner->GetSecureParams(),
             TaskRunner->GetTaskParams(),
             TaskRunner->GetReadRanges(),
+            TaskRunner->GetMemoryLimits(),
             TaskRunner->GetRandomProvider()
         );
     }

ydb/library/yql/dq/actors/task_runner/events.h

@@ -179,6 +179,7 @@ struct TEvTaskRunnerCreateFinished
         const NKikimr::NMiniKQL::THolderFactory& holderFactory,
         std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
         THashMap<ui64, std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>>&& inputTransforms,
+        const TDqTaskRunnerMemoryLimits& memoryLimits,
         const TTaskRunnerActorSensors& sensors = {}
     )
         : Sensors(sensors)
@@ -189,6 +190,7 @@ struct TEvTaskRunnerCreateFinished
         , HolderFactory(holderFactory)
         , Alloc(alloc)
         , InputTransforms(std::move(inputTransforms))
+        , MemoryLimits(memoryLimits)
     { 
         Y_ABORT_UNLESS(inputTransforms.empty() || Alloc);
     }
@@ -210,6 +212,7 @@ struct TEvTaskRunnerCreateFinished
     const NKikimr::NMiniKQL::THolderFactory& HolderFactory;
     std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
     THashMap<ui64, std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> InputTransforms; //can'not be const, because we need to explicitly clear it in destructor
+    const TDqTaskRunnerMemoryLimits MemoryLimits;
 };
 
 struct TEvTaskRunFinished

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

@@ -467,7 +467,8 @@ class TLocalTaskRunnerActor
             TaskRunner->GetTypeEnv(),
             TaskRunner->GetHolderFactory(),
             Alloc,
-            std::move(inputTransforms)
+            std::move(inputTransforms),
+            TaskRunner->GetMemoryLimits()
         );
 
         Send(

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

@@ -539,6 +539,7 @@ class TDqTaskRunner : public IDqTaskRunner {
     void Prepare(const TDqTaskSettings& task, const TDqTaskRunnerMemoryLimits& memoryLimits,
         const IDqTaskRunnerExecutionContext& execCtx) override
     {
+        MemoryLimits = memoryLimits;
         TaskId = task.GetId();
         auto entry = BuildTask(task);
 
@@ -901,6 +902,10 @@ class TDqTaskRunner : public IDqTaskRunner {
         return Stats.get();
     }
 
+    const TDqTaskRunnerMemoryLimits& GetMemoryLimits() const override {
+        return MemoryLimits;
+    }
+
     TString Save() const override {
         return AllocatedHolder->ProgramParsed.CompGraph->SaveGraphState();
     }
@@ -998,6 +1003,7 @@ class TDqTaskRunner : public IDqTaskRunner {
     ui64 TaskId = 0;
     TDqTaskRunnerContext Context;
     TDqTaskRunnerSettings Settings;
+    TDqTaskRunnerMemoryLimits MemoryLimits;
     TLogFunc LogFunc;
     std::unique_ptr<NUdf::ISecureParamsProvider> SecureParamsProvider;
     TDqTaskCountersProvider CountersProvider;

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

@@ -437,6 +437,7 @@ class IDqTaskRunner : public TSimpleRefCount<IDqTaskRunner>, private TNonCopyabl
     virtual const NKikimr::NMiniKQL::TTypeEnvironment& GetTypeEnv() const = 0;
     virtual const NKikimr::NMiniKQL::THolderFactory& GetHolderFactory() const = 0;
     virtual NKikimr::NMiniKQL::TScopedAlloc& GetAllocator() const = 0;
+    virtual const TDqTaskRunnerMemoryLimits& GetMemoryLimits() const = 0;
 
     virtual const THashMap<TString, TString>& GetSecureParams() const = 0;
     virtual const THashMap<TString, TString>& GetTaskParams() const = 0;

ydb/library/yql/providers/dq/actors/worker_actor.cpp

@@ -300,6 +300,7 @@ class TDqWorker: public TRichActor<TDqWorker>
             const auto& readRanges = ev->Get()->ReadRanges;
             const auto& typeEnv = ev->Get()->TypeEnv;
             const auto& holderFactory = ev->Get()->HolderFactory;
+            const auto& memoryLimits = ev->Get()->MemoryLimits;
 
             Stat.Measure<void>("PrepareChannels", [&](){
                 auto& inputs = Task.GetInputs();
@@ -323,7 +324,8 @@ class TDqWorker: public TRichActor<TDqWorker>
                                 .TypeEnv = typeEnv,
                                 .HolderFactory = holderFactory,
                                 .ProgramBuilder = *source.ProgramBuilder,
-                                .MemoryQuotaManager = MemoryQuotaManager
+                                .MemoryQuotaManager = MemoryQuotaManager,
+                                .MemoryLimits = memoryLimits
                             });
                         RegisterLocalChild(source.Actor);
                     } else {

ydb/library/yql/providers/dq/task_runner/tasks_runner_local.cpp

@@ -181,6 +181,10 @@ class TLocalTaskRunner: public ITaskRunner {
         return {0, ""};
     }
 
+    const NDq::TDqTaskRunnerMemoryLimits& GetMemoryLimits() const override {
+        return Runner->GetMemoryLimits();
+    }
+
 private:
     void UpdateStats() {
         QueryStat.AddTaskRunnerStats(*Runner->GetStats(), Task.GetId(), Task.GetStageId());