diff --git a/docs/source/python/csv.rst b/docs/source/python/csv.rst index 0f61259dda97..127475cb4cbd 100644 --- a/docs/source/python/csv.rst +++ b/docs/source/python/csv.rst @@ -115,6 +115,7 @@ Available parsing options are: ~ParseOptions.newlines_in_values ~ParseOptions.ignore_empty_lines ~ParseOptions.invalid_row_handler + ~ParseOptions.pad_short_rows .. seealso:: diff --git a/python/pyarrow/_csv.pyx b/python/pyarrow/_csv.pyx index 4789b9532e2e..98bb5b6488bb 100644 --- a/python/pyarrow/_csv.pyx +++ b/python/pyarrow/_csv.pyx @@ -379,6 +379,8 @@ cdef class ParseOptions(_Weakrefable): parsing (because of a mismatching number of columns). It should accept a single InvalidRow argument and return either "skip" or "error" depending on the desired outcome. + pad_short_rows : bool, optional (default False) + Whether rows with fewer columns than expected are padded with nulls. Examples -------- @@ -432,7 +434,8 @@ cdef class ParseOptions(_Weakrefable): def __init__(self, *, delimiter=None, quote_char=None, double_quote=None, escape_char=None, newlines_in_values=None, - ignore_empty_lines=None, invalid_row_handler=None): + ignore_empty_lines=None, invalid_row_handler=None, + pad_short_rows=None): if delimiter is not None: self.delimiter = delimiter if quote_char is not None: @@ -447,6 +450,8 @@ cdef class ParseOptions(_Weakrefable): self.ignore_empty_lines = ignore_empty_lines if invalid_row_handler is not None: self.invalid_row_handler = invalid_row_handler + if pad_short_rows is not None: + self.pad_short_rows = pad_short_rows @property def delimiter(self): @@ -556,6 +561,17 @@ cdef class ParseOptions(_Weakrefable): def ignore_empty_lines(self, value): deref(self.options).ignore_empty_lines = value + @property + def pad_short_rows(self): + """ + Whether rows with fewer columns than expected are padded with nulls. + """ + return deref(self.options).pad_short_rows + + @pad_short_rows.setter + def pad_short_rows(self, value): + deref(self.options).pad_short_rows = value + def validate(self): check_status(deref(self.options).Validate()) @@ -576,7 +592,8 @@ cdef class ParseOptions(_Weakrefable): self.escape_char == other.escape_char and self.newlines_in_values == other.newlines_in_values and self.ignore_empty_lines == other.ignore_empty_lines and - self._invalid_row_handler == other._invalid_row_handler + self._invalid_row_handler == other._invalid_row_handler and + self.pad_short_rows == other.pad_short_rows ) @staticmethod @@ -588,12 +605,20 @@ cdef class ParseOptions(_Weakrefable): def __getstate__(self): return (self.delimiter, self.quote_char, self.double_quote, self.escape_char, self.newlines_in_values, - self.ignore_empty_lines, self.invalid_row_handler) + self.ignore_empty_lines, self.invalid_row_handler, + self.pad_short_rows) def __setstate__(self, state): - (self.delimiter, self.quote_char, self.double_quote, - self.escape_char, self.newlines_in_values, - self.ignore_empty_lines, self.invalid_row_handler) = state + if len(state) == 7: + (self.delimiter, self.quote_char, self.double_quote, + self.escape_char, self.newlines_in_values, + self.ignore_empty_lines, self.invalid_row_handler) = state + self.pad_short_rows = False + else: + (self.delimiter, self.quote_char, self.double_quote, + self.escape_char, self.newlines_in_values, + self.ignore_empty_lines, self.invalid_row_handler, + self.pad_short_rows) = state def __eq__(self, other): try: @@ -610,7 +635,8 @@ cdef class ParseOptions(_Weakrefable): newlines_in_values={self.newlines_in_values}, ignore_empty_lines={self.ignore_empty_lines}, invalid_row_handler={getattr(self.invalid_row_handler, '__name__', - self.invalid_row_handler)}""") + self.invalid_row_handler)}, + pad_short_rows={self.pad_short_rows}""") def __repr__(self): return (f"({self._repr_base()})") diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 296148f52816..300d64d32d63 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -2057,6 +2057,7 @@ cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil: c_bool newlines_in_values c_bool ignore_empty_lines function[CInvalidRowHandler] invalid_row_handler + c_bool pad_short_rows CCSVParseOptions() CCSVParseOptions(CCSVParseOptions) diff --git a/python/pyarrow/tests/test_csv.py b/python/pyarrow/tests/test_csv.py index 275a7e808ddd..870747bd9f6b 100644 --- a/python/pyarrow/tests/test_csv.py +++ b/python/pyarrow/tests/test_csv.py @@ -236,7 +236,8 @@ def test_parse_options(pickle_module): double_quote=[True, False], newlines_in_values=[False, True], ignore_empty_lines=[True, False], - invalid_row_handler=[None, skip_handler]) + invalid_row_handler=[None, skip_handler], + pad_short_rows=[False, True]) check_options_class_pickling(cls, pickler=pickle_module, delimiter='x', @@ -245,7 +246,8 @@ def test_parse_options(pickle_module): double_quote=False, newlines_in_values=True, ignore_empty_lines=False, - invalid_row_handler=skip_handler) + invalid_row_handler=skip_handler, + pad_short_rows=True) cls().validate() opts = cls() @@ -292,11 +294,24 @@ def test_parse_options(pickle_module): escape_char='\r', newlines_in_values=False, ignore_empty_lines=True, - invalid_row_handler=None""" + invalid_row_handler=None, + pad_short_rows=False""" assert repr(opts) == f"({expected_repr_inner})" assert str(opts) == f"ParseOptions({expected_repr_inner})" + # Verify backward compatibility with legacy 7-element state + legacy_state = (',', '"', True, False, False, True, None) + opts_legacy = ParseOptions() + opts_legacy.__setstate__(legacy_state) + assert opts_legacy.pad_short_rows is False + assert opts_legacy.delimiter == ',' + assert opts_legacy.quote_char == '"' + + # Verify equality differentiation + assert ParseOptions() != ParseOptions(pad_short_rows=True) + assert ParseOptions() == ParseOptions(pad_short_rows=False) + def test_convert_options(pickle_module): cls = ConvertOptions @@ -766,6 +781,86 @@ def row_num(x): 'b': ["e", "j"], } + def test_pad_short_rows(self): + # Default behavior: error on short rows + rows = b"a,b,c\n1,2,3\n4,5\n6\n" + with pytest.raises(ValueError, match="Expected 3 columns, got 2: 4,5"): + self.read_bytes(rows, parse_options=ParseOptions(pad_short_rows=False)) + + # Pad short rows with nulls (inferred as int64) + table = self.read_bytes(rows, parse_options=ParseOptions(pad_short_rows=True)) + assert table.to_pydict() == { + "a": [1, 4, 6], + "b": [2, 5, None], + "c": [3, None, None], + } + + # Pad short rows with string values + rows_str = b"a,b,c\nfoo,bar,baz\nqux,quux\ncorge\n" + table_str = self.read_bytes( + rows_str, parse_options=ParseOptions(pad_short_rows=True) + ) + assert table_str.to_pydict() == { + "a": ["foo", "qux", "corge"], + "b": ["bar", "quux", None], + "c": ["baz", None, None], + } + + # Missing trailing fields vs explicitly empty trailing fields + rows_empty = b"a,b,c\nfoo,bar,baz\nqux,quux,\ncorge\n" + table_empty = self.read_bytes( + rows_empty, + parse_options=ParseOptions(pad_short_rows=True), + convert_options=ConvertOptions(strings_can_be_null=False), + ) + assert table_empty.to_pydict() == { + "a": ["foo", "qux", "corge"], + "b": ["bar", "quux", None], + "c": ["baz", "", None], + } + + # Pad short rows with explicit column names + rows_no_header = b"1,2,3\n4,5\n6\n" + read_opts = ReadOptions(column_names=["x", "y", "z"]) + table = self.read_bytes( + rows_no_header, + read_options=read_opts, + parse_options=ParseOptions(pad_short_rows=True), + ) + assert table.to_pydict() == { + "x": [1, 4, 6], + "y": [2, 5, None], + "z": [3, None, None], + } + + # Pad short rows with typed converters + convert_opts = ConvertOptions( + column_types={"a": pa.int64(), "b": pa.float64(), "c": pa.string()} + ) + table = self.read_bytes( + rows, + convert_options=convert_opts, + parse_options=ParseOptions(pad_short_rows=True), + ) + assert table.to_pydict() == { + "a": [1, 4, 6], + "b": [2.0, 5.0, None], + "c": ["3", None, None], + } + + # Rows with excess columns still raise error when ignore_extra_columns=False + rows_extra = b"a,b\n1,2\n3,4,5\n" + with pytest.raises(ValueError, match="Expected 2 columns, got 3"): + self.read_bytes(rows_extra, parse_options=ParseOptions(pad_short_rows=True)) + + # Header-only input + rows_header = b"a,b,c\n" + table_header = self.read_bytes( + rows_header, parse_options=ParseOptions(pad_short_rows=True) + ) + assert table_header.schema.names == ["a", "b", "c"] + assert table_header.num_rows == 0 + def test_chunker_out_of_sync(self): # GH-39892: if there are newlines in values, the parser may become # out of sync with the chunker. In this case, we try to produce an diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 9ddfdfdc9c49..b4f7c982c880 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -1074,6 +1074,7 @@ def test_file_format_pickling(pickle_module): formats = [ ds.IpcFileFormat(), ds.CsvFileFormat(), + ds.CsvFileFormat(pa.csv.ParseOptions(pad_short_rows=True)), ds.CsvFileFormat(pa.csv.ParseOptions(delimiter='\t', ignore_empty_lines=True)), ds.CsvFileFormat(read_options=pa.csv.ReadOptions( @@ -3567,6 +3568,29 @@ def test_csv_format_options_generate_columns(tempdir, dataset_reader): 'f3': pa.array([1])})) +def test_csv_format_pad_short_rows(tempdir, dataset_reader): + path = str(tempdir / 'short.csv') + with open(path, 'w') as sink: + sink.write('a,b,c\n1,2,3\n4,5\n6\n') + + # Default format fails on short rows + dataset_default = ds.dataset(path, format='csv') + with pytest.raises(pa.ArrowInvalid, match="Expected 3 columns, got 2: 4,5"): + dataset_reader.to_table(dataset_default) + + # With pad_short_rows=True, short rows are padded with nulls + format_padded = ds.CsvFileFormat( + parse_options=pa.csv.ParseOptions(pad_short_rows=True) + ) + dataset_padded = ds.dataset(path, format=format_padded) + result = dataset_reader.to_table(dataset_padded) + assert result.to_pydict() == { + 'a': [1, 4, 6], + 'b': [2, 5, None], + 'c': [3, None, None], + } + + def test_csv_fragment_options(tempdir, dataset_reader): path = str(tempdir / 'test.csv') with open(path, 'w') as sink: