Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 4 additions & 0 deletions cpp/src/arrow/csv/options.h
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,10 @@ struct ARROW_EXPORT WriteOptions {
char delimiter = ',';

/// \brief The string to write for null values. Quotes are not allowed in this string.
///
/// If this is empty and the schema has a single column, nulls are written as `""`
/// (unless `quoting_style` is `QuotingStyle::None`), so that a null row is not
/// written as an empty line, which readers skip by default.
std::string null_string;

/// \brief IO context for writing.
Expand Down
11 changes: 8 additions & 3 deletions cpp/src/arrow/csv/writer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -512,10 +512,15 @@ class CSVWriterImpl : public ipc::RecordBatchWriter {
return Status::Invalid("Null string cannot contain quotes.");
}

const std::string null_string_value =
schema->num_fields() == 1 && options.null_string.empty() &&
options.quoting_style != QuotingStyle::None
? std::string(R"("")")
: options.null_string;
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> null_string,
arrow::AllocateBuffer(options.null_string.length()));
memcpy(null_string->mutable_data(), options.null_string.data(),
options.null_string.length());
arrow::AllocateBuffer(null_string_value.length()));
memcpy(null_string->mutable_data(), null_string_value.data(),
null_string_value.length());

std::vector<std::unique_ptr<ColumnPopulator>> populators(schema->num_fields());
std::string delimiter(1, options.delimiter);
Expand Down
51 changes: 44 additions & 7 deletions cpp/src/arrow/csv/writer_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
#include <vector>

#include "arrow/buffer.h"
#include "arrow/csv/reader.h"
#include "arrow/csv/writer.h"
#include "arrow/io/memory.h"
#include "arrow/ipc/writer.h"
Expand Down Expand Up @@ -384,13 +385,49 @@ TEST_P(TestWriteCSV, TestWrite) {
INSTANTIATE_TEST_SUITE_P(MultiColumnWriteCSVTest, TestWriteCSV,
::testing::ValuesIn(GenerateTestCases()));

INSTANTIATE_TEST_SUITE_P(SingleColumnWriteCSVTest, TestWriteCSV,
::testing::Values(WriterTestParams(
schema({field("int64", int64())}),
R"([{ "int64": 9999}, {}, { "int64": -15}])", WriteOptions(),
R"("int64")"
"\n9999\n\n-15\n",
Status::OK())));
INSTANTIATE_TEST_SUITE_P(
SingleColumnWriteCSVTest, TestWriteCSV,
::testing::Values(
WriterTestParams(schema({field("int64", int64())}),
R"([{ "int64": 9999}, {}, { "int64": -15}])", WriteOptions(),
R"("int64")"
"\n9999\n\"\"\n-15\n",
Status::OK()),
WriterTestParams(schema({field("int64", int64())}), R"([{}, { "int64": 1}, {}])",
DefaultTestOptions(), "\"\"\n1\n\"\"\n"),
WriterTestParams(schema({field("int64", int64())}), R"([{}, { "int64": 1}, {}])",
DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
QuotingStyle::AllValid),
"\"\"\n\"1\"\n\"\"\n"),
WriterTestParams(schema({field("int64", int64())}), R"([{}, { "int64": 1}, {}])",
DefaultTestOptions(/*include_header=*/false, /*null_string=*/"",
QuotingStyle::None),
"\n1\n\n"),
WriterTestParams(schema({field("int64", int64())}), R"([{}, { "int64": 1}, {}])",
DefaultTestOptions(/*include_header=*/false,
/*null_string=*/"NA"),
"NA\n1\nNA\n"),
WriterTestParams(schema({field("s", utf8())}), R"([{ "s": "a"}, {}, { "s": ""}])",
DefaultTestOptions(), "\"a\"\n\"\"\n\"\"\n"),
WriterTestParams(schema({field("a", int64()), field("b", int64())}),
R"([{ "a": 1}, {}])", DefaultTestOptions(), "1,\n,\n")));

TEST(TestWriteCSV, SingleColumnNullsRoundTrip) {
auto schema = arrow::schema({field("c", int64())});
auto batch = RecordBatchFromJSON(schema, R"([{"c": 1}, {}, {"c": 3}, {}, {"c": 5}])");
ASSERT_OK_AND_ASSIGN(auto out, io::BufferOutputStream::Create());
ASSERT_OK(WriteCSV(*batch, WriteOptions::Defaults(), out.get()));
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Buffer> csv, out->Finish());

auto input = std::make_shared<io::BufferReader>(csv);
ASSERT_OK_AND_ASSIGN(
auto reader,
TableReader::Make(io::default_io_context(), input, ReadOptions::Defaults(),
ParseOptions::Defaults(), ConvertOptions::Defaults()));
ASSERT_OK_AND_ASSIGN(auto table, reader->Read());
ASSERT_OK_AND_ASSIGN(auto expected, Table::FromRecordBatches({batch}));
AssertTablesEqual(*expected, *table);
}

#ifndef _WIN32
// TODO(ARROW-13168):
Expand Down
Loading