Repository: apache/arrow
File: cpp/src/arrow/csv/writer_test.cc (C++)

Code before:
       /*expected_output*/ "", expected_status_illegal_delimiter('\r')},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n", /*delimiter=*/'\n'),
       /*expected_output*/ "", expected_status_illegal_delimiter('\n')},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n", /*delimiter=*/'"'),
       /*expected_output*/ "", expected_status_illegal_delimiter('"')},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/";", /*delimiter=*/';'),
       /*expected_output*/ "", expected_status_illegal_delimiter(';')}};
}

class TestWriteCSV : public ::testing::TestWithParam<WriterTestParams> {
 protected:
  template <typename Data>
  Result<std::string> ToCsvString(const Data& data, const WriteOptions& options) {
    std::shared_ptr<io::BufferOutputStream> out;
    ARROW_ASSIGN_OR_RAISE(out, io::BufferOutputStream::Create());

    RETURN_NOT_OK(WriteCSV(data, options, out.get()));
    ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> buffer, out->Finish());
    return std::string(reinterpret_cast<const char*>(buffer->data()), buffer->size());
  }

  Result<std::string> ToCsvStringUsingWriter(const Table& data,
                                             const WriteOptions& options) {
    std::shared_ptr<io::BufferOutputStream> out;
    ARROW_ASSIGN_OR_RAISE(out, io::BufferOutputStream::Create());
    // Write row-by-row
    ARROW_ASSIGN_OR_RAISE(auto writer, MakeCSVWriter(out, data.schema(), options));
    TableBatchReader reader(data);
    reader.set_chunksize(1);
    std::shared_ptr<RecordBatch> batch;
    RETURN_NOT_OK(reader.ReadNext(&batch));
    while (batch != nullptr) {
      RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
      RETURN_NOT_OK(reader.ReadNext(&batch));
    }
    RETURN_NOT_OK(writer->Close());
    EXPECT_EQ(data.num_rows(), writer->stats().num_record_batches);
    ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> buffer, out->Finish());
    return std::string(reinterpret_cast<const char*>(buffer->data()), buffer->size());
  }
};

TEST_P(TestWriteCSV, TestWrite) {
  ASSERT_OK_AND_ASSIGN(std::shared_ptr<io::BufferOutputStream> out,
                       io::BufferOutputStream::Create());

Review comment (suggestion):
It seems that all three `QuotingStyle` values should be tested with a mix of column names requiring quoting and column names that don't need quoting.

Code after:
      reject_structural_params({"0123456\\r789", nullptr, "abcdef"}, "0123456\r789"),
      reject_structural_params({"0123456789", nullptr, "abcde,", nullptr}, "abcde,"),
      reject_structural_params({"0123456789", nullptr, "abcdef,", nullptr}, "abcdef,"),
      reject_structural_params({nullptr, nullptr, ",0123456789", "abcde"}, ",0123456789"),
      reject_structural_params({"0123456", nullptr, "7\\\"89", ",abcdef"}, "7\"89"),
      // exercise custom delimiter
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, "\n", /*delimiter=*/'\t'),
       expected_output_delimiter_tabs},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n", /*delimiter=*/'\r'),
       /*expected_output*/ "", expected_status_illegal_delimiter('\r')},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n", /*delimiter=*/'\n'),
       /*expected_output*/ "", expected_status_illegal_delimiter('\n')},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n", /*delimiter=*/'"'),
       /*expected_output*/ "", expected_status_illegal_delimiter('"')},
      {schema_custom_delimiter, batch_custom_delimiter,
       DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/";", /*delimiter=*/';'),
       /*expected_output*/ "", expected_status_illegal_delimiter(';')},
      {header_without_structural_charaters, "[]",
       DefaultTestOptions(/*include_header=*/true, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n",
                          /*delimiter=*/',', /*batch_size=*/5,
                          /*quoting_header=*/QuotingStyle::None),
       expected_header_without_structural_charaters},
      {abc_schema, "[]",
       DefaultTestOptions(/*include_header=*/true, /*null_string=*/"",
                          QuotingStyle::Needed, /*eol=*/"\n",
                          /*delimiter=*/',', /*batch_size=*/5,
                          /*quoting_header=*/QuotingStyle::None),
       "", expected_status_no_quotes_with_structural_in_header("b\"")},
  };
}

class TestWriteCSV : public ::testing::TestWithParam<WriterTestParams> {
 protected:
  template <typename Data>
  Result<std::string> ToCsvString(const Data& data, const WriteOptions& options) {
    std::shared_ptr<io::BufferOutputStream> out;
    ARROW_ASSIGN_OR_RAISE(out, io::BufferOutputStream::Create());

    RETURN_NOT_OK(WriteCSV(data, options, out.get()));
    ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> buffer, out->Finish());
    return std::string(reinterpret_cast<const char*>(buffer->data()), buffer->size());