Skip to content

ARROW-12161: [C++][Dataset] Revert async CSV reader in datasets #10018

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
lidavidm wants to merge 1 commit into from
Closed

ARROW-12161: [C++][Dataset] Revert async CSV reader in datasets #10018

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 21 additions & 51 deletions cpp/src/arrow/dataset/file_csv.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@
#include "arrow/io/compressed.h"
#include "arrow/result.h"
#include "arrow/type.h"
#include "arrow/util/async_generator.h"
#include "arrow/util/iterator.h"
#include "arrow/util/logging.h"

Expand Down Expand Up @@ -114,53 +113,35 @@ static inline Result<csv::ReadOptions> GetReadOptions(
return read_options;
}

static inline Future<std::shared_ptr<csv::StreamingReader>> OpenReaderAsync(
static inline Result<std::shared_ptr<csv::StreamingReader>> OpenReader(
const FileSource& source, const CsvFileFormat& format,
internal::Executor* cpu_executor,
const std::shared_ptr<ScanOptions>& scan_options = nullptr,
MemoryPool* pool = default_memory_pool()) {
ARROW_ASSIGN_OR_RAISE(auto reader_options, GetReadOptions(format, scan_options));

util::string_view first_block;
ARROW_ASSIGN_OR_RAISE(auto input, source.OpenCompressed());
ARROW_ASSIGN_OR_RAISE(
input, io::BufferedInputStream::Create(reader_options.block_size,
default_memory_pool(), std::move(input)));
ARROW_ASSIGN_OR_RAISE(first_block, input->Peek(reader_options.block_size));

auto peek_fut = DeferNotOk(input->io_context().executor()->Submit(
[input, reader_options] { return input->Peek(reader_options.block_size); }));

return peek_fut.Then([=](const util::string_view& first_block)
-> Future<std::shared_ptr<csv::StreamingReader>> {
const auto& parse_options = format.parse_options;
auto convert_options = csv::ConvertOptions::Defaults();
if (scan_options != nullptr) {
ARROW_ASSIGN_OR_RAISE(convert_options,
GetConvertOptions(format, scan_options, first_block, pool));
}

return csv::StreamingReader::MakeAsync(io::default_io_context(), std::move(input),
cpu_executor, reader_options, parse_options,
convert_options)
.Then(
[](const std::shared_ptr<csv::StreamingReader>& maybe_reader)
-> Result<std::shared_ptr<csv::StreamingReader>> { return maybe_reader; },
[source](const Status& err) -> Result<std::shared_ptr<csv::StreamingReader>> {
return err.WithMessage("Could not open CSV input source '", source.path(),
"': ", err);
});
});
}
const auto& parse_options = format.parse_options;
auto convert_options = csv::ConvertOptions::Defaults();
if (scan_options != nullptr) {
ARROW_ASSIGN_OR_RAISE(convert_options,
GetConvertOptions(format, scan_options, first_block, pool));
}

static inline Result<std::shared_ptr<csv::StreamingReader>> OpenReader(
const FileSource& source, const CsvFileFormat& format,
const std::shared_ptr<ScanOptions>& scan_options = nullptr,
MemoryPool* pool = default_memory_pool()) {
bool use_threads = (scan_options != nullptr && scan_options->use_threads);
return internal::RunSynchronously<std::shared_ptr<csv::StreamingReader>>(
[&](Executor* executor) {
return OpenReaderAsync(source, format, executor, scan_options, pool);
},
use_threads);
auto maybe_reader =
csv::StreamingReader::Make(io::IOContext(pool), std::move(input), reader_options,
parse_options, convert_options);
if (!maybe_reader.ok()) {
return maybe_reader.status().WithMessage("Could not open CSV input source '",
source.path(), "': ", maybe_reader.status());
}

return std::move(maybe_reader).ValueOrDie();
}

/// \brief A ScanTask backed by an Csv file.
Expand All @@ -174,20 +155,9 @@ class CsvScanTask : public ScanTask {
source_(fragment->source()) {}

Result<RecordBatchIterator> Execute() override {
ARROW_ASSIGN_OR_RAISE(auto gen, ExecuteAsync(internal::GetCpuThreadPool()));
return MakeGeneratorIterator(std::move(gen));
}

bool supports_async() const override { return true; }

Result<RecordBatchGenerator> ExecuteAsync(internal::Executor* cpu_executor) override {
auto reader_fut =
OpenReaderAsync(source_, *format_, cpu_executor, options(), options()->pool);
auto generator_fut = reader_fut.Then(
[](const std::shared_ptr<csv::StreamingReader>& reader) -> RecordBatchGenerator {
return [reader]() { return reader->ReadNextAsync(); };
});
return MakeFromFuture(generator_fut);
ARROW_ASSIGN_OR_RAISE(auto reader,
OpenReader(source_, *format_, options(), options()->pool));
return IteratorFromReader(std::move(reader));
}

private:
Expand Down