async ca/task runner: get rid of channel/etc id lists in events

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

@@ -88,14 +88,8 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
 
     void DoBootstrap() {
         NActors::IActor* actor;
-        THashSet<ui32> inputWithDisabledCheckpointing;
-        for (const auto&[idx, inputInfo]: InputChannelsMap) {
-            if (inputInfo.CheckpointingMode == NDqProto::CHECKPOINTING_MODE_DISABLED) {
-                inputWithDisabledCheckpointing.insert(idx);
-            }
-        }
         std::tie(TaskRunnerActor, actor) = TaskRunnerActorFactory->Create(
-            this, TBase::GetAllocatorPtr(), GetTxId(), Task.GetId(), std::move(inputWithDisabledCheckpointing), InitMemoryQuota());
+            this, TBase::GetAllocatorPtr(), GetTxId(), Task.GetId(), InitMemoryQuota());
         TaskRunnerActorId = RegisterWithSameMailbox(actor);
 
         TDqTaskRunnerMemoryLimits limits;
@@ -199,13 +193,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
             DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), CheckpointRequest, .Defined());
             DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), WatermarkRequest, .Defined());
             DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), MemLimit);
-            for (const auto& sinkId: ContinueRunEvent->SinkIds) {
-                html << "ContinueRunEvent.SinkIds: " << sinkId << "<br />";
-            }
-
-            for (const auto& inputTransformId: ContinueRunEvent->InputTransformIds) {
-                html << "ContinueRunEvent.InputTransformIds: " << inputTransformId << "<br />";
-            }
         }
 
         DUMP((*this), ContinueRunStartWaitTime, .ToString());
@@ -444,7 +431,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
     void OnStateRequest(TEvDqCompute::TEvStateRequest::TPtr& ev) {
         CA_LOG_T("Got TEvStateRequest from actor " << ev->Sender << " PingCookie: " << ev->Cookie);
         if (!SentStatsRequest) {
-            Send(TaskRunnerActorId, new NTaskRunnerActor::TEvStatistics(GetIds(SinksMap), GetIds(InputTransformsMap)));
+            Send(TaskRunnerActorId, new NTaskRunnerActor::TEvStatistics());
             SentStatsRequest = true;
         }
         WaitingForStateResponse.push_back({ev->Sender, ev->Cookie});
@@ -1066,15 +1053,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         return nullptr;
     }
 
-    template<typename TSecond>
-    TVector<ui32> GetIds(const THashMap<ui64, TSecond>& collection) {
-        TVector<ui32> ids;
-        std::transform(collection.begin(), collection.end(), std::back_inserter(ids), [](const auto& p) {
-            return p.first;
-        });
-        return ids;
-    }
-
     void InjectBarrierToOutputs(const NDqProto::TCheckpoint&) override {
         Y_ABORT_UNLESS(CheckpointingMode != NDqProto::CHECKPOINTING_MODE_DISABLED);
         // already done in task_runner_actor
@@ -1139,28 +1117,19 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         if (!ContinueRunEvent) {
             ContinueRunStartWaitTime = TInstant::Now();
             ContinueRunEvent = std::make_unique<NTaskRunnerActor::TEvContinueRun>();
-            ContinueRunEvent->SinkIds = GetIds(SinksMap);
-            ContinueRunEvent->InputTransformIds = GetIds(InputTransformsMap);
         }
         ContinueRunEvent->CheckpointOnly = checkpointOnly;
         if (TMaybe<TInstant> watermarkRequest = GetWatermarkRequest()) {
             if (!ContinueRunEvent->WatermarkRequest) {
                 ContinueRunEvent->WatermarkRequest.ConstructInPlace();
                 ContinueRunEvent->WatermarkRequest->Watermark = *watermarkRequest;
-
-                ContinueRunEvent->WatermarkRequest->ChannelIds.reserve(OutputChannelsMap.size());
-                for (const auto& [channelId, info] : OutputChannelsMap) {
-                    if (info.WatermarksMode != NDqProto::EWatermarksMode::WATERMARKS_MODE_DISABLED) {
-                        ContinueRunEvent->WatermarkRequest->ChannelIds.emplace_back(channelId);
-                    }
-                }
             } else {
                 ContinueRunEvent->WatermarkRequest->Watermark = Max(ContinueRunEvent->WatermarkRequest->Watermark, *watermarkRequest);
             }
         }
         if (checkpointRequest) {
             if (!ContinueRunEvent->CheckpointRequest) {
-                ContinueRunEvent->CheckpointRequest.ConstructInPlace(GetIds(OutputChannelsMap), GetIds(SinksMap), *checkpointRequest);
+                ContinueRunEvent->CheckpointRequest.ConstructInPlace(*checkpointRequest);
             } else {
                 Y_ABORT_UNLESS(ContinueRunEvent->CheckpointRequest->Checkpoint.GetGeneration() == checkpointRequest->GetGeneration());
                 Y_ABORT_UNLESS(ContinueRunEvent->CheckpointRequest->Checkpoint.GetId() == checkpointRequest->GetId());

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

@@ -299,25 +299,19 @@ struct TEvOutputChannelData
 struct TWatermarkRequest {
     TWatermarkRequest() = default;
 
-    TWatermarkRequest(TVector<ui32>&& channelIds, TInstant watermark)
-        : ChannelIds(std::move(channelIds))
-        , Watermark(watermark) {
+    explicit TWatermarkRequest(TInstant watermark)
+        : Watermark(watermark) {
     }
 
-    TVector<ui32> ChannelIds;
     TInstant Watermark;
 };
 
 // Holds info required to inject barriers to outputs
 struct TCheckpointRequest {
-    TCheckpointRequest(TVector<ui32>&& channelIds, TVector<ui32>&& sinkIds, const NDqProto::TCheckpoint& checkpoint)
-        : ChannelIds(std::move(channelIds))
-        , SinkIds(std::move(sinkIds))
-        , Checkpoint(checkpoint) {
+    explicit TCheckpointRequest(const NDqProto::TCheckpoint& checkpoint)
+        : Checkpoint(checkpoint) {
     }
 
-    TVector<ui32> ChannelIds;
-    TVector<ui32> SinkIds;
     NDqProto::TCheckpoint Checkpoint;
 };
 
@@ -337,20 +331,18 @@ struct TEvContinueRun
         , CheckpointOnly(checkpointOnly)
     { }
 
-    TEvContinueRun(THashSet<ui32>&& inputChannels, ui64 memLimit)
+    TEvContinueRun(TVector<ui32>&& inputChannels, ui64 memLimit)
         : AskFreeSpace(false)
         , InputChannels(std::move(inputChannels))
         , MemLimit(memLimit)
     { }
 
     bool AskFreeSpace = true;
-    const THashSet<ui32> InputChannels;
+    const TVector<ui32> InputChannels;
     ui64 MemLimit;
     TMaybe<TWatermarkRequest> WatermarkRequest = Nothing();
     TMaybe<TCheckpointRequest> CheckpointRequest = Nothing();
     bool CheckpointOnly = false;
-    TVector<ui32> SinkIds;
-    TVector<ui32> InputTransformIds;
 };
 
 //Sent by TaskRunnerActor to ComputeActor as an acknowledgement in AsyncInputPush method call
@@ -427,14 +419,10 @@ struct TEvLoadTaskRunnerFromStateDone : NActors::TEventLocal<TEvLoadTaskRunnerFr
 
 struct TEvStatistics : NActors::TEventLocal<TEvStatistics, TTaskRunnerEvents::EvStatistics>
 {
-    explicit TEvStatistics(TVector<ui32>&& sinkIds, TVector<ui32>&& inputTransformIds)
-        : SinkIds(std::move(sinkIds))
-        , InputTransformIds(std::move(inputTransformIds))
-        , Stats() {
+    explicit TEvStatistics()
+        : Stats() {
     }
 
-    TVector<ui32> SinkIds;
-    TVector<ui32> InputTransformIds;
     NDq::TDqTaskRunnerStatsView Stats;
 };
 

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

@@ -46,7 +46,6 @@ struct ITaskRunnerActorFactory {
         std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
         const TTxId& txId,
         ui64 taskId,
-        THashSet<ui32>&& inputChannelsWithDisabledCheckpoints = {},
         THolder<NYql::NDq::TDqMemoryQuota>&& memoryQuota = {}) = 0;
 };
 

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

@@ -40,14 +40,13 @@ class TLocalTaskRunnerActor
 public:
     static constexpr char ActorName[] = "YQL_DQ_TASK_RUNNER";
 
-    TLocalTaskRunnerActor(ITaskRunnerActor::ICallbacks* parent, const TTaskRunnerFactory& factory, std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const TTxId& txId, ui64 taskId, THashSet<ui32>&& inputChannelsWithDisabledCheckpoints, THolder<NYql::NDq::TDqMemoryQuota>&& memoryQuota)
+    TLocalTaskRunnerActor(ITaskRunnerActor::ICallbacks* parent, const TTaskRunnerFactory& factory, std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const TTxId& txId, ui64 taskId, THolder<NYql::NDq::TDqMemoryQuota>&& memoryQuota)
         : TActor<TLocalTaskRunnerActor>(&TLocalTaskRunnerActor::Handler)
         , Alloc(alloc)
         , Parent(parent)
         , Factory(factory)
         , TxId(txId)
         , TaskId(taskId)
-        , InputChannelsWithDisabledCheckpoints(std::move(inputChannelsWithDisabledCheckpoints))
         , MemoryQuota(std::move(memoryQuota))
     {
     }
@@ -91,12 +90,12 @@ class TLocalTaskRunnerActor
     void OnStatisticsRequest(TEvStatistics::TPtr& ev) {
 
         THashMap<ui32, const IDqAsyncOutputBuffer*> sinks;
-        for (const auto sinkId : ev->Get()->SinkIds) {
+        for (const auto sinkId : Sinks) {
             sinks[sinkId] = TaskRunner->GetSink(sinkId).Get();
         }
 
         THashMap<ui32, const IDqAsyncInputBuffer*> inputTransforms;
-        for (const auto inputTransformId : ev->Get()->InputTransformIds) {
+        for (const auto inputTransformId : InputTransforms) {
             inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId)->second.Get();
         }
 
@@ -132,11 +131,7 @@ class TLocalTaskRunnerActor
     }
 
     bool ReadyToCheckpoint() {
-        for (const auto inputChannelId: Inputs) {
-            if (InputChannelsWithDisabledCheckpoints.contains(inputChannelId)) {
-                continue;
-            }
-
+        for (const auto inputChannelId: InputsWithCheckpoints) {
             const auto input = TaskRunner->GetInputChannel(inputChannelId);
             if (!input->IsPaused()) {
                 return false;
@@ -162,7 +157,7 @@ class TLocalTaskRunnerActor
 
     void OnContinueRun(TEvContinueRun::TPtr& ev) {
         auto guard = TaskRunner->BindAllocator(MemoryQuota ? MemoryQuota->GetMkqlMemoryLimit() : ev->Get()->MemLimit);
-        auto inputMap = ev->Get()->AskFreeSpace
+        const auto& inputMap = ev->Get()->AskFreeSpace
             ? Inputs
             : ev->Get()->InputChannels;
 
@@ -198,9 +193,9 @@ class TLocalTaskRunnerActor
             if (shouldHandleWatermark) {
                 const auto watermarkRequested = ev->Get()->WatermarkRequest->Watermark;
                 LOG_T("Task runner. Watermarks. Injecting requested watermark " << watermarkRequested
-                    << " to " << ev->Get()->WatermarkRequest->ChannelIds.size() << " outputs ");
+                    << " to " << OutputsWithWatermarks.size() << " outputs ");
 
-                for (const auto& channelId : ev->Get()->WatermarkRequest->ChannelIds) {
+                for (const auto& channelId : OutputsWithWatermarks) {
                     NDqProto::TWatermark watermark;
                     watermark.SetTimestampUs(watermarkRequested.MicroSeconds());
                     TaskRunner->GetOutputChannel(channelId)->Push(std::move(watermark));
@@ -218,10 +213,10 @@ class TLocalTaskRunnerActor
                     data.Blob = TaskRunner->Save();
                     // inject barriers
                     // todo:(whcrc) barriers are injected even if source state save failed
-                    for (const auto& channelId : ev->Get()->CheckpointRequest->ChannelIds) {
+                    for (const auto& channelId : Outputs) {
                         TaskRunner->GetOutputChannel(channelId)->Push(NDqProto::TCheckpoint(ev->Get()->CheckpointRequest->Checkpoint));
                     }
-                    for (const auto& sinkId : ev->Get()->CheckpointRequest->SinkIds) {
+                    for (const auto& sinkId : Sinks) {
                         TaskRunner->GetSink(sinkId)->Push(NDqProto::TCheckpoint(ev->Get()->CheckpointRequest->Checkpoint));
                     }
                 } catch (const std::exception& e) {
@@ -236,15 +231,15 @@ class TLocalTaskRunnerActor
         }
 
         {
-            auto st = MakeHolder<TEvStatistics>(std::move(ev->Get()->SinkIds), std::move(ev->Get()->InputTransformIds));
+            auto st = MakeHolder<TEvStatistics>();
 
             THashMap<ui32, const IDqAsyncOutputBuffer*> sinks;
-            for (const auto sinkId : st->SinkIds) {
+            for (const auto sinkId : Sinks) {
                 sinks[sinkId] = TaskRunner->GetSink(sinkId).Get();
             }
 
             THashMap<ui32, const IDqAsyncInputBuffer*> inputTransforms;
-            for (const auto inputTransformId : st->InputTransformIds) { // TODO
+            for (const auto inputTransformId : InputTransforms) { // TODO
                 inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId)->second.Get();
             }
 
@@ -428,13 +423,35 @@ class TLocalTaskRunnerActor
         for (auto inputId = 0; inputId < inputs.size(); inputId++) {
             auto& input = inputs[inputId];
             if (input.HasSource()) {
-                Sources.emplace(inputId);
+                Sources.emplace_back(inputId);
             } else {
                 for (auto& channel : input.GetChannels()) {
-                    Inputs.emplace(channel.GetId());
+                    Inputs.emplace_back(channel.GetId());
+                    if (channel.GetCheckpointingMode() != NDqProto::CHECKPOINTING_MODE_DISABLED) {
+                        InputsWithCheckpoints.emplace_back(channel.GetId());
+                    }
                 }
             }
         }
+        std::sort(Inputs.begin(), Inputs.end());
+        Y_ENSURE(std::unique(Inputs.begin(), Inputs.end()) == Inputs.end());
+
+        auto& outputs = settings.GetOutputs();
+        for (auto outputId = 0; outputId < outputs.size(); outputId++) {
+            auto& output = outputs[outputId];
+            if (output.HasSink()) {
+                Sinks.emplace_back(outputId);
+            } else {
+                for (auto& channel : output.GetChannels()) {
+                    Outputs.emplace_back(channel.GetId());
+                    if (channel.GetWatermarksMode() != NDqProto::WATERMARKS_MODE_DISABLED) {
+                        OutputsWithWatermarks.emplace_back(channel.GetId());
+                    }
+                }
+            }
+        }
+        std::sort(Outputs.begin(), Outputs.end());
+        Y_ENSURE(std::unique(Outputs.begin(), Outputs.end()) == Outputs.end());
 
         auto guard = TaskRunner->BindAllocator(MemoryQuota ? TMaybe<ui64>(MemoryQuota->GetMkqlMemoryLimit()) : Nothing());
         if (MemoryQuota) {
@@ -456,8 +473,11 @@ class TLocalTaskRunnerActor
         THashMap<ui64, std::pair<NUdf::TUnboxedValue, IDqAsyncInputBuffer::TPtr>> inputTransforms;
         for (auto i = 0; i != inputs.size(); ++i) {
             if (auto t = TaskRunner->GetInputTransform(i)) {
+                Y_ENSURE(inputs[i].HasTransform());
                 inputTransforms[i] = *t;
-                InputTransforms.emplace(i);
+                InputTransforms.emplace_back(i);
+            } else {
+                Y_ENSURE(!inputs[i].HasTransform());
             }
         }
 
@@ -504,11 +524,14 @@ class TLocalTaskRunnerActor
     TTaskRunnerFactory Factory;
     const TTxId TxId;
     const ui64 TaskId;
-    THashSet<ui32> Inputs;
-    THashSet<ui32> InputTransforms;
-    THashSet<ui32> Sources;
+    TVector<ui32> Inputs;
+    TVector<ui32> InputsWithCheckpoints;
+    TVector<ui32> InputTransforms;
+    TVector<ui32> Sources;
+    TVector<ui32> Sinks;
+    TVector<ui32> Outputs;
+    TVector<ui32> OutputsWithWatermarks;
     TIntrusivePtr<NDq::IDqTaskRunner> TaskRunner;
-    THashSet<ui32> InputChannelsWithDisabledCheckpoints;
     THolder<TDqMemoryQuota> MemoryQuota;
     ui64 ActorElapsedTicks = 0;
 };
@@ -523,10 +546,9 @@ struct TLocalTaskRunnerActorFactory: public ITaskRunnerActorFactory {
         std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
         const TTxId& txId,
         ui64 taskId,
-        THashSet<ui32>&& inputChannelsWithDisabledCheckpoints,
         THolder<NYql::NDq::TDqMemoryQuota>&& memoryQuota) override
     {
-        auto* actor = new TLocalTaskRunnerActor(parent, Factory, alloc, txId, taskId, std::move(inputChannelsWithDisabledCheckpoints), std::move(memoryQuota));
+        auto* actor = new TLocalTaskRunnerActor(parent, Factory, alloc, txId, taskId, std::move(memoryQuota));
         return std::make_tuple(
             static_cast<ITaskRunnerActor*>(actor),
             static_cast<NActors::IActor*>(actor)

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

@@ -568,11 +568,11 @@ class TDqWorker: public TRichActor<TDqWorker>
             return;
         }
 
-        THashSet<ui32> inputChannels;
+        TVector<ui32> inputChannels;
         for (auto& input : InputMap) {
             auto& channel = input.second;
             if (!channel.Requested && !channel.Finished) {
-                inputChannels.insert(channel.ChannelId);
+                inputChannels.push_back(channel.ChannelId);
             }
         }
 

ydb/library/yql/providers/dq/task_runner_actor/task_runner_actor.cpp

@@ -589,10 +589,10 @@ class TTaskRunnerActor
         for (auto inputId = 0; inputId < inputs.size(); inputId++) {
             auto& input = inputs[inputId];
             if (input.HasSource()) {
-                Sources.emplace(inputId);
+                Sources.emplace_back(inputId);
             } else {
                 for (auto& channel : input.GetChannels()) {
-                    Inputs.emplace(channel.GetId());
+                    Inputs.emplace_back(channel.GetId());
                 }
             }
         }
@@ -756,8 +756,8 @@ class TTaskRunnerActor
     NTaskRunnerProxy::ITaskRunner::TPtr TaskRunner;
     ITaskRunnerInvoker::TPtr Invoker;
     bool Local;
-    THashSet<ui32> Inputs;
-    THashSet<ui32> Sources;
+    TVector<ui32> Inputs;
+    TVector<ui32> Sources;
     TIntrusivePtr<TDqConfiguration> Settings;
     NDqProto::EDataTransportVersion DataTransportVersion;
     ui64 StageId;
@@ -784,7 +784,6 @@ class TTaskRunnerActorFactory: public ITaskRunnerActorFactory {
         std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
         const TTxId& txId,
         ui64 taskId,
-        THashSet<ui32>&&,
         THolder<NYql::NDq::TDqMemoryQuota>&&) override
     {
         auto* actor = new TTaskRunnerActor(parent, alloc, ProxyFactory, InvokerFactory->Create(), txId, taskId, RuntimeData);