public:

using FormatType = CsvFileFormat;

static Result<std::shared_ptr<Buffer>> Write(RecordBatchReader* reader) {

ARROW_ASSIGN_OR_RAISE(auto sink, io::BufferOutputStream::Create());

ARROW_ASSIGN_OR_RAISE(auto table, reader->ToTable());

auto options = csv::WriteOptions::Defaults();

RETURN_NOT_OK(csv::WriteCSV(*table, options, sink.get()));

return sink->Finish();

}

static std::shared_ptr<CsvFileFormat> MakeFormat() {

auto format = std::make_shared<CsvFileFormat>();

// Required for CountRows (since the test generates data with nulls that get written

// as empty lines)

format->parse_options.ignore_empty_lines = false;

return format;

}

public ::testing::WithParamInterface<CsvFileFormatParams> {

public:

bool UseScanV2() { return GetParam().use_new_scan_v2; }

Compression::type GetCompression() { return GetParam().compression_type; }

void SetUp() {

internal::Initialize();

auto fragment_scan_options = std::make_shared<CsvFragmentScanOptions>();

fragment_scan_options->parse_options.ignore_empty_lines = false;

opts_->fragment_scan_options = fragment_scan_options;

}

std::unique_ptr<FileSource> GetFileSource(std::string csv) {

if (GetCompression() == Compression::UNCOMPRESSED) {

return std::make_unique<FileSource>(Buffer::FromString(std::move(csv)));

}

std::string path = "test.csv";

switch (GetCompression()) {

case Compression::type::GZIP:

path += ".gz";

break;

case Compression::type::ZSTD:

path += ".zstd";

break;

case Compression::type::LZ4_FRAME:

path += ".lz4";

break;

case Compression::type::BZ2:

path += ".bz2";

break;

default:

// No known extension

break;

}

EXPECT_OK_AND_ASSIGN(auto fs, fs::internal::MockFileSystem::Make(fs::kNoTime, {}));

EXPECT_OK_AND_ASSIGN(auto codec, util::Codec::Create(GetCompression()));

EXPECT_OK_AND_ASSIGN(auto buffer_writer, fs->OpenOutputStream(path));

EXPECT_OK_AND_ASSIGN(auto stream,

io::CompressedOutputStream::Make(codec.get(), buffer_writer));

ARROW_EXPECT_OK(stream->Write(csv));

ARROW_EXPECT_OK(stream->Close());

EXPECT_OK_AND_ASSIGN(auto info, fs->GetFileInfo(path));

return std::make_unique<FileSource>(info, fs, GetCompression());

}

CsvFragmentScanOptions MakeDefaultFormatOptions() {

CsvFragmentScanOptions scan_opts;

scan_opts.parse_options.ignore_empty_lines = false;

return scan_opts;

}

ScanV2Options MigrateLegacyOptions(std::shared_ptr<Fragment> fragment) {

std::shared_ptr<Dataset> dataset = std::make_shared<FragmentDataset>(

opts_->dataset_schema, FragmentVector{std::move(fragment)});

ScanV2Options updated(std::move(dataset));

updated.format_options = opts_->fragment_scan_options.get();

updated.filter = opts_->filter;

for (const auto& field : opts_->projected_schema->fields()) {

auto field_name = field->name();

EXPECT_OK_AND_ASSIGN(FieldPath field_path,

FieldRef(field_name).FindOne(*opts_->dataset_schema));

updated.columns.push_back(field_path);

}

return updated;

}

RecordBatchIterator Batches(const std::shared_ptr<Fragment>& fragment) {

if (UseScanV2()) {

ScanV2Options v2_options = MigrateLegacyOptions(fragment);

EXPECT_TRUE(ScanV2Options::AddFieldsNeededForFilter(&v2_options).ok());

EXPECT_OK_AND_ASSIGN(std::unique_ptr<RecordBatchReader> reader,

compute::DeclarationToReader(

compute::Declaration("scan2", std::move(v2_options)),

/*use_threads=*/false));

struct ReaderIterator {

Result<std::shared_ptr<RecordBatch>> Next() { return reader->Next(); }

std::unique_ptr<RecordBatchReader> reader;

};

return RecordBatchIterator(ReaderIterator{std::move(reader)});

} else {

EXPECT_OK_AND_ASSIGN(auto batch_gen, fragment->ScanBatchesAsync(opts_));

return MakeGeneratorIterator(batch_gen);

}

}

// ARROW-17382: quoted headers after a BOM should be parsed correctly

auto source = GetFileSource("\xef\xbb\xbf\"ab\",\"cd\"\nef,gh\nij,kl\n");

auto fields = {field("ab", utf8()), field("cd", utf8())};

SetSchema(fields);

auto fragment = MakeFragment(*source);

int64_t row_count = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

AssertSchemaEqual(batch->schema(), schema(fields));

row_count += batch->num_rows();

}

ASSERT_EQ(row_count, 2);

auto source = GetFileSource(R"(f64

2)");

SetSchema({field("f64", float64())});

auto fragment = MakeFragment(*source);

int64_t row_count = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

row_count += batch->num_rows();

}

ASSERT_EQ(row_count, 4);

auto source = GetFileSource(R"(str

bar)");

SetSchema({field("str", utf8())});

auto fragment = MakeFragment(*source);

auto fragment_scan_options =

static_cast<CsvFragmentScanOptions*>(opts_->fragment_scan_options.get());

fragment_scan_options->convert_options.null_values = {"MYNULL"};

fragment_scan_options->convert_options.strings_can_be_null = true;

int64_t null_count = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

null_count += batch->GetColumnByName("str")->null_count();

}

ASSERT_EQ(null_count, 1);

auto source = GetFileSource(R"(header_skipped

bar)");

{

SetSchema({field("str", utf8())});

auto defaults = std::make_shared<CsvFragmentScanOptions>();

defaults->read_options.skip_rows = 1;

format_->default_fragment_scan_options = defaults;

opts_->fragment_scan_options = nullptr;

auto fragment = MakeFragment(*source);

ASSERT_OK_AND_ASSIGN(auto physical_schema, fragment->ReadPhysicalSchema());

AssertSchemaEqual(opts_->dataset_schema, physical_schema);

int64_t rows = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

rows += batch->GetColumnByName("str")->length();

}

ASSERT_EQ(rows, 4);

}

{

SetSchema({field("header_skipped", utf8())});

// These options completely override the default ones

opts_->fragment_scan_options = std::make_shared<CsvFragmentScanOptions>();

auto fragment_scan_options =

static_cast<CsvFragmentScanOptions*>(opts_->fragment_scan_options.get());

fragment_scan_options->read_options.block_size = 1 << 22;

int64_t rows = 0;

auto fragment = MakeFragment(*source);

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

rows += batch->GetColumnByName("header_skipped")->length();

}

ASSERT_EQ(rows, 5);

}

{

SetSchema({field("custom_header", utf8())});

auto defaults = std::make_shared<CsvFragmentScanOptions>();

defaults->read_options.column_names = {"custom_header"};

format_->default_fragment_scan_options = defaults;

opts_->fragment_scan_options = nullptr;

int64_t rows = 0;

auto fragment = MakeFragment(*source);

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

rows += batch->GetColumnByName("custom_header")->length();

}

ASSERT_EQ(rows, 6);

}

auto source = GetFileSource("1,1\n2,3");

SetSchema({field("ints_1", int64()), field("ints_2", int64())});

auto defaults = std::make_shared<CsvFragmentScanOptions>();

defaults->read_options.column_names = {"ints_1", "ints_2"};

format_->default_fragment_scan_options = defaults;

opts_->fragment_scan_options = nullptr;

auto fragment = MakeFragment(*source);

ASSERT_OK_AND_ASSIGN(auto physical_schema, fragment->ReadPhysicalSchema());

AssertSchemaEqual(opts_->dataset_schema, physical_schema);

int64_t rows = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

rows += batch->num_rows();

}

ASSERT_EQ(rows, 2);

defaults->read_options.column_names = {"same", "same"};

format_->default_fragment_scan_options = defaults;

fragment = MakeFragment(*source);

SetSchema({field("same", int64())});

if (UseScanV2()) {

// V2 scan method's basic evolution strategy builds ds_to_frag_map and just finds

// the first instance of a matching column and doesn't check further to see if

// there are duplicates. So in this case it would grab the first column.

//

// Not clear if this is a good thing or not.

rows = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

rows += batch->num_rows();

}

ASSERT_EQ(rows, 2);

} else {

// Legacy scan method does not support CSV columns with duplicate names

EXPECT_RAISES_WITH_MESSAGE_THAT(

Invalid, ::testing::HasSubstr("CSV file contained multiple columns named same"),

Batches(fragment).Next());

}

auto source = GetFileSource(R"(f64

2)");

// NB: dataset_schema includes a column not present in the file

SetSchema({field("f64", float64()), field("virtual", int32())});

auto fragment = MakeFragment(*source);

ASSERT_OK_AND_ASSIGN(auto physical_schema, fragment->ReadPhysicalSchema());

AssertSchemaEqual(Schema({field("f64", float64())}), *physical_schema);

int64_t row_count = 0;

for (auto maybe_batch : Batches(fragment)) {

ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);

if (UseScanV2()) {

// In the new scan, evolution happens and inserts a null column in place of the

// virtual column

AssertSchemaEqual(*batch->schema(), *opts_->dataset_schema);

} else {

AssertSchemaEqual(*batch->schema(), *physical_schema);

}

row_count += batch->num_rows();

}

ASSERT_EQ(row_count, 4);

TestInspect();

auto source = GetFileSource(R"(f64

2)");

ASSERT_OK_AND_ASSIGN(auto actual, format_->Inspect(*source.get()));

EXPECT_EQ(*actual, Schema({field("f64", float64())}));

// Regression test for ARROW-12083

auto source = GetFileSource(R"(actually_string

2)");

auto defaults = std::make_shared<CsvFragmentScanOptions>();

format_->default_fragment_scan_options = defaults;

ASSERT_OK_AND_ASSIGN(auto actual, format_->Inspect(*source.get()));

// Default type inferred

EXPECT_EQ(*actual, Schema({field("actually_string", float64())}));

// Override the inferred type

defaults->convert_options.column_types["actually_string"] = utf8();

ASSERT_OK_AND_ASSIGN(actual, format_->Inspect(*source.get()));

EXPECT_EQ(*actual, Schema({field("actually_string", utf8())}));

TestIsSupported();

bool supported;

auto source = GetFileSource("");

ASSERT_OK_AND_ASSIGN(supported, format_->IsSupported(*source));

ASSERT_EQ(supported, false);

source = GetFileSource(R"(declare,two

1,2,3)");

ASSERT_OK_AND_ASSIGN(supported, format_->IsSupported(*source));

ASSERT_EQ(supported, false);

source = GetFileSource(R"(f64

2)");

ASSERT_OK_AND_ASSIGN(supported, format_->IsSupported(*source));

EXPECT_EQ(supported, true);

auto source = GetFileSource(R"(betrayal_not_really_f64,str

2,baz)");

auto fragment = MakeFragment(*source);

ASSERT_OK_AND_ASSIGN(auto physical_schema, fragment->ReadPhysicalSchema());

AssertSchemaEqual(

Schema({field("betrayal_not_really_f64", float64()), field("str", utf8())}),

*physical_schema);

// CSV is a text format, so it is valid to read column betrayal_not_really_f64 as string

// rather than double

auto not_float64 = utf8();

auto dataset_schema =

schema({field("betrayal_not_really_f64", not_float64), field("str", utf8())});

ScannerBuilder builder(dataset_schema, fragment, opts_);

// This filter is valid with declared schema, but would *not* be valid

// if betrayal_not_really_f64 were read as double rather than string.

ASSERT_OK(

builder.Filter(equal(field_ref("betrayal_not_really_f64"), field_ref("str"))));

// project only "str"

ASSERT_OK(builder.Project({"str"}));

ASSERT_OK_AND_ASSIGN(auto scanner, builder.Finish());

ASSERT_OK_AND_ASSIGN(auto batch_it, scanner->ScanBatches());

// Run through the scan checking for errors to ensure that "f64" is read with the

// specified type and does not revert to the inferred type (if it reverts to

// inferring float64 then evaluation of the comparison expression should break)

ASSERT_OK(batch_it.Visit([](TaggedRecordBatch) { return Status::OK(); }));

auto options =

checked_pointer_cast<CsvFileWriteOptions>(format_->DefaultWriteOptions());

options->write_options->include_header = false;

auto data_schema = schema({field("f64", float64())});

ASSERT_OK_AND_ASSIGN(auto sink, GetFileSink());

ASSERT_OK_AND_ASSIGN(auto fs, fs::internal::MockFileSystem::Make(fs::kNoTime, {}));

ASSERT_OK_AND_ASSIGN(auto writer,

format_->MakeWriter(sink, data_schema, options, {fs, "<buffer>"}));

ASSERT_OK(writer->Write(ConstantArrayGenerator::Zeroes(5, data_schema)));

ASSERT_FINISHES_OK(writer->Finish());

ASSERT_OK_AND_ASSIGN(auto written, sink->Finish());

ASSERT_EQ("0\n0\n0\n0\n0\n", written->ToString());

void SetUp() override {

internal::Initialize();

scan_options_.parse_options.ignore_empty_lines = false;

}

const FragmentScanOptions* GetFormatOptions() override { return &scan_options_; }

protected:

CsvFragmentScanOptions scan_options_;