diff --git a/cpp/src/arrow/csv/options.h b/cpp/src/arrow/csv/options.h index 41d0b63dacbb..fd2ec58e5861 100644 --- a/cpp/src/arrow/csv/options.h +++ b/cpp/src/arrow/csv/options.h @@ -183,8 +183,9 @@ enum class ARROW_EXPORT QuotingStyle { /// Only enclose values in quotes which need them, because their CSV rendering can /// contain quotes itself (e.g. strings or binary values) Needed, - /// Enclose all valid values in quotes. Nulls are not quoted. May cause readers to - /// interpret all values as strings if schema is inferred. + /// Enclose all valid values in quotes. Nulls are not quoted, except as described for + /// `WriteOptions::null_string`. May cause readers to interpret all values as strings + /// if schema is inferred. AllValid, /// Do not enclose any values in quotes. Prevents values from containing quotes ("), /// cell delimiters (,) or line endings (\\r, \\n), (following RFC4180). If values @@ -206,6 +207,13 @@ 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 whose valid values are never + /// written as `""` (anything but a string, binary or zero-width fixed size binary + /// 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. Other columns keep the empty line, since a valid empty + /// value is already written as `""` there. std::string null_string; /// \brief IO context for writing. diff --git a/cpp/src/arrow/csv/writer.cc b/cpp/src/arrow/csv/writer.cc index 2db0dba2de71..72520e6ea6ab 100644 --- a/cpp/src/arrow/csv/writer.cc +++ b/cpp/src/arrow/csv/writer.cc @@ -36,6 +36,7 @@ namespace arrow { +using internal::checked_cast; using internal::checked_pointer_cast; namespace csv { @@ -438,6 +439,19 @@ class QuotedColumnPopulator : public ColumnPopulator { std::vector row_needs_escaping_; }; +// Whether a valid value of this type can be written as an empty quoted string, in +// which case a null written as "" could not be told apart from it. +bool CanWriteEmptyQuotedValue(const DataType& type) { + if (type.id() == Type::DICTIONARY) { + return CanWriteEmptyQuotedValue( + *checked_cast(type).value_type()); + } + if (type.id() == Type::FIXED_SIZE_BINARY) { + return checked_cast(type).byte_width() == 0; + } + return is_base_binary_like(type.id()) || is_binary_view_like(type.id()); +} + Result> MakePopulator( const DataType& type, const std::string& end_chars, const char delimiter, const std::shared_ptr& null_string, QuotingStyle quoting_style, @@ -512,10 +526,15 @@ class CSVWriterImpl : public ipc::RecordBatchWriter { return Status::Invalid("Null string cannot contain quotes."); } + const bool quote_nulls = schema->num_fields() == 1 && options.null_string.empty() && + options.quoting_style != QuotingStyle::None && + !CanWriteEmptyQuotedValue(*schema->field(0)->type()); + const std::string null_string_value = + quote_nulls ? std::string(R"("")") : options.null_string; ARROW_ASSIGN_OR_RAISE(std::shared_ptr 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> populators(schema->num_fields()); std::string delimiter(1, options.delimiter); diff --git a/cpp/src/arrow/csv/writer_test.cc b/cpp/src/arrow/csv/writer_test.cc index ce4d8ab16d01..417419c29114 100644 --- a/cpp/src/arrow/csv/writer_test.cc +++ b/cpp/src/arrow/csv/writer_test.cc @@ -23,6 +23,7 @@ #include #include "arrow/buffer.h" +#include "arrow/csv/reader.h" #include "arrow/csv/writer.h" #include "arrow/io/memory.h" #include "arrow/ipc/writer.h" @@ -384,13 +385,88 @@ 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("s", utf8())}), R"([{ "s": "a"}, {}, { "s": ""}])", + DefaultTestOptions(/*include_header=*/false, /*null_string=*/"", + QuotingStyle::AllValid), + "\"a\"\n\n\"\"\n"), + WriterTestParams(schema({field("d", dictionary(int32(), utf8()))}), + R"([{ "d": "a"}, {}])", DefaultTestOptions(), "\"a\"\n\n"), + WriterTestParams(schema({field("a", int64()), field("b", int64())}), + R"([{ "a": 1}, {}])", DefaultTestOptions(), "1,\n,\n"))); + +Result> RoundTripCSV(const RecordBatch& batch, + const WriteOptions& write_options, + const ParseOptions& parse_options, + const ConvertOptions& convert_options) { + ARROW_ASSIGN_OR_RAISE(auto out, io::BufferOutputStream::Create()); + RETURN_NOT_OK(WriteCSV(batch, write_options, out.get())); + ARROW_ASSIGN_OR_RAISE(std::shared_ptr csv, out->Finish()); + + auto input = std::make_shared(csv); + ARROW_ASSIGN_OR_RAISE(auto reader, TableReader::Make(io::default_io_context(), input, + ReadOptions::Defaults(), + parse_options, convert_options)); + return reader->Read(); +} + +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 expected, Table::FromRecordBatches({batch})); + + for (auto quoting_style : {QuotingStyle::Needed, QuotingStyle::AllValid}) { + auto write_options = WriteOptions::Defaults(); + write_options.quoting_style = quoting_style; + ASSERT_OK_AND_ASSIGN(auto table, + RoundTripCSV(*batch, write_options, ParseOptions::Defaults(), + ConvertOptions::Defaults())); + AssertTablesEqual(*expected, *table); + } +} + +TEST(TestWriteCSV, SingleColumnStringNullsRoundTrip) { + auto schema = arrow::schema({field("s", utf8())}); + auto batch = RecordBatchFromJSON(schema, R"([{"s": "a"}, {}, {"s": ""}])"); + ASSERT_OK_AND_ASSIGN(auto expected, Table::FromRecordBatches({batch})); + + auto parse_options = ParseOptions::Defaults(); + parse_options.ignore_empty_lines = false; + auto convert_options = ConvertOptions::Defaults(); + convert_options.strings_can_be_null = true; + convert_options.quoted_strings_can_be_null = false; + + for (auto quoting_style : {QuotingStyle::Needed, QuotingStyle::AllValid}) { + auto write_options = WriteOptions::Defaults(); + write_options.quoting_style = quoting_style; + ASSERT_OK_AND_ASSIGN( + auto table, RoundTripCSV(*batch, write_options, parse_options, convert_options)); + AssertTablesEqual(*expected, *table); + } +} #ifndef _WIN32 // TODO(ARROW-13168):