diff --git a/python/pyarrow/_csv.pyx b/python/pyarrow/_csv.pyx index f2cefb8ff3fb..4789b9532e2e 100644 --- a/python/pyarrow/_csv.pyx +++ b/python/pyarrow/_csv.pyx @@ -1472,6 +1472,10 @@ cdef class WriteOptions(_Weakrefable): CSV data delimiter : 1-character string, optional (default ",") The character delimiting individual cells in the CSV data. + eol : str, optional (default "\\n") + The end of line character to use for ending rows. + null_string : str, optional (default "") + The string to write for null values. Quotes are not allowed in this string. quoting_style : str, optional (default "needed") Whether to quote values, and if so, which quoting style to use. The following values are accepted: @@ -1490,7 +1494,8 @@ cdef class WriteOptions(_Weakrefable): __slots__ = () def __init__(self, *, include_header=None, batch_size=None, - delimiter=None, quoting_style=None, quoting_header=None): + delimiter=None, eol=None, null_string=None, + quoting_style=None, quoting_header=None): self.options.reset(new CCSVWriteOptions(CCSVWriteOptions.Defaults())) if include_header is not None: self.include_header = include_header @@ -1498,6 +1503,10 @@ cdef class WriteOptions(_Weakrefable): self.batch_size = batch_size if delimiter is not None: self.delimiter = delimiter + if eol is not None: + self.eol = eol + if null_string is not None: + self.null_string = null_string if quoting_style is not None: self.quoting_style = quoting_style if quoting_header is not None: @@ -1537,6 +1546,28 @@ cdef class WriteOptions(_Weakrefable): def delimiter(self, value): deref(self.options).delimiter = _single_char(value) + @property + def eol(self): + """ + The end of line character to use for ending rows. + """ + return frombytes(deref(self.options).eol) + + @eol.setter + def eol(self, value): + deref(self.options).eol = tobytes(value) + + @property + def null_string(self): + """ + The string to write for null values. Quotes are not allowed in this string. + """ + return frombytes(deref(self.options).null_string) + + @null_string.setter + def null_string(self, value): + deref(self.options).null_string = tobytes(value) + @property def quoting_style(self): """ diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index ffc02ffd79ae..d592d4024aa9 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -2183,6 +2183,8 @@ cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil: unsigned char delimiter CQuotingStyle quoting_style CQuotingStyle quoting_header + c_string eol + c_string null_string CIOContext io_context CCSVWriteOptions() diff --git a/python/pyarrow/tests/test_csv.py b/python/pyarrow/tests/test_csv.py index ac9012ebdf63..275a7e808ddd 100644 --- a/python/pyarrow/tests/test_csv.py +++ b/python/pyarrow/tests/test_csv.py @@ -417,6 +417,7 @@ def test_write_options(): check_options_class( cls, include_header=[True, False], delimiter=[',', '\t', '|'], + eol=['\n', '\r\n'], null_string=['', 'NA'], quoting_style=['needed', 'none', 'all_valid']) assert opts.batch_size > 0 @@ -2138,6 +2139,41 @@ def test_write_quoting_header(): buf.seek(0) +def test_write_eol(): + t = pa.Table.from_arrays([[1, 2, 3], ["a", "b", "c"]], ["c1", "c2"]) + buf = io.BytesIO() + for write_options, res in [ + (WriteOptions(), b'"c1","c2"\n1,"a"\n2,"b"\n3,"c"\n'), + (WriteOptions(eol='\n'), b'"c1","c2"\n1,"a"\n2,"b"\n3,"c"\n'), + (WriteOptions(eol='\r\n'), + b'"c1","c2"\r\n1,"a"\r\n2,"b"\r\n3,"c"\r\n'), + (WriteOptions(eol='*'), b'"c1","c2"*1,"a"*2,"b"*3,"c"*'), + ]: + with CSVWriter(buf, t.schema, write_options=write_options) as writer: + writer.write_table(t) + assert buf.getvalue() == res + buf.seek(0) + buf.truncate() + + +def test_write_null_string(): + t = pa.Table.from_arrays([[1, 2, None], ["a", None, "c"]], ["c1", "c2"]) + buf = io.BytesIO() + for write_options, res in [ + (WriteOptions(), b'"c1","c2"\n1,"a"\n2,\n,"c"\n'), + (WriteOptions(null_string=''), b'"c1","c2"\n1,"a"\n2,\n,"c"\n'), + (WriteOptions(null_string='NA'), + b'"c1","c2"\n1,"a"\n2,NA\nNA,"c"\n'), + (WriteOptions(null_string='N/A'), + b'"c1","c2"\n1,"a"\n2,N/A\nN/A,"c"\n'), + ]: + with CSVWriter(buf, t.schema, write_options=write_options) as writer: + writer.write_table(t) + assert buf.getvalue() == res + buf.seek(0) + buf.truncate() + + def test_read_csv_reference_cycle(): # ARROW-13187 def inner():