| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import bz2 |
| from datetime import datetime |
| from decimal import Decimal |
| import gzip |
| import io |
| import itertools |
| import os |
| import shutil |
| import string |
| import tempfile |
| import unittest |
| |
| import pytest |
| |
| import numpy as np |
| |
| import pyarrow as pa |
| from pyarrow.csv import read_csv, ReadOptions, ParseOptions, ConvertOptions |
| |
| |
| def generate_col_names(): |
| # 'a', 'b'... 'z', then 'aa', 'ab'... |
| letters = string.ascii_lowercase |
| for letter in letters: |
| yield letter |
| for first in letter: |
| for second in letter: |
| yield first + second |
| |
| |
| def make_random_csv(num_cols=2, num_rows=10, linesep=u'\r\n'): |
| arr = np.random.RandomState(42).randint(0, 1000, size=(num_cols, num_rows)) |
| col_names = list(itertools.islice(generate_col_names(), num_cols)) |
| csv = io.StringIO() |
| csv.write(u",".join(col_names)) |
| csv.write(linesep) |
| for row in arr.T: |
| csv.write(u",".join(map(str, row))) |
| csv.write(linesep) |
| csv = csv.getvalue().encode() |
| columns = [pa.array(a, type=pa.int64()) for a in arr] |
| expected = pa.Table.from_arrays(columns, col_names) |
| return csv, expected |
| |
| |
| def test_read_options(): |
| cls = ReadOptions |
| opts = cls() |
| |
| assert opts.block_size > 0 |
| opts.block_size = 12345 |
| assert opts.block_size == 12345 |
| |
| assert opts.use_threads is True |
| opts.use_threads = False |
| assert opts.use_threads is False |
| |
| opts = cls(block_size=1234, use_threads=False) |
| assert opts.block_size == 1234 |
| assert opts.use_threads is False |
| |
| |
| def test_parse_options(): |
| cls = ParseOptions |
| opts = cls() |
| assert opts.delimiter == ',' |
| assert opts.quote_char == '"' |
| assert opts.double_quote is True |
| assert opts.escape_char is False |
| assert opts.header_rows == 1 |
| assert opts.newlines_in_values is False |
| assert opts.ignore_empty_lines is True |
| |
| opts.delimiter = 'x' |
| assert opts.delimiter == 'x' |
| assert opts.quote_char == '"' |
| |
| opts.escape_char = 'z' |
| assert opts.escape_char == 'z' |
| assert opts.quote_char == '"' |
| |
| opts.quote_char = False |
| assert opts.quote_char is False |
| assert opts.escape_char == 'z' |
| |
| opts.escape_char = False |
| assert opts.escape_char is False |
| assert opts.quote_char is False |
| |
| opts.newlines_in_values = True |
| assert opts.newlines_in_values is True |
| |
| opts.ignore_empty_lines = False |
| assert opts.ignore_empty_lines is False |
| |
| opts.header_rows = 2 |
| assert opts.header_rows == 2 |
| |
| opts = cls(delimiter=';', quote_char='%', double_quote=False, |
| escape_char='\\', header_rows=2, newlines_in_values=True, |
| ignore_empty_lines=False) |
| assert opts.delimiter == ';' |
| assert opts.quote_char == '%' |
| assert opts.double_quote is False |
| assert opts.escape_char == '\\' |
| assert opts.header_rows == 2 |
| assert opts.newlines_in_values is True |
| assert opts.ignore_empty_lines is False |
| |
| |
| def test_convert_options(): |
| cls = ConvertOptions |
| opts = cls() |
| |
| assert opts.check_utf8 is True |
| opts.check_utf8 = False |
| assert opts.check_utf8 is False |
| |
| assert opts.strings_can_be_null is False |
| opts.strings_can_be_null = True |
| assert opts.strings_can_be_null is True |
| |
| assert opts.column_types == {} |
| # Pass column_types as mapping |
| opts.column_types = {'b': pa.int16(), 'c': pa.float32()} |
| assert opts.column_types == {'b': pa.int16(), 'c': pa.float32()} |
| opts.column_types = {'v': 'int16', 'w': 'null'} |
| assert opts.column_types == {'v': pa.int16(), 'w': pa.null()} |
| # Pass column_types as schema |
| schema = pa.schema([('a', pa.int32()), ('b', pa.string())]) |
| opts.column_types = schema |
| assert opts.column_types == {'a': pa.int32(), 'b': pa.string()} |
| # Pass column_types as sequence |
| opts.column_types = [('x', pa.binary())] |
| assert opts.column_types == {'x': pa.binary()} |
| |
| with pytest.raises(TypeError, match='DataType expected'): |
| opts.column_types = {'a': None} |
| with pytest.raises(TypeError): |
| opts.column_types = 0 |
| |
| assert isinstance(opts.null_values, list) |
| assert '' in opts.null_values |
| assert 'N/A' in opts.null_values |
| opts.null_values = ['xxx', 'yyy'] |
| assert opts.null_values == ['xxx', 'yyy'] |
| |
| assert isinstance(opts.true_values, list) |
| opts.true_values = ['xxx', 'yyy'] |
| assert opts.true_values == ['xxx', 'yyy'] |
| |
| assert isinstance(opts.false_values, list) |
| opts.false_values = ['xxx', 'yyy'] |
| assert opts.false_values == ['xxx', 'yyy'] |
| |
| opts = cls(check_utf8=False, column_types={'a': pa.null()}, |
| null_values=['N', 'nn'], true_values=['T', 'tt'], |
| false_values=['F', 'ff'], strings_can_be_null=True) |
| assert opts.check_utf8 is False |
| assert opts.column_types == {'a': pa.null()} |
| assert opts.null_values == ['N', 'nn'] |
| assert opts.false_values == ['F', 'ff'] |
| assert opts.true_values == ['T', 'tt'] |
| assert opts.strings_can_be_null is True |
| |
| |
| class BaseTestCSVRead: |
| |
| def read_bytes(self, b, **kwargs): |
| return self.read_csv(pa.py_buffer(b), **kwargs) |
| |
| def check_names(self, table, names): |
| assert table.num_columns == len(names) |
| assert [c.name for c in table.columns] == names |
| |
| def test_file_object(self): |
| data = b"a,b\n1,2\n" |
| expected_data = {'a': [1], 'b': [2]} |
| bio = io.BytesIO(data) |
| table = self.read_csv(bio) |
| assert table.to_pydict() == expected_data |
| # Text files not allowed |
| sio = io.StringIO(data.decode()) |
| with pytest.raises(TypeError): |
| self.read_csv(sio) |
| |
| def test_header(self): |
| rows = b"abc,def,gh\n" |
| table = self.read_bytes(rows) |
| assert isinstance(table, pa.Table) |
| self.check_names(table, ["abc", "def", "gh"]) |
| assert table.num_rows == 0 |
| |
| def test_bom(self): |
| rows = b"\xef\xbb\xbfa,b\n1,2\n" |
| expected_data = {'a': [1], 'b': [2]} |
| table = self.read_bytes(rows) |
| assert table.to_pydict() == expected_data |
| |
| def test_simple_ints(self): |
| # Infer integer columns |
| rows = b"a,b,c\n1,2,3\n4,5,6\n" |
| table = self.read_bytes(rows) |
| schema = pa.schema([('a', pa.int64()), |
| ('b', pa.int64()), |
| ('c', pa.int64())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [1, 4], |
| 'b': [2, 5], |
| 'c': [3, 6], |
| } |
| |
| def test_simple_varied(self): |
| # Infer various kinds of data |
| rows = b"a,b,c,d\n1,2,3,0\n4.0,-5,foo,True\n" |
| table = self.read_bytes(rows) |
| schema = pa.schema([('a', pa.float64()), |
| ('b', pa.int64()), |
| ('c', pa.string()), |
| ('d', pa.bool_())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [1.0, 4.0], |
| 'b': [2, -5], |
| 'c': [u"3", u"foo"], |
| 'd': [False, True], |
| } |
| |
| def test_simple_nulls(self): |
| # Infer various kinds of data, with nulls |
| rows = (b"a,b,c,d,e,f\n" |
| b"1,2,,,3,N/A\n" |
| b"nan,-5,foo,,nan,TRUE\n" |
| b"4.5,#N/A,nan,,\xff,false\n") |
| table = self.read_bytes(rows) |
| schema = pa.schema([('a', pa.float64()), |
| ('b', pa.int64()), |
| ('c', pa.string()), |
| ('d', pa.null()), |
| ('e', pa.binary()), |
| ('f', pa.bool_())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [1.0, None, 4.5], |
| 'b': [2, -5, None], |
| 'c': [u"", u"foo", u"nan"], |
| 'd': [None, None, None], |
| 'e': [b"3", b"nan", b"\xff"], |
| 'f': [None, True, False], |
| } |
| |
| def test_simple_timestamps(self): |
| # Infer a timestamp column |
| rows = b"a,b\n1970,1970-01-01\n1989,1989-07-14\n" |
| table = self.read_bytes(rows) |
| schema = pa.schema([('a', pa.int64()), |
| ('b', pa.timestamp('s'))]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [1970, 1989], |
| 'b': [datetime(1970, 1, 1), datetime(1989, 7, 14)], |
| } |
| |
| def test_custom_nulls(self): |
| # Infer nulls with custom values |
| opts = ConvertOptions(null_values=['Xxx', 'Zzz']) |
| rows = b"a,b,c,d\nZzz,Xxx,1,2\nXxx,#N/A,,Zzz\n" |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.null()), |
| ('b', pa.string()), |
| ('c', pa.string()), |
| ('d', pa.int64())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [None, None], |
| 'b': [u"Xxx", u"#N/A"], |
| 'c': [u"1", u""], |
| 'd': [2, None], |
| } |
| |
| opts = ConvertOptions(null_values=['Xxx', 'Zzz'], |
| strings_can_be_null=True) |
| table = self.read_bytes(rows, convert_options=opts) |
| assert table.to_pydict() == { |
| 'a': [None, None], |
| 'b': [None, u"#N/A"], |
| 'c': [u"1", u""], |
| 'd': [2, None], |
| } |
| |
| opts = ConvertOptions(null_values=[]) |
| rows = b"a,b\n#N/A,\n" |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.string()), |
| ('b', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [u"#N/A"], |
| 'b': [u""], |
| } |
| |
| def test_custom_bools(self): |
| # Infer booleans with custom values |
| opts = ConvertOptions(true_values=['T', 'yes'], |
| false_values=['F', 'no']) |
| rows = (b"a,b,c\n" |
| b"True,T,t\n" |
| b"False,F,f\n" |
| b"True,yes,yes\n" |
| b"False,no,no\n" |
| b"N/A,N/A,N/A\n") |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.string()), |
| ('b', pa.bool_()), |
| ('c', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': ["True", "False", "True", "False", "N/A"], |
| 'b': [True, False, True, False, None], |
| 'c': ["t", "f", "yes", "no", "N/A"], |
| } |
| |
| def test_column_types(self): |
| # Ask for specific column types in ConvertOptions |
| opts = ConvertOptions(column_types={'b': 'float32', |
| 'c': 'string', |
| 'd': 'boolean', |
| 'e': pa.decimal128(11, 2), |
| 'zz': 'null'}) |
| rows = b"a,b,c,d,e\n1,2,3,true,1.0\n4,-5,6,false,0\n" |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.int64()), |
| ('b', pa.float32()), |
| ('c', pa.string()), |
| ('d', pa.bool_()), |
| ('e', pa.decimal128(11, 2))]) |
| expected = { |
| 'a': [1, 4], |
| 'b': [2.0, -5.0], |
| 'c': ["3", "6"], |
| 'd': [True, False], |
| 'e': [Decimal("1.00"), Decimal("0.00")] |
| } |
| assert table.schema == schema |
| assert table.to_pydict() == expected |
| # Pass column_types as schema |
| opts = ConvertOptions( |
| column_types=pa.schema([('b', pa.float32()), |
| ('c', pa.string()), |
| ('d', pa.bool_()), |
| ('e', pa.decimal128(11, 2)), |
| ('zz', pa.bool_())])) |
| table = self.read_bytes(rows, convert_options=opts) |
| assert table.schema == schema |
| assert table.to_pydict() == expected |
| # One of the columns in column_types fails converting |
| rows = b"a,b,c,d,e\n1,XXX,3,true,5\n4,-5,6,false,7\n" |
| with pytest.raises(pa.ArrowInvalid) as exc: |
| self.read_bytes(rows, convert_options=opts) |
| err = str(exc.value) |
| assert "In column #1: " in err |
| assert "CSV conversion error to float: invalid value 'XXX'" in err |
| |
| def test_no_ending_newline(self): |
| # No \n after last line |
| rows = b"a,b,c\n1,2,3\n4,5,6" |
| table = self.read_bytes(rows) |
| assert table.to_pydict() == { |
| 'a': [1, 4], |
| 'b': [2, 5], |
| 'c': [3, 6], |
| } |
| |
| def test_trivial(self): |
| # A bit pointless, but at least it shouldn't crash |
| rows = b",\n\n" |
| table = self.read_bytes(rows) |
| assert table.to_pydict() == {'': []} |
| |
| def test_invalid_csv(self): |
| # Various CSV errors |
| rows = b"a,b,c\n1,2\n4,5,6\n" |
| with pytest.raises(pa.ArrowInvalid, match="Expected 3 columns, got 2"): |
| self.read_bytes(rows) |
| rows = b"a,b,c\n1,2,3\n4" |
| with pytest.raises(pa.ArrowInvalid, match="Expected 3 columns, got 1"): |
| self.read_bytes(rows) |
| for rows in [b"", b"\n", b"\r\n", b"\r", b"\n\n"]: |
| with pytest.raises(pa.ArrowInvalid, match="Empty CSV file"): |
| self.read_bytes(rows) |
| |
| def test_options_delimiter(self): |
| rows = b"a;b,c\nde,fg;eh\n" |
| table = self.read_bytes(rows) |
| assert table.to_pydict() == { |
| 'a;b': [u'de'], |
| 'c': [u'fg;eh'], |
| } |
| opts = ParseOptions(delimiter=';') |
| table = self.read_bytes(rows, parse_options=opts) |
| assert table.to_pydict() == { |
| 'a': [u'de,fg'], |
| 'b,c': [u'eh'], |
| } |
| |
| def test_small_random_csv(self): |
| csv, expected = make_random_csv(num_cols=2, num_rows=10) |
| table = self.read_bytes(csv) |
| assert table.schema == expected.schema |
| assert table.equals(expected) |
| assert table.to_pydict() == expected.to_pydict() |
| |
| def test_stress_block_sizes(self): |
| # Test a number of small block sizes to stress block stitching |
| csv_base, expected = make_random_csv(num_cols=2, num_rows=500) |
| block_sizes = [11, 12, 13, 17, 37, 111] |
| csvs = [csv_base, csv_base.rstrip(b'\r\n')] |
| for csv in csvs: |
| for block_size in block_sizes: |
| read_options = ReadOptions(block_size=block_size) |
| table = self.read_bytes(csv, read_options=read_options) |
| assert table.schema == expected.schema |
| if not table.equals(expected): |
| # Better error output |
| assert table.to_pydict() == expected.to_pydict() |
| |
| |
| class TestSerialCSVRead(BaseTestCSVRead, unittest.TestCase): |
| |
| def read_csv(self, *args, **kwargs): |
| read_options = kwargs.setdefault('read_options', ReadOptions()) |
| read_options.use_threads = False |
| table = read_csv(*args, **kwargs) |
| table._validate() |
| return table |
| |
| |
| class TestParallelCSVRead(BaseTestCSVRead, unittest.TestCase): |
| |
| def read_csv(self, *args, **kwargs): |
| read_options = kwargs.setdefault('read_options', ReadOptions()) |
| read_options.use_threads = True |
| table = read_csv(*args, **kwargs) |
| table._validate() |
| return table |
| |
| |
| class BaseTestCompressedCSVRead: |
| |
| def setUp(self): |
| self.tmpdir = tempfile.mkdtemp(prefix='arrow-csv-test-') |
| |
| def tearDown(self): |
| shutil.rmtree(self.tmpdir) |
| |
| def test_random_csv(self): |
| csv, expected = make_random_csv(num_cols=2, num_rows=100) |
| csv_path = os.path.join(self.tmpdir, self.csv_filename) |
| self.write_file(csv_path, csv) |
| try: |
| table = read_csv(csv_path) |
| except pa.ArrowNotImplementedError as e: |
| pytest.skip(str(e)) |
| return |
| table._validate() |
| assert table.schema == expected.schema |
| assert table.equals(expected) |
| assert table.to_pydict() == expected.to_pydict() |
| |
| |
| class TestGZipCSVRead(BaseTestCompressedCSVRead, unittest.TestCase): |
| csv_filename = "compressed.csv.gz" |
| |
| def write_file(self, path, contents): |
| with gzip.open(path, 'wb', 3) as f: |
| f.write(contents) |
| |
| |
| class TestBZ2CSVRead(BaseTestCompressedCSVRead, unittest.TestCase): |
| csv_filename = "compressed.csv.bz2" |
| |
| def write_file(self, path, contents): |
| with bz2.BZ2File(path, 'w') as f: |
| f.write(contents) |
| |
| |
| def test_read_csv_does_not_close_passed_file_handles(): |
| # ARROW-4823 |
| buf = io.BytesIO(b"a,b,c\n1,2,3\n4,5,6") |
| read_csv(buf) |
| assert not buf.closed |
