westonpace commented on a change in pull request #9644:
URL: https://github.com/apache/arrow/pull/9644#discussion_r601455438
##########
File path: cpp/src/arrow/csv/reader.cc
##########
@@ -672,101 +687,142 @@ class BaseStreamingReader : public ReaderMixin, public
csv::StreamingReader {
std::vector<std::shared_ptr<ColumnDecoder>> column_decoders_;
std::shared_ptr<Schema> schema_;
std::shared_ptr<RecordBatch> pending_batch_;
- Iterator<std::shared_ptr<Buffer>> buffer_iterator_;
+ AsyncGenerator<std::shared_ptr<Buffer>> buffer_generator_;
+ Executor* cpu_executor_;
bool eof_ = false;
};
/////////////////////////////////////////////////////////////////////////
// Serial StreamingReader implementation
-class SerialStreamingReader : public BaseStreamingReader {
+class SerialStreamingReader : public BaseStreamingReader,
+ public
std::enable_shared_from_this<SerialStreamingReader> {
public:
using BaseStreamingReader::BaseStreamingReader;
- Status Init() override {
+ Future<std::shared_ptr<csv::StreamingReader>> Init() override {
ARROW_ASSIGN_OR_RAISE(auto istream_it,
io::MakeInputStreamIterator(input_,
read_options_.block_size));
- // Since we're converting serially, no need to readahead more than one
block
- int32_t block_queue_size = 1;
- ARROW_ASSIGN_OR_RAISE(auto rh_it,
- MakeReadaheadIterator(std::move(istream_it),
block_queue_size));
- buffer_iterator_ = CSVBufferIterator::Make(std::move(rh_it));
- task_group_ = internal::TaskGroup::MakeSerial(stop_token_);
+ ARROW_ASSIGN_OR_RAISE(auto bg_it,
MakeBackgroundGenerator(std::move(istream_it),
+
io_context_.executor()));
+
+ auto rh_it = MakeSerialReadaheadGenerator(std::move(bg_it), 8);
+
+ auto transferred_it = MakeTransferredGenerator(rh_it, cpu_executor_);
+ buffer_generator_ =
CSVBufferIterator::MakeAsync(std::move(transferred_it));
+ task_group_ = internal::TaskGroup::MakeSerial(io_context_.stop_token());
+
+ auto self = shared_from_this();
// Read schema from first batch
- ARROW_ASSIGN_OR_RAISE(pending_batch_, ReadNext());
- DCHECK_NE(schema_, nullptr);
- return Status::OK();
+ return ReadNextAsync().Then([self](const std::shared_ptr<RecordBatch>&
first_batch)
+ ->
Result<std::shared_ptr<csv::StreamingReader>> {
+ self->pending_batch_ = first_batch;
+ DCHECK_NE(self->schema_, nullptr);
+ return self;
+ });
}
- protected:
- Result<std::shared_ptr<RecordBatch>> ReadNext() override {
- if (eof_) {
- return nullptr;
- }
- if (stop_token_.IsStopRequested()) {
- eof_ = true;
- return stop_token_.Poll();
- }
- if (!block_iterator_) {
- Status st = SetupReader();
- if (!st.ok()) {
- // Can't setup reader => bail out
- eof_ = true;
- return st;
- }
+ Result<std::shared_ptr<RecordBatch>> DecodeBatchAndUpdateSchema() {
+ auto maybe_batch = DecodeNextBatch();
Review comment:
Hmm, I thought that was only for the parallel CSV reader (the streaming
reader is exclusively serial). I'm not sure how it could work otherwise. It
seems this would block indefinitely.
--
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:
us...@infra.apache.org
