YQ-2068 support input transforms in async Compute Actor (#3459)

zverevgeny · web-flow · commit 6e5100ed51ff · 2024-04-04T18:43:36.000+03:00
diff --git a/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp b/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp
@@ -465,6 +465,9 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
             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);
 
         {
diff --git a/ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h b/ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h
@@ -216,7 +216,7 @@ class TDqSyncComputeActorBase: public TDqComputeActorBase<TDerived, TComputeActo
         }
 
         for (auto& [inputIndex, transform] : this->InputTransformsMap) {
-            std::tie(transform.Input, transform.Buffer) = TaskRunner->GetInputTransform(inputIndex);
+            std::tie(transform.Input, transform.Buffer) = *TaskRunner->GetInputTransform(inputIndex);
         }
 
         for (auto& [channelId, channel] : this->OutputChannelsMap) {
diff --git a/ydb/library/yql/dq/actors/task_runner/events.h b/ydb/library/yql/dq/actors/task_runner/events.h
@@ -28,12 +28,12 @@ struct TTaskRunnerEvents {
 
         EvOutputChannelDataRequest,
         EvOutputChannelData,
-
-	EvInputChannelData,
+        
+        EvInputChannelData,
         EvInputChannelDataAck,
-
-	// EvContinueRun -> TaskRunner->Run() -> TEvTaskRunFinished
-        EvContinueRun,
+        
+        // EvContinueRun -> TaskRunner->Run() -> TEvTaskRunFinished
+        EvContinueRun, 
         EvRunFinished,
 
         EvSourceDataAck,
@@ -171,21 +171,34 @@ struct TEvTaskRunnerCreateFinished
     : NActors::TEventLocal<TEvTaskRunnerCreateFinished, TTaskRunnerEvents::EvCreateFinished>
 {
 
-    TEvTaskRunnerCreateFinished() = default;
     TEvTaskRunnerCreateFinished(
         const THashMap<TString, TString>& secureParams,
         const THashMap<TString, TString>& taskParams,
         const TVector<TString>& readRanges,
         const NKikimr::NMiniKQL::TTypeEnvironment& typeEnv,
         const NKikimr::NMiniKQL::THolderFactory& holderFactory,
-        const TTaskRunnerActorSensors& sensors = {})
+        std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
+        THashMap<ui64, std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>>&& inputTransforms,
+        const TTaskRunnerActorSensors& sensors = {}
+    )
         : Sensors(sensors)
         , SecureParams(secureParams)
         , TaskParams(taskParams)
         , ReadRanges(readRanges)
         , TypeEnv(typeEnv)
         , HolderFactory(holderFactory)
-    { }
+        , Alloc(alloc)
+        , InputTransforms(std::move(inputTransforms))
+    { 
+        Y_ABORT_UNLESS(inputTransforms.empty() || Alloc);
+    }
+
+    ~TEvTaskRunnerCreateFinished() {
+        if (!InputTransforms.empty()) {
+            auto guard = Guard(*Alloc);
+            InputTransforms.clear();
+        }
+    }
 
     TTaskRunnerActorSensors Sensors;
 
@@ -195,6 +208,8 @@ struct TEvTaskRunnerCreateFinished
     const TVector<TString>& ReadRanges;
     const NKikimr::NMiniKQL::TTypeEnvironment& TypeEnv;
     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
 };
 
 struct TEvTaskRunFinished
diff --git a/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp b/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp
@@ -95,7 +95,7 @@ class TLocalTaskRunnerActor
 
         THashMap<ui32, const IDqAsyncInputBuffer*> inputTransforms;
         for (const auto inputTransformId : ev->Get()->InputTransformIds) {
-            inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId).second.Get();
+            inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId)->second.Get();
         }
 
         ev->Get()->Stats = TDqTaskRunnerStatsView(TaskRunner->GetStats(), std::move(sinks), std::move(inputTransforms));
@@ -231,7 +231,7 @@ class TLocalTaskRunnerActor
 
             THashMap<ui32, const IDqAsyncInputBuffer*> inputTransforms;
             for (const auto inputTransformId : st->InputTransformIds) { // TODO
-                inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId).second.Get();
+                inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId)->second.Get();
             }
 
             st->Stats = TDqTaskRunnerStatsView(TaskRunner->GetStats(), std::move(sinks), std::move(inputTransforms));
@@ -427,6 +427,14 @@ 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));
 
@@ -435,7 +443,10 @@ class TLocalTaskRunnerActor
             TaskRunner->GetTaskParams(),
             TaskRunner->GetReadRanges(),
             TaskRunner->GetTypeEnv(),
-            TaskRunner->GetHolderFactory());
+            TaskRunner->GetHolderFactory(),
+            Alloc,
+            std::move(inputTransforms)
+        );
 
         Send(
             ParentId,
diff --git a/ydb/library/yql/dq/runtime/dq_tasks_runner.cpp b/ydb/library/yql/dq/runtime/dq_tasks_runner.cpp
@@ -779,10 +779,12 @@ class TDqTaskRunner : public IDqTaskRunner {
         return *ptr;
     }
 
-    std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr> GetInputTransform(ui64 inputIndex) override {
-        auto ptr = AllocatedHolder->InputTransforms.FindPtr(inputIndex);
-        YQL_ENSURE(ptr, "task: " << TaskId << " does not have input index: " << inputIndex << " or such transform");
-        return {ptr->TransformInput, ptr->TransformOutput};
+    std::optional<std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> GetInputTransform(ui64 inputIndex) override {
+        if (auto ptr = AllocatedHolder->InputTransforms.FindPtr(inputIndex)) {
+            return {std::pair{ptr->TransformInput, ptr->TransformOutput}};
+        } else {
+            return std::nullopt;
+        }
     }
 
     std::pair<IDqAsyncOutputBuffer::TPtr, IDqOutputConsumer::TPtr> GetOutputTransform(ui64 outputIndex) override {
diff --git a/ydb/library/yql/dq/runtime/dq_tasks_runner.h b/ydb/library/yql/dq/runtime/dq_tasks_runner.h
@@ -373,7 +373,7 @@ class IDqTaskRunner : public TSimpleRefCount<IDqTaskRunner>, private TNonCopyabl
     virtual IDqAsyncInputBuffer::TPtr GetSource(ui64 inputIndex) = 0;
     virtual IDqOutputChannel::TPtr GetOutputChannel(ui64 channelId) = 0;
     virtual IDqAsyncOutputBuffer::TPtr GetSink(ui64 outputIndex) = 0;
-    virtual std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr> GetInputTransform(ui64 inputIndex) = 0;
+    virtual std::optional<std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> GetInputTransform(ui64 inputIndex) = 0;
     virtual std::pair<IDqAsyncOutputBuffer::TPtr, IDqOutputConsumer::TPtr> GetOutputTransform(ui64 outputIndex) = 0;
 
     virtual IRandomProvider* GetRandomProvider() const = 0;
diff --git a/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp b/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp
@@ -1715,7 +1715,7 @@ class TDqTaskRunner: public NDq::IDqTaskRunner {
         return sink;
     }
 
-    std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr> GetInputTransform(ui64 /*inputIndex*/) override {
+    std::optional<std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> GetInputTransform(ui64 /*inputIndex*/) override {
         return {};
     }
 
diff --git a/ydb/library/yql/providers/dq/task_runner_actor/task_runner_actor.cpp b/ydb/library/yql/providers/dq/task_runner_actor/task_runner_actor.cpp
@@ -636,6 +636,8 @@ class TTaskRunnerActor
                     taskRunner->GetReadRanges(),
                     taskRunner->GetTypeEnv(),
                     taskRunner->GetHolderFactory(),
+                    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc>{},
+                    THashMap<ui64, std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>>{},
                     sensors);
 
                 actorSystem->Send(
diff --git a/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.cpp b/ydb/library/yql/providers/yt/gateway/file/yql_yt_file.cpp
@@ -759,7 +759,7 @@ class TYtFileGateway : public IYtGateway {
             writer.SetSpecs(spec);
 
             TStringStream err;
-            auto type = BuildType(*tableInfo.RowSpec->GetType(), typeBuilder, err);//
+            auto type = BuildType(*tableInfo.RowSpec->GetType(), typeBuilder, err);
             TValuePacker packer(true, type);
             for (auto& c: content) {
                 auto val = packer.Unpack(c, holderFactory);

Original file line number	Diff line number	Diff line change
`@@ -465,6 +465,9 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC`
`465`	`465`	`Stat->AddCounters2(ev->Get()->Sensors);`
`466`	`466`	`}`
`467`	`467`	`TypeEnv = const_cast<NKikimr::NMiniKQL::TTypeEnvironment*>(&typeEnv);`
	`468`	`+ for (auto& [inputIndex, transform] : this->InputTransformsMap) {`
	`469`	`+ std::tie(transform.Input, transform.Buffer) = ev->Get()->InputTransforms.at(inputIndex);`
	`470`	`+ }`
`468`	`471`	`FillIoMaps(holderFactory, typeEnv, secureParams, taskParams, readRanges, nullptr);`
`469`	`472`
`470`	`473`	`{`
Original file line number	Diff line number	Diff line change
`@@ -216,7 +216,7 @@ class TDqSyncComputeActorBase: public TDqComputeActorBase<TDerived, TComputeActo`
`216`	`216`	`}`
`217`	`217`
`218`	`218`	`for (auto& [inputIndex, transform] : this->InputTransformsMap) {`
`219`		`- std::tie(transform.Input, transform.Buffer) = TaskRunner->GetInputTransform(inputIndex);`
	`219`	`+ std::tie(transform.Input, transform.Buffer) = *TaskRunner->GetInputTransform(inputIndex);`
`220`	`220`	`}`
`221`	`221`
`222`	`222`	`for (auto& [channelId, channel] : this->OutputChannelsMap) {`
Original file line number	Diff line number	Diff line change
`@@ -1715,7 +1715,7 @@ class TDqTaskRunner: public NDq::IDqTaskRunner {`
`1715`	`1715`	`return sink;`
`1716`	`1716`	`}`
`1717`	`1717`
`1718`		`- std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr> GetInputTransform(ui64 /inputIndex/) override {`
	`1718`	`+ std::optional<std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> GetInputTransform(ui64 /inputIndex/) override {`
`1719`	`1719`	`return {};`
`1720`	`1720`	`}`
`1721`	`1721`