| 185 | options.executor = plan()->query_context()->executor(); |
| 186 | options.should_schedule = ShouldSchedule::IfDifferentExecutor; |
| 187 | ARROW_ASSIGN_OR_RAISE(Future<> scan_task, plan_->query_context()->BeginExternalTask( |
| 188 | "SourceNode::DatasetScan")); |
| 189 | if (!scan_task.is_valid()) { |
| 190 | // Plan has already been aborted, no need to start scanning |
| 191 | return Status::OK(); |
| 192 | } |
| 193 | auto fut = Loop([this, options] { |
| 194 | std::unique_lock<std::mutex> lock(mutex_); |
| 195 | if (stop_requested_) { |
| 196 | return Future<ControlFlow<int>>::MakeFinished(Break(batch_count_)); |
| 197 | } |
| 198 | lock.unlock(); |
| 199 | |
| 200 | arrow::util::tracing::Span fetch_batch_span; |
| 201 | auto fetch_batch_scope = |
| 202 | START_SCOPED_SPAN(fetch_batch_span, "SourceNode::ReadBatch"); |
| 203 | return generator_().Then( |
| 204 | [this]( |
| 205 | const std::optional<ExecBatch>& morsel_or_end) -> Future<ControlFlow<int>> { |
| 206 | std::unique_lock<std::mutex> lock(mutex_); |
| 207 | if (IsIterationEnd(morsel_or_end) || stop_requested_) { |
| 208 | return Break(batch_count_); |
| 209 | } |
| 210 | lock.unlock(); |
| 211 | SliceAndDeliverMorsel(*morsel_or_end); |
| 212 | lock.lock(); |
| 213 | if (!backpressure_future_.is_finished()) { |
| 214 | EVENT_ON_CURRENT_SPAN("SourceNode::BackpressureApplied"); |
| 215 | return backpressure_future_.Then( |
| 216 | []() -> ControlFlow<int> { return Continue(); }); |
| 217 | } |
| 218 | return Future<ControlFlow<int>>::MakeFinished(Continue()); |
| 219 | }, |
| 220 | [](const Status& err) -> Future<ControlFlow<int>> { return err; }, options); |
| 221 | }); |
| 222 | fut.AddCallback( |
| 223 | [this, scan_task](Result<int> maybe_total_batches) mutable { |
| 224 | if (maybe_total_batches.ok()) { |
| 225 | plan_->query_context()->ScheduleTask( |
| 226 | [this, total_batches = *maybe_total_batches] { |
| 227 | return output_->InputFinished(this, total_batches); |
| 228 | }, |
| 229 | "SourceNode::InputFinished"); |
| 230 | } |
| 231 | scan_task.MarkFinished(maybe_total_batches.status()); |
| 232 | }, |
| 233 | options); |
| 234 | return Status::OK(); |
| 235 | } |
| 236 | |
| 237 | const Ordering& ordering() const override { return ordering_; } |
| 238 | |
| 239 | void PauseProducing(ExecNode* output, int32_t counter) override { |
| 240 | std::lock_guard<std::mutex> lg(mutex_); |
nothing calls this directly
no test coverage detected