| 2540 | : queue_size_{queue_size}, mapper_{std::move(mapper)} {} |
| 2541 | |
| 2542 | Result<RecordBatchWithMetadata> next() { |
| 2543 | TaskWithMetadata task; |
| 2544 | { |
| 2545 | std::unique_lock<std::mutex> lock(mutex_); |
| 2546 | cv_.wait(lock, |
| 2547 | [&] { return !error_.ok() || !batches_.empty() || end_of_stream_; }); |
| 2548 | if (!error_.ok()) { |
| 2549 | return error_; |
| 2550 | } |
| 2551 | |
| 2552 | if (batches_.empty() && end_of_stream_) { |
| 2553 | return IterationEnd<RecordBatchWithMetadata>(); |
| 2554 | } |
| 2555 | |
| 2556 | task = std::move(batches_.front()); |
| 2557 | batches_.pop(); |
| 2558 | } |
| 2559 | |
| 2560 | producer_->request(producer_, 1); |
| 2561 | ArrowDeviceArray out; |
| 2562 | if (task.task_.extract_data(&task.task_, &out) != 0) { |
| 2563 | std::unique_lock<std::mutex> lock(mutex_); |
| 2564 | cv_.wait(lock, [&] { return !error_.ok(); }); |
| 2565 | return error_; |
| 2566 | } |
| 2567 | |
| 2568 | ARROW_ASSIGN_OR_RAISE(auto batch, ImportDeviceRecordBatch(&out, schema_, mapper_)); |
| 2569 | return RecordBatchWithMetadata{std::move(batch), std::move(task.metadata_)}; |
| 2570 | } |
| 2571 | |
| 2572 | const uint64_t queue_size_; |
| 2573 | const DeviceMemoryMapper mapper_; |
no test coverage detected