westonpace commented on a change in pull request #9095:
URL: https://github.com/apache/arrow/pull/9095#discussion_r564745648
##########
File path: cpp/src/arrow/csv/reader.cc
##########
@@ -837,20 +885,121 @@ class ThreadedTableReader : public BaseTableReader {
}
protected:
- ThreadPool* thread_pool_;
+ Executor* thread_pool_;
+ Iterator<std::shared_ptr<Buffer>> buffer_iterator_;
+};
+
+class AsyncThreadedTableReader
+ : public BaseTableReader,
+ public std::enable_shared_from_this<AsyncThreadedTableReader> {
+ public:
+ using BaseTableReader::BaseTableReader;
+
+ AsyncThreadedTableReader(MemoryPool* pool, std::shared_ptr<io::InputStream>
input,
+ const ReadOptions& read_options,
+ const ParseOptions& parse_options,
+ const ConvertOptions& convert_options, Executor*
thread_pool)
+ : BaseTableReader(pool, input, read_options, parse_options,
convert_options),
+ thread_pool_(thread_pool) {}
+
+ ~AsyncThreadedTableReader() override {
+ if (task_group_) {
+ // In case of error, make sure all pending tasks are finished before
+ // we start destroying BaseTableReader members
+ ARROW_UNUSED(task_group_->Finish());
+ }
+ }
+
+ Status Init() override {
+ ARROW_ASSIGN_OR_RAISE(auto istream_it,
+ io::MakeInputStreamIterator(input_,
read_options_.block_size));
+
+ ARROW_ASSIGN_OR_RAISE(auto bg_it,
+ MakeBackgroundIterator(std::move(istream_it),
thread_pool_));
+
+ int32_t block_queue_size = thread_pool_->GetCapacity();
+ auto rh_it = AddReadahead(bg_it, block_queue_size);
+ buffer_generator_ = CSVBufferIterator::MakeAsync(std::move(rh_it));
+ return Status::OK();
+ }
+
+ Result<std::shared_ptr<Table>> Read() override { return
ReadAsync().result(); }
+
+ Future<std::shared_ptr<Table>> ReadAsync() override {
+ task_group_ = internal::TaskGroup::MakeThreaded(thread_pool_);
+
+ auto self = shared_from_this();
+ return ProcessFirstBuffer().Then([self](const std::shared_ptr<Buffer>
first_buffer) {
+ auto block_generator = ThreadedBlockReader::MakeAsyncIterator(
+ self->buffer_generator_, MakeChunker(self->parse_options_),
+ std::move(first_buffer));
+
+ std::function<Status(util::optional<CSVBlock>)> block_visitor =
+ [self](util::optional<CSVBlock> maybe_block) -> Status {
+ DCHECK(!maybe_block->consume_bytes);
Review comment:
See previous comment. The Visit function is never passed the end token,
there is no need. However, since the user is allowed to specify their own end
token, this use of util::optional is unavoidable. For example, a user could
create an iterator of TestInt where the end token is simply the TestInt with
value -999. Now there is no "wrapping" going on (in the way that
`util::optional` wraps `CSVBlock`) so there is nothing to unwrap when calling
the `Visit` function.
On the other hand, if we assume `util::optional` is always used for
iterators of non-pointer types we could redefine visit to take in T instead.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
