| 164 | } |
| 165 | |
| 166 | Status StartProducing() override { |
| 167 | NoteStartProducing(ToStringExtra()); |
| 168 | { |
| 169 | // If another exec node encountered an error during its StartProducing call |
| 170 | // it might have already called StopProducing on all of its inputs (including this |
| 171 | // node). |
| 172 | // |
| 173 | std::unique_lock<std::mutex> lock(mutex_); |
| 174 | if (stop_requested_) { |
| 175 | return Status::OK(); |
| 176 | } |
| 177 | started_ = true; |
| 178 | } |
| 179 | |
| 180 | CallbackOptions options; |
| 181 | // These options will transfer execution to the desired Executor if necessary. |
| 182 | // This can happen for in-memory scans where batches don't require |
| 183 | // any CPU work to decode. Otherwise, parsing etc should have already |
| 184 | // been placed us on the desired Executor and no queues will be pushed to. |
| 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 { |
nothing calls this directly
no test coverage detected