Skip to content
Merged
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
33 changes: 32 additions & 1 deletion python/pyarrow/_csv.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -1490,14 +1494,19 @@ 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
if batch_size is not None:
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:
Expand Down Expand Up @@ -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):
"""
Expand Down
2 changes: 2 additions & 0 deletions python/pyarrow/includes/libarrow.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
36 changes: 36 additions & 0 deletions python/pyarrow/tests/test_csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand Down