| # 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 abc |
| import bz2 |
| from datetime import date, datetime |
| from decimal import Decimal |
| import gc |
| import gzip |
| import io |
| import itertools |
| import os |
| import pickle |
| import select |
| import shutil |
| import signal |
| import string |
| import tempfile |
| import threading |
| import time |
| import unittest |
| import weakref |
| |
| import pytest |
| |
| import numpy as np |
| |
| import pyarrow as pa |
| from pyarrow.csv import ( |
| open_csv, read_csv, ReadOptions, ParseOptions, ConvertOptions, ISO8601, |
| write_csv, WriteOptions, CSVWriter, InvalidRow) |
| from pyarrow.tests import util |
| |
| |
| def generate_col_names(): |
| # 'a', 'b'... 'z', then 'aa', 'ab'... |
| letters = string.ascii_lowercase |
| yield from letters |
| for first in letters: |
| for second in letters: |
| yield first + second |
| |
| |
| def make_random_csv(num_cols=2, num_rows=10, linesep='\r\n', write_names=True): |
| arr = np.random.RandomState(42).randint(0, 1000, size=(num_cols, num_rows)) |
| csv = io.StringIO() |
| col_names = list(itertools.islice(generate_col_names(), num_cols)) |
| if write_names: |
| csv.write(",".join(col_names)) |
| csv.write(linesep) |
| for row in arr.T: |
| csv.write(",".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 make_empty_csv(column_names): |
| csv = io.StringIO() |
| csv.write(",".join(column_names)) |
| csv.write("\n") |
| return csv.getvalue().encode() |
| |
| |
| def check_options_class(cls, **attr_values): |
| """ |
| Check setting and getting attributes of an *Options class. |
| """ |
| opts = cls() |
| |
| for name, values in attr_values.items(): |
| assert getattr(opts, name) == values[0], \ |
| "incorrect default value for " + name |
| for v in values: |
| setattr(opts, name, v) |
| assert getattr(opts, name) == v, "failed setting value" |
| |
| with pytest.raises(AttributeError): |
| opts.zzz_non_existent = True |
| |
| # Check constructor named arguments |
| non_defaults = {name: values[1] for name, values in attr_values.items()} |
| opts = cls(**non_defaults) |
| for name, value in non_defaults.items(): |
| assert getattr(opts, name) == value |
| |
| |
| # The various options classes need to be picklable for dataset |
| def check_options_class_pickling(cls, **attr_values): |
| opts = cls(**attr_values) |
| new_opts = pickle.loads(pickle.dumps(opts, |
| protocol=pickle.HIGHEST_PROTOCOL)) |
| for name, value in attr_values.items(): |
| assert getattr(new_opts, name) == value |
| |
| |
| class InvalidRowHandler: |
| def __init__(self, result): |
| self.result = result |
| self.rows = [] |
| |
| def __call__(self, row): |
| self.rows.append(row) |
| return self.result |
| |
| def __eq__(self, other): |
| return (isinstance(other, InvalidRowHandler) and |
| other.result == self.result) |
| |
| def __ne__(self, other): |
| return (not isinstance(other, InvalidRowHandler) or |
| other.result != self.result) |
| |
| |
| def test_read_options(): |
| cls = ReadOptions |
| opts = cls() |
| |
| check_options_class(cls, use_threads=[True, False], |
| skip_rows=[0, 3], |
| column_names=[[], ["ab", "cd"]], |
| autogenerate_column_names=[False, True], |
| encoding=['utf8', 'utf16'], |
| skip_rows_after_names=[0, 27]) |
| |
| check_options_class_pickling(cls, use_threads=True, |
| skip_rows=3, |
| column_names=["ab", "cd"], |
| autogenerate_column_names=False, |
| encoding='utf16', |
| skip_rows_after_names=27) |
| |
| assert opts.block_size > 0 |
| opts.block_size = 12345 |
| assert opts.block_size == 12345 |
| |
| opts = cls(block_size=1234) |
| assert opts.block_size == 1234 |
| |
| opts.validate() |
| |
| match = "ReadOptions: block_size must be at least 1: 0" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.block_size = 0 |
| opts.validate() |
| |
| match = "ReadOptions: skip_rows cannot be negative: -1" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.skip_rows = -1 |
| opts.validate() |
| |
| match = "ReadOptions: skip_rows_after_names cannot be negative: -1" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.skip_rows_after_names = -1 |
| opts.validate() |
| |
| match = "ReadOptions: autogenerate_column_names cannot be true when" \ |
| " column_names are provided" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.autogenerate_column_names = True |
| opts.column_names = ('a', 'b') |
| opts.validate() |
| |
| |
| def test_parse_options(): |
| cls = ParseOptions |
| skip_handler = InvalidRowHandler('skip') |
| |
| check_options_class(cls, delimiter=[',', 'x'], |
| escape_char=[False, 'y'], |
| quote_char=['"', 'z', False], |
| double_quote=[True, False], |
| newlines_in_values=[False, True], |
| ignore_empty_lines=[True, False], |
| invalid_row_handler=[None, skip_handler]) |
| |
| check_options_class_pickling(cls, delimiter='x', |
| escape_char='y', |
| quote_char=False, |
| double_quote=False, |
| newlines_in_values=True, |
| ignore_empty_lines=False, |
| invalid_row_handler=skip_handler) |
| |
| cls().validate() |
| opts = cls() |
| opts.delimiter = "\t" |
| opts.validate() |
| |
| match = "ParseOptions: delimiter cannot be \\\\r or \\\\n" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.delimiter = "\n" |
| opts.validate() |
| |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.delimiter = "\r" |
| opts.validate() |
| |
| match = "ParseOptions: quote_char cannot be \\\\r or \\\\n" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.quote_char = "\n" |
| opts.validate() |
| |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.quote_char = "\r" |
| opts.validate() |
| |
| match = "ParseOptions: escape_char cannot be \\\\r or \\\\n" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.escape_char = "\n" |
| opts.validate() |
| |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.escape_char = "\r" |
| opts.validate() |
| |
| |
| def test_convert_options(): |
| cls = ConvertOptions |
| opts = cls() |
| |
| check_options_class( |
| cls, check_utf8=[True, False], |
| strings_can_be_null=[False, True], |
| quoted_strings_can_be_null=[True, False], |
| decimal_point=['.', ','], |
| include_columns=[[], ['def', 'abc']], |
| include_missing_columns=[False, True], |
| auto_dict_encode=[False, True], |
| timestamp_parsers=[[], [ISO8601, '%y-%m']]) |
| |
| check_options_class_pickling( |
| cls, check_utf8=False, |
| strings_can_be_null=True, |
| quoted_strings_can_be_null=False, |
| decimal_point=',', |
| include_columns=['def', 'abc'], |
| include_missing_columns=False, |
| auto_dict_encode=True, |
| timestamp_parsers=[ISO8601, '%y-%m']) |
| |
| with pytest.raises(ValueError): |
| opts.decimal_point = '..' |
| |
| assert opts.auto_dict_max_cardinality > 0 |
| opts.auto_dict_max_cardinality = 99999 |
| assert opts.auto_dict_max_cardinality == 99999 |
| |
| 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'] |
| |
| assert opts.timestamp_parsers == [] |
| opts.timestamp_parsers = [ISO8601] |
| assert opts.timestamp_parsers == [ISO8601] |
| |
| opts = cls(column_types={'a': pa.null()}, |
| null_values=['N', 'nn'], true_values=['T', 'tt'], |
| false_values=['F', 'ff'], auto_dict_max_cardinality=999, |
| timestamp_parsers=[ISO8601, '%Y-%m-%d']) |
| 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.auto_dict_max_cardinality == 999 |
| assert opts.timestamp_parsers == [ISO8601, '%Y-%m-%d'] |
| |
| |
| def test_write_options(): |
| cls = WriteOptions |
| opts = cls() |
| |
| check_options_class( |
| cls, include_header=[True, False]) |
| |
| assert opts.batch_size > 0 |
| opts.batch_size = 12345 |
| assert opts.batch_size == 12345 |
| |
| opts = cls(batch_size=9876) |
| assert opts.batch_size == 9876 |
| |
| opts.validate() |
| |
| match = "WriteOptions: batch_size must be at least 1: 0" |
| with pytest.raises(pa.ArrowInvalid, match=match): |
| opts = cls() |
| opts.batch_size = 0 |
| opts.validate() |
| |
| |
| class BaseTestCSV(abc.ABC): |
| """Common tests which are shared by streaming and non streaming readers""" |
| |
| @abc.abstractmethod |
| def read_bytes(self, b, **kwargs): |
| """ |
| :param b: bytes to be parsed |
| :param kwargs: arguments passed on to open the csv file |
| :return: b parsed as a single RecordBatch |
| """ |
| raise NotImplementedError |
| |
| @property |
| @abc.abstractmethod |
| def use_threads(self): |
| """Whether this test is multi-threaded""" |
| raise NotImplementedError |
| |
| @staticmethod |
| def check_names(table, names): |
| assert table.num_columns == len(names) |
| assert table.column_names == names |
| |
| def test_header_skip_rows(self): |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| opts = ReadOptions() |
| opts.skip_rows = 1 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["ef", "gh"]) |
| assert table.to_pydict() == { |
| "ef": ["ij", "mn"], |
| "gh": ["kl", "op"], |
| } |
| |
| opts.skip_rows = 3 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["mn", "op"]) |
| assert table.to_pydict() == { |
| "mn": [], |
| "op": [], |
| } |
| |
| opts.skip_rows = 4 |
| with pytest.raises(pa.ArrowInvalid): |
| # Not enough rows |
| table = self.read_bytes(rows, read_options=opts) |
| |
| # Can skip rows with a different number of columns |
| rows = b"abcd\n,,,,,\nij,kl\nmn,op\n" |
| opts.skip_rows = 2 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["ij", "kl"]) |
| assert table.to_pydict() == { |
| "ij": ["mn"], |
| "kl": ["op"], |
| } |
| |
| # Can skip all rows exactly when columns are given |
| opts.skip_rows = 4 |
| opts.column_names = ['ij', 'kl'] |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["ij", "kl"]) |
| assert table.to_pydict() == { |
| "ij": [], |
| "kl": [], |
| } |
| |
| def test_skip_rows_after_names(self): |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| opts = ReadOptions() |
| opts.skip_rows_after_names = 1 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["ab", "cd"]) |
| assert table.to_pydict() == { |
| "ab": ["ij", "mn"], |
| "cd": ["kl", "op"], |
| } |
| |
| # Can skip exact number of rows |
| opts.skip_rows_after_names = 3 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["ab", "cd"]) |
| assert table.to_pydict() == { |
| "ab": [], |
| "cd": [], |
| } |
| |
| # Can skip beyond all rows |
| opts.skip_rows_after_names = 4 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["ab", "cd"]) |
| assert table.to_pydict() == { |
| "ab": [], |
| "cd": [], |
| } |
| |
| # Can skip rows with a different number of columns |
| rows = b"abcd\n,,,,,\nij,kl\nmn,op\n" |
| opts.skip_rows_after_names = 2 |
| opts.column_names = ["f0", "f1"] |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["f0", "f1"]) |
| assert table.to_pydict() == { |
| "f0": ["ij", "mn"], |
| "f1": ["kl", "op"], |
| } |
| opts = ReadOptions() |
| |
| # Can skip rows with new lines in the value |
| rows = b'ab,cd\n"e\nf","g\n\nh"\n"ij","k\nl"\nmn,op' |
| opts.skip_rows_after_names = 2 |
| parse_opts = ParseOptions() |
| parse_opts.newlines_in_values = True |
| table = self.read_bytes(rows, read_options=opts, |
| parse_options=parse_opts) |
| self.check_names(table, ["ab", "cd"]) |
| assert table.to_pydict() == { |
| "ab": ["mn"], |
| "cd": ["op"], |
| } |
| |
| # Can skip rows when block ends in middle of quoted value |
| opts.skip_rows_after_names = 2 |
| opts.block_size = 26 |
| table = self.read_bytes(rows, read_options=opts, |
| parse_options=parse_opts) |
| self.check_names(table, ["ab", "cd"]) |
| assert table.to_pydict() == { |
| "ab": ["mn"], |
| "cd": ["op"], |
| } |
| opts = ReadOptions() |
| |
| # Can skip rows that are beyond the first block without lexer |
| rows, expected = make_random_csv(num_cols=5, num_rows=1000) |
| opts.skip_rows_after_names = 900 |
| opts.block_size = len(rows) / 11 |
| table = self.read_bytes(rows, read_options=opts) |
| assert table.schema == expected.schema |
| assert table.num_rows == 100 |
| table_dict = table.to_pydict() |
| for name, values in expected.to_pydict().items(): |
| assert values[900:] == table_dict[name] |
| |
| # Can skip rows that are beyond the first block with lexer |
| table = self.read_bytes(rows, read_options=opts, |
| parse_options=parse_opts) |
| assert table.schema == expected.schema |
| assert table.num_rows == 100 |
| table_dict = table.to_pydict() |
| for name, values in expected.to_pydict().items(): |
| assert values[900:] == table_dict[name] |
| |
| # Skip rows and skip rows after names |
| rows, expected = make_random_csv(num_cols=5, num_rows=200, |
| write_names=False) |
| opts = ReadOptions() |
| opts.skip_rows = 37 |
| opts.skip_rows_after_names = 41 |
| opts.column_names = expected.schema.names |
| table = self.read_bytes(rows, read_options=opts, |
| parse_options=parse_opts) |
| assert table.schema == expected.schema |
| assert (table.num_rows == |
| expected.num_rows - opts.skip_rows - |
| opts.skip_rows_after_names) |
| table_dict = table.to_pydict() |
| for name, values in expected.to_pydict().items(): |
| assert (values[opts.skip_rows + opts.skip_rows_after_names:] == |
| table_dict[name]) |
| |
| def test_row_number_offset_in_errors(self): |
| # Row numbers are only correctly counted in serial reads |
| def format_msg(msg_format, row, *args): |
| if self.use_threads: |
| row_info = "" |
| else: |
| row_info = "Row #{}: ".format(row) |
| return msg_format.format(row_info, *args) |
| |
| csv, _ = make_random_csv(4, 100, write_names=True) |
| |
| read_options = ReadOptions() |
| read_options.block_size = len(csv) / 3 |
| convert_options = ConvertOptions() |
| convert_options.column_types = {"a": pa.int32()} |
| |
| # Test without skip_rows and column names in the csv |
| csv_bad_columns = csv + b"1,2\r\n" |
| message_columns = format_msg("{}Expected 4 columns, got 2", 102) |
| with pytest.raises(pa.ArrowInvalid, match=message_columns): |
| self.read_bytes(csv_bad_columns, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| csv_bad_type = csv + b"a,b,c,d\r\n" |
| message_value = format_msg( |
| "In CSV column #0: {}" |
| "CSV conversion error to int32: invalid value 'a'", |
| 102, csv) |
| with pytest.raises(pa.ArrowInvalid, match=message_value): |
| self.read_bytes(csv_bad_type, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| long_row = (b"this is a long row" * 15) + b",3\r\n" |
| csv_bad_columns_long = csv + long_row |
| message_long = format_msg("{}Expected 4 columns, got 2: {} ...", 102, |
| long_row[0:96].decode("utf-8")) |
| with pytest.raises(pa.ArrowInvalid, match=message_long): |
| self.read_bytes(csv_bad_columns_long, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| # Test skipping rows after the names |
| read_options.skip_rows_after_names = 47 |
| |
| with pytest.raises(pa.ArrowInvalid, match=message_columns): |
| self.read_bytes(csv_bad_columns, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| with pytest.raises(pa.ArrowInvalid, match=message_value): |
| self.read_bytes(csv_bad_type, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| with pytest.raises(pa.ArrowInvalid, match=message_long): |
| self.read_bytes(csv_bad_columns_long, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| read_options.skip_rows_after_names = 0 |
| |
| # Test without skip_rows and column names not in the csv |
| csv, _ = make_random_csv(4, 100, write_names=False) |
| read_options.column_names = ["a", "b", "c", "d"] |
| csv_bad_columns = csv + b"1,2\r\n" |
| message_columns = format_msg("{}Expected 4 columns, got 2", 101) |
| with pytest.raises(pa.ArrowInvalid, match=message_columns): |
| self.read_bytes(csv_bad_columns, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| csv_bad_columns_long = csv + long_row |
| message_long = format_msg("{}Expected 4 columns, got 2: {} ...", 101, |
| long_row[0:96].decode("utf-8")) |
| with pytest.raises(pa.ArrowInvalid, match=message_long): |
| self.read_bytes(csv_bad_columns_long, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| csv_bad_type = csv + b"a,b,c,d\r\n" |
| message_value = format_msg( |
| "In CSV column #0: {}" |
| "CSV conversion error to int32: invalid value 'a'", |
| 101) |
| message_value = message_value.format(len(csv)) |
| with pytest.raises(pa.ArrowInvalid, match=message_value): |
| self.read_bytes(csv_bad_type, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| # Test with skip_rows and column names not in the csv |
| read_options.skip_rows = 23 |
| with pytest.raises(pa.ArrowInvalid, match=message_columns): |
| self.read_bytes(csv_bad_columns, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| with pytest.raises(pa.ArrowInvalid, match=message_value): |
| self.read_bytes(csv_bad_type, |
| read_options=read_options, |
| convert_options=convert_options) |
| |
| def test_invalid_row_handler(self): |
| rows = b"a,b\nc\nd,e\nf,g,h\ni,j\n" |
| parse_opts = ParseOptions() |
| with pytest.raises( |
| ValueError, |
| match="Expected 2 columns, got 1: c"): |
| self.read_bytes(rows, parse_options=parse_opts) |
| |
| # Skip requested |
| parse_opts.invalid_row_handler = InvalidRowHandler('skip') |
| table = self.read_bytes(rows, parse_options=parse_opts) |
| assert table.to_pydict() == { |
| 'a': ["d", "i"], |
| 'b': ["e", "j"], |
| } |
| |
| def row_num(x): |
| return None if self.use_threads else x |
| expected_rows = [ |
| InvalidRow(2, 1, row_num(2), "c"), |
| InvalidRow(2, 3, row_num(4), "f,g,h"), |
| ] |
| assert parse_opts.invalid_row_handler.rows == expected_rows |
| |
| # Error requested |
| parse_opts.invalid_row_handler = InvalidRowHandler('error') |
| with pytest.raises( |
| ValueError, |
| match="Expected 2 columns, got 1: c"): |
| self.read_bytes(rows, parse_options=parse_opts) |
| expected_rows = [InvalidRow(2, 1, row_num(2), "c")] |
| assert parse_opts.invalid_row_handler.rows == expected_rows |
| |
| |
| class BaseCSVTableRead(BaseTestCSV): |
| |
| def read_csv(self, csv, *args, validate_full=True, **kwargs): |
| """ |
| Reads the CSV file into memory using pyarrow's read_csv |
| csv The CSV bytes |
| args Positional arguments to be forwarded to pyarrow's read_csv |
| validate_full Whether or not to fully validate the resulting table |
| kwargs Keyword arguments to be forwarded to pyarrow's read_csv |
| """ |
| assert isinstance(self.use_threads, bool) # sanity check |
| read_options = kwargs.setdefault('read_options', ReadOptions()) |
| read_options.use_threads = self.use_threads |
| table = read_csv(csv, *args, **kwargs) |
| table.validate(full=validate_full) |
| return table |
| |
| def read_bytes(self, b, **kwargs): |
| return self.read_csv(pa.py_buffer(b), **kwargs) |
| |
| 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_one_chunk(self): |
| # ARROW-7661: lack of newline at end of file should not produce |
| # an additional chunk. |
| rows = [b"a,b", b"1,2", b"3,4", b"56,78"] |
| for line_ending in [b'\n', b'\r', b'\r\n']: |
| for file_ending in [b'', line_ending]: |
| data = line_ending.join(rows) + file_ending |
| table = self.read_bytes(data) |
| assert len(table.to_batches()) == 1 |
| assert table.to_pydict() == { |
| "a": [1, 3, 56], |
| "b": [2, 4, 78], |
| } |
| |
| def test_header_column_names(self): |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| opts = ReadOptions() |
| opts.column_names = ["x", "y"] |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["x", "y"]) |
| assert table.to_pydict() == { |
| "x": ["ab", "ef", "ij", "mn"], |
| "y": ["cd", "gh", "kl", "op"], |
| } |
| |
| opts.skip_rows = 3 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["x", "y"]) |
| assert table.to_pydict() == { |
| "x": ["mn"], |
| "y": ["op"], |
| } |
| |
| opts.skip_rows = 4 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["x", "y"]) |
| assert table.to_pydict() == { |
| "x": [], |
| "y": [], |
| } |
| |
| opts.skip_rows = 5 |
| with pytest.raises(pa.ArrowInvalid): |
| # Not enough rows |
| table = self.read_bytes(rows, read_options=opts) |
| |
| # Unexpected number of columns |
| opts.skip_rows = 0 |
| opts.column_names = ["x", "y", "z"] |
| with pytest.raises(pa.ArrowInvalid, |
| match="Expected 3 columns, got 2"): |
| table = self.read_bytes(rows, read_options=opts) |
| |
| # Can skip rows with a different number of columns |
| rows = b"abcd\n,,,,,\nij,kl\nmn,op\n" |
| opts.skip_rows = 2 |
| opts.column_names = ["x", "y"] |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["x", "y"]) |
| assert table.to_pydict() == { |
| "x": ["ij", "mn"], |
| "y": ["kl", "op"], |
| } |
| |
| def test_header_autogenerate_column_names(self): |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| opts = ReadOptions() |
| opts.autogenerate_column_names = True |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["f0", "f1"]) |
| assert table.to_pydict() == { |
| "f0": ["ab", "ef", "ij", "mn"], |
| "f1": ["cd", "gh", "kl", "op"], |
| } |
| |
| opts.skip_rows = 3 |
| table = self.read_bytes(rows, read_options=opts) |
| self.check_names(table, ["f0", "f1"]) |
| assert table.to_pydict() == { |
| "f0": ["mn"], |
| "f1": ["op"], |
| } |
| |
| # Not enough rows, impossible to infer number of columns |
| opts.skip_rows = 4 |
| with pytest.raises(pa.ArrowInvalid): |
| table = self.read_bytes(rows, read_options=opts) |
| |
| def test_include_columns(self): |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| convert_options = ConvertOptions() |
| convert_options.include_columns = ['ab'] |
| table = self.read_bytes(rows, convert_options=convert_options) |
| self.check_names(table, ["ab"]) |
| assert table.to_pydict() == { |
| "ab": ["ef", "ij", "mn"], |
| } |
| |
| # Order of include_columns is respected, regardless of CSV order |
| convert_options.include_columns = ['cd', 'ab'] |
| table = self.read_bytes(rows, convert_options=convert_options) |
| schema = pa.schema([('cd', pa.string()), |
| ('ab', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| "cd": ["gh", "kl", "op"], |
| "ab": ["ef", "ij", "mn"], |
| } |
| |
| # Include a column not in the CSV file => raises by default |
| convert_options.include_columns = ['xx', 'ab', 'yy'] |
| with pytest.raises(KeyError, |
| match="Column 'xx' in include_columns " |
| "does not exist in CSV file"): |
| self.read_bytes(rows, convert_options=convert_options) |
| |
| def test_include_missing_columns(self): |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| read_options = ReadOptions() |
| convert_options = ConvertOptions() |
| convert_options.include_columns = ['xx', 'ab', 'yy'] |
| convert_options.include_missing_columns = True |
| table = self.read_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| schema = pa.schema([('xx', pa.null()), |
| ('ab', pa.string()), |
| ('yy', pa.null())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| "xx": [None, None, None], |
| "ab": ["ef", "ij", "mn"], |
| "yy": [None, None, None], |
| } |
| |
| # Combining with `column_names` |
| read_options.column_names = ["xx", "yy"] |
| convert_options.include_columns = ["yy", "cd"] |
| table = self.read_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| schema = pa.schema([('yy', pa.string()), |
| ('cd', pa.null())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| "yy": ["cd", "gh", "kl", "op"], |
| "cd": [None, None, None, None], |
| } |
| |
| # And with `column_types` as well |
| convert_options.column_types = {"yy": pa.binary(), |
| "cd": pa.int32()} |
| table = self.read_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| schema = pa.schema([('yy', pa.binary()), |
| ('cd', pa.int32())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| "yy": [b"cd", b"gh", b"kl", b"op"], |
| "cd": [None, None, None, None], |
| } |
| |
| 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': ["3", "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': ["", "foo", "nan"], |
| 'd': [None, None, None], |
| 'e': [b"3", b"nan", b"\xff"], |
| 'f': [None, True, False], |
| } |
| |
| def test_decimal_point(self): |
| # Infer floats with a custom decimal point |
| parse_options = ParseOptions(delimiter=';') |
| rows = b"a;b\n1.25;2,5\nNA;-3\n-4;NA" |
| |
| table = self.read_bytes(rows, parse_options=parse_options) |
| schema = pa.schema([('a', pa.float64()), |
| ('b', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [1.25, None, -4.0], |
| 'b': ["2,5", "-3", "NA"], |
| } |
| |
| convert_options = ConvertOptions(decimal_point=',') |
| table = self.read_bytes(rows, parse_options=parse_options, |
| convert_options=convert_options) |
| schema = pa.schema([('a', pa.string()), |
| ('b', pa.float64())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': ["1.25", "NA", "-4"], |
| 'b': [2.5, -3.0, None], |
| } |
| |
| def test_simple_timestamps(self): |
| # Infer a timestamp column |
| rows = (b"a,b,c\n" |
| b"1970,1970-01-01 00:00:00,1970-01-01 00:00:00.123\n" |
| b"1989,1989-07-14 01:00:00,1989-07-14 01:00:00.123456\n") |
| table = self.read_bytes(rows) |
| schema = pa.schema([('a', pa.int64()), |
| ('b', pa.timestamp('s')), |
| ('c', pa.timestamp('ns'))]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [1970, 1989], |
| 'b': [datetime(1970, 1, 1), datetime(1989, 7, 14, 1)], |
| 'c': [datetime(1970, 1, 1, 0, 0, 0, 123000), |
| datetime(1989, 7, 14, 1, 0, 0, 123456)], |
| } |
| |
| def test_timestamp_parsers(self): |
| # Infer timestamps with custom parsers |
| rows = b"a,b\n1970/01/01,1980-01-01 00\n1970/01/02,1980-01-02 00\n" |
| opts = ConvertOptions() |
| |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.string()), |
| ('b', pa.timestamp('s'))]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': ['1970/01/01', '1970/01/02'], |
| 'b': [datetime(1980, 1, 1), datetime(1980, 1, 2)], |
| } |
| |
| opts.timestamp_parsers = ['%Y/%m/%d'] |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.timestamp('s')), |
| ('b', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [datetime(1970, 1, 1), datetime(1970, 1, 2)], |
| 'b': ['1980-01-01 00', '1980-01-02 00'], |
| } |
| |
| opts.timestamp_parsers = ['%Y/%m/%d', ISO8601] |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.timestamp('s')), |
| ('b', pa.timestamp('s'))]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [datetime(1970, 1, 1), datetime(1970, 1, 2)], |
| 'b': [datetime(1980, 1, 1), datetime(1980, 1, 2)], |
| } |
| |
| def test_dates(self): |
| # Dates are inferred as date32 by default |
| rows = b"a,b\n1970-01-01,1970-01-02\n1971-01-01,1971-01-02\n" |
| table = self.read_bytes(rows) |
| schema = pa.schema([('a', pa.date32()), |
| ('b', pa.date32())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [date(1970, 1, 1), date(1971, 1, 1)], |
| 'b': [date(1970, 1, 2), date(1971, 1, 2)], |
| } |
| |
| # Can ask for date types explicitly |
| opts = ConvertOptions() |
| opts.column_types = {'a': pa.date32(), 'b': pa.date64()} |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.date32()), |
| ('b', pa.date64())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [date(1970, 1, 1), date(1971, 1, 1)], |
| 'b': [date(1970, 1, 2), date(1971, 1, 2)], |
| } |
| |
| # Can ask for timestamp types explicitly |
| opts = ConvertOptions() |
| opts.column_types = {'a': pa.timestamp('s'), 'b': pa.timestamp('ms')} |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.timestamp('s')), |
| ('b', pa.timestamp('ms'))]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [datetime(1970, 1, 1), datetime(1971, 1, 1)], |
| 'b': [datetime(1970, 1, 2), datetime(1971, 1, 2)], |
| } |
| |
| def test_times(self): |
| # Times are inferred as time32[s] by default |
| from datetime import time |
| |
| rows = b"a,b\n12:34:56,12:34:56.789\n23:59:59,23:59:59.999\n" |
| table = self.read_bytes(rows) |
| # Column 'b' has subseconds, so cannot be inferred as time32[s] |
| schema = pa.schema([('a', pa.time32('s')), |
| ('b', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [time(12, 34, 56), time(23, 59, 59)], |
| 'b': ["12:34:56.789", "23:59:59.999"], |
| } |
| |
| # Can ask for time types explicitly |
| opts = ConvertOptions() |
| opts.column_types = {'a': pa.time64('us'), 'b': pa.time32('ms')} |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.time64('us')), |
| ('b', pa.time32('ms'))]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'a': [time(12, 34, 56), time(23, 59, 59)], |
| 'b': [time(12, 34, 56, 789000), time(23, 59, 59, 999000)], |
| } |
| |
| def test_auto_dict_encode(self): |
| opts = ConvertOptions(auto_dict_encode=True) |
| rows = "a,b\nab,1\ncdé,2\ncdé,3\nab,4".encode() |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.dictionary(pa.int32(), pa.string())), |
| ('b', pa.int64())]) |
| expected = { |
| 'a': ["ab", "cdé", "cdé", "ab"], |
| 'b': [1, 2, 3, 4], |
| } |
| assert table.schema == schema |
| assert table.to_pydict() == expected |
| |
| opts.auto_dict_max_cardinality = 2 |
| table = self.read_bytes(rows, convert_options=opts) |
| assert table.schema == schema |
| assert table.to_pydict() == expected |
| |
| # Cardinality above max => plain-encoded |
| opts.auto_dict_max_cardinality = 1 |
| table = self.read_bytes(rows, convert_options=opts) |
| assert table.schema == pa.schema([('a', pa.string()), |
| ('b', pa.int64())]) |
| assert table.to_pydict() == expected |
| |
| # With invalid UTF8, not checked |
| opts.auto_dict_max_cardinality = 50 |
| opts.check_utf8 = False |
| rows = b"a,b\nab,1\ncd\xff,2\nab,3" |
| table = self.read_bytes(rows, convert_options=opts, |
| validate_full=False) |
| assert table.schema == schema |
| dict_values = table['a'].chunk(0).dictionary |
| assert len(dict_values) == 2 |
| assert dict_values[0].as_py() == "ab" |
| assert dict_values[1].as_buffer() == b"cd\xff" |
| |
| # With invalid UTF8, checked |
| opts.check_utf8 = True |
| table = self.read_bytes(rows, convert_options=opts) |
| schema = pa.schema([('a', pa.dictionary(pa.int32(), pa.binary())), |
| ('b', pa.int64())]) |
| expected = { |
| 'a': [b"ab", b"cd\xff", b"ab"], |
| 'b': [1, 2, 3], |
| } |
| assert table.schema == schema |
| assert table.to_pydict() == expected |
| |
| 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': ["Xxx", "#N/A"], |
| 'c': ["1", ""], |
| '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, "#N/A"], |
| 'c': ["1", ""], |
| 'd': [2, None], |
| } |
| opts.quoted_strings_can_be_null = False |
| table = self.read_bytes(rows, convert_options=opts) |
| assert table.to_pydict() == { |
| 'a': [None, None], |
| 'b': ["Xxx", "#N/A"], |
| 'c': ["1", ""], |
| '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': ["#N/A"], |
| 'b': [""], |
| } |
| |
| 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 CSV column #1: " in err |
| assert "CSV conversion error to float: invalid value 'XXX'" in err |
| |
| def test_column_types_dict(self): |
| # Ask for dict-encoded column types in ConvertOptions |
| column_types = [ |
| ('a', pa.dictionary(pa.int32(), pa.utf8())), |
| ('b', pa.dictionary(pa.int32(), pa.int64())), |
| ('c', pa.dictionary(pa.int32(), pa.decimal128(11, 2))), |
| ('d', pa.dictionary(pa.int32(), pa.large_utf8()))] |
| |
| opts = ConvertOptions(column_types=dict(column_types)) |
| rows = (b"a,b,c,d\n" |
| b"abc,123456,1.0,zz\n" |
| b"defg,123456,0.5,xx\n" |
| b"abc,N/A,1.0,xx\n") |
| table = self.read_bytes(rows, convert_options=opts) |
| |
| schema = pa.schema(column_types) |
| expected = { |
| 'a': ["abc", "defg", "abc"], |
| 'b': [123456, 123456, None], |
| 'c': [Decimal("1.00"), Decimal("0.50"), Decimal("1.00")], |
| 'd': ["zz", "xx", "xx"], |
| } |
| assert table.schema == schema |
| assert table.to_pydict() == expected |
| |
| # Unsupported index type |
| column_types[0] = ('a', pa.dictionary(pa.int8(), pa.utf8())) |
| |
| opts = ConvertOptions(column_types=dict(column_types)) |
| with pytest.raises(NotImplementedError): |
| table = self.read_bytes(rows, convert_options=opts) |
| |
| def test_column_types_with_column_names(self): |
| # When both `column_names` and `column_types` are given, names |
| # in `column_types` should refer to names in `column_names` |
| rows = b"a,b\nc,d\ne,f\n" |
| read_options = ReadOptions(column_names=['x', 'y']) |
| convert_options = ConvertOptions(column_types={'x': pa.binary()}) |
| table = self.read_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| schema = pa.schema([('x', pa.binary()), |
| ('y', pa.string())]) |
| assert table.schema == schema |
| assert table.to_pydict() == { |
| 'x': [b'a', b'c', b'e'], |
| 'y': ['b', 'd', 'f'], |
| } |
| |
| 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_empty_lines(self): |
| rows = b"a,b\n\r1,2\r\n\r\n3,4\r\n" |
| table = self.read_bytes(rows) |
| assert table.to_pydict() == { |
| 'a': [1, 3], |
| 'b': [2, 4], |
| } |
| parse_options = ParseOptions(ignore_empty_lines=False) |
| table = self.read_bytes(rows, parse_options=parse_options) |
| assert table.to_pydict() == { |
| 'a': [None, 1, None, 3], |
| 'b': [None, 2, None, 4], |
| } |
| read_options = ReadOptions(skip_rows=2) |
| table = self.read_bytes(rows, parse_options=parse_options, |
| read_options=read_options) |
| assert table.to_pydict() == { |
| '1': [None, 3], |
| '2': [None, 4], |
| } |
| |
| 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': ['de'], |
| 'c': ['fg;eh'], |
| } |
| opts = ParseOptions(delimiter=';') |
| table = self.read_bytes(rows, parse_options=opts) |
| assert table.to_pydict() == { |
| 'a': ['de,fg'], |
| 'b,c': ['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() |
| |
| def test_stress_convert_options_blowup(self): |
| # ARROW-6481: A convert_options with a very large number of columns |
| # should not blow memory and CPU time. |
| try: |
| clock = time.thread_time |
| except AttributeError: |
| clock = time.time |
| num_columns = 10000 |
| col_names = ["K{}".format(i) for i in range(num_columns)] |
| csv = make_empty_csv(col_names) |
| t1 = clock() |
| convert_options = ConvertOptions( |
| column_types={k: pa.string() for k in col_names[::2]}) |
| table = self.read_bytes(csv, convert_options=convert_options) |
| dt = clock() - t1 |
| # Check that processing time didn't blow up. |
| # This is a conservative check (it takes less than 300 ms |
| # in debug mode on my local machine). |
| assert dt <= 10.0 |
| # Check result |
| assert table.num_columns == num_columns |
| assert table.num_rows == 0 |
| assert table.column_names == col_names |
| |
| def test_cancellation(self): |
| if (threading.current_thread().ident != |
| threading.main_thread().ident): |
| pytest.skip("test only works from main Python thread") |
| # Skips test if not available |
| raise_signal = util.get_raise_signal() |
| signum = signal.SIGINT |
| |
| def signal_from_thread(): |
| # Give our workload a chance to start up |
| time.sleep(0.2) |
| raise_signal(signum) |
| |
| # We start with a small CSV reading workload and increase its size |
| # until it's large enough to get an interruption during it, even in |
| # release mode on fast machines. |
| last_duration = 0.0 |
| workload_size = 100_000 |
| attempts = 0 |
| |
| while last_duration < 5.0 and attempts < 10: |
| print("workload size:", workload_size) |
| large_csv = b"a,b,c\n" + b"1,2,3\n" * workload_size |
| exc_info = None |
| |
| try: |
| # We use a signal fd to reliably ensure that the signal |
| # has been delivered to Python, regardless of how exactly |
| # it was caught. |
| with util.signal_wakeup_fd() as sigfd: |
| try: |
| t = threading.Thread(target=signal_from_thread) |
| t.start() |
| t1 = time.time() |
| try: |
| self.read_bytes(large_csv) |
| except KeyboardInterrupt as e: |
| exc_info = e |
| last_duration = time.time() - t1 |
| finally: |
| # Wait for signal to arrive if it didn't already, |
| # to avoid getting a KeyboardInterrupt after the |
| # `except` block below. |
| select.select([sigfd], [], [sigfd], 10.0) |
| |
| except KeyboardInterrupt: |
| # KeyboardInterrupt didn't interrupt `read_bytes` above. |
| pass |
| |
| if exc_info is not None: |
| # We managed to get `self.read_bytes` interrupted, see if it |
| # was actually interrupted inside Arrow C++ or in the Python |
| # scaffolding. |
| if exc_info.__context__ is not None: |
| # Interrupted inside Arrow C++, we're satisfied now |
| break |
| |
| # Increase workload size to get a better chance |
| workload_size = workload_size * 3 |
| |
| if exc_info is None: |
| pytest.fail("Failed to get an interruption during CSV reading") |
| |
| # Interruption should have arrived timely |
| assert last_duration <= 1.0 |
| e = exc_info.__context__ |
| assert isinstance(e, pa.ArrowCancelled) |
| assert e.signum == signum |
| |
| def test_cancellation_disabled(self): |
| # ARROW-12622: reader would segfault when the cancelling signal |
| # handler was not enabled (e.g. if disabled, or if not on the |
| # main thread) |
| t = threading.Thread( |
| target=lambda: self.read_bytes(b"f64\n0.1")) |
| t.start() |
| t.join() |
| |
| |
| class TestSerialCSVTableRead(BaseCSVTableRead): |
| @property |
| def use_threads(self): |
| return False |
| |
| |
| class TestThreadedCSVTableRead(BaseCSVTableRead): |
| @property |
| def use_threads(self): |
| return True |
| |
| |
| class BaseStreamingCSVRead(BaseTestCSV): |
| |
| def open_csv(self, csv, *args, **kwargs): |
| """ |
| Reads the CSV file into memory using pyarrow's open_csv |
| csv The CSV bytes |
| args Positional arguments to be forwarded to pyarrow's open_csv |
| kwargs Keyword arguments to be forwarded to pyarrow's open_csv |
| """ |
| read_options = kwargs.setdefault('read_options', ReadOptions()) |
| read_options.use_threads = self.use_threads |
| return open_csv(csv, *args, **kwargs) |
| |
| def open_bytes(self, b, **kwargs): |
| return self.open_csv(pa.py_buffer(b), **kwargs) |
| |
| def check_reader(self, reader, expected_schema, expected_data): |
| assert reader.schema == expected_schema |
| batches = list(reader) |
| assert len(batches) == len(expected_data) |
| for batch, expected_batch in zip(batches, expected_data): |
| batch.validate(full=True) |
| assert batch.schema == expected_schema |
| assert batch.to_pydict() == expected_batch |
| |
| def read_bytes(self, b, **kwargs): |
| return self.open_bytes(b, **kwargs).read_all() |
| |
| def test_file_object(self): |
| data = b"a,b\n1,2\n3,4\n" |
| expected_data = {'a': [1, 3], 'b': [2, 4]} |
| bio = io.BytesIO(data) |
| reader = self.open_csv(bio) |
| expected_schema = pa.schema([('a', pa.int64()), |
| ('b', pa.int64())]) |
| self.check_reader(reader, expected_schema, [expected_data]) |
| |
| def test_header(self): |
| rows = b"abc,def,gh\n" |
| reader = self.open_bytes(rows) |
| expected_schema = pa.schema([('abc', pa.null()), |
| ('def', pa.null()), |
| ('gh', pa.null())]) |
| self.check_reader(reader, expected_schema, []) |
| |
| def test_inference(self): |
| # Inference is done on first block |
| rows = b"a,b\n123,456\nabc,de\xff\ngh,ij\n" |
| expected_schema = pa.schema([('a', pa.string()), |
| ('b', pa.binary())]) |
| |
| read_options = ReadOptions() |
| read_options.block_size = len(rows) |
| reader = self.open_bytes(rows, read_options=read_options) |
| self.check_reader(reader, expected_schema, |
| [{'a': ['123', 'abc', 'gh'], |
| 'b': [b'456', b'de\xff', b'ij']}]) |
| |
| read_options.block_size = len(rows) - 1 |
| reader = self.open_bytes(rows, read_options=read_options) |
| self.check_reader(reader, expected_schema, |
| [{'a': ['123', 'abc'], |
| 'b': [b'456', b'de\xff']}, |
| {'a': ['gh'], |
| 'b': [b'ij']}]) |
| |
| def test_inference_failure(self): |
| # Inference on first block, then conversion failure on second block |
| rows = b"a,b\n123,456\nabc,de\xff\ngh,ij\n" |
| read_options = ReadOptions() |
| read_options.block_size = len(rows) - 7 |
| reader = self.open_bytes(rows, read_options=read_options) |
| expected_schema = pa.schema([('a', pa.int64()), |
| ('b', pa.int64())]) |
| assert reader.schema == expected_schema |
| assert reader.read_next_batch().to_pydict() == { |
| 'a': [123], 'b': [456] |
| } |
| # Second block |
| with pytest.raises(ValueError, |
| match="CSV conversion error to int64"): |
| reader.read_next_batch() |
| # EOF |
| with pytest.raises(StopIteration): |
| reader.read_next_batch() |
| |
| def test_invalid_csv(self): |
| # CSV errors on first block |
| rows = b"a,b\n1,2,3\n4,5\n6,7\n" |
| read_options = ReadOptions() |
| read_options.block_size = 10 |
| with pytest.raises(pa.ArrowInvalid, |
| match="Expected 2 columns, got 3"): |
| reader = self.open_bytes( |
| rows, read_options=read_options) |
| |
| # CSV errors on second block |
| rows = b"a,b\n1,2\n3,4,5\n6,7\n" |
| read_options.block_size = 8 |
| reader = self.open_bytes(rows, read_options=read_options) |
| assert reader.read_next_batch().to_pydict() == {'a': [1], 'b': [2]} |
| with pytest.raises(pa.ArrowInvalid, |
| match="Expected 2 columns, got 3"): |
| reader.read_next_batch() |
| # Cannot continue after a parse error |
| with pytest.raises(StopIteration): |
| reader.read_next_batch() |
| |
| def test_options_delimiter(self): |
| rows = b"a;b,c\nde,fg;eh\n" |
| reader = self.open_bytes(rows) |
| expected_schema = pa.schema([('a;b', pa.string()), |
| ('c', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'a;b': ['de'], |
| 'c': ['fg;eh']}]) |
| |
| opts = ParseOptions(delimiter=';') |
| reader = self.open_bytes(rows, parse_options=opts) |
| expected_schema = pa.schema([('a', pa.string()), |
| ('b,c', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'a': ['de,fg'], |
| 'b,c': ['eh']}]) |
| |
| def test_no_ending_newline(self): |
| # No \n after last line |
| rows = b"a,b,c\n1,2,3\n4,5,6" |
| reader = self.open_bytes(rows) |
| expected_schema = pa.schema([('a', pa.int64()), |
| ('b', pa.int64()), |
| ('c', pa.int64())]) |
| self.check_reader(reader, expected_schema, |
| [{'a': [1, 4], |
| 'b': [2, 5], |
| 'c': [3, 6]}]) |
| |
| def test_empty_file(self): |
| with pytest.raises(ValueError, match="Empty CSV file"): |
| self.open_bytes(b"") |
| |
| def test_column_options(self): |
| # With column_names |
| rows = b"1,2,3\n4,5,6" |
| read_options = ReadOptions() |
| read_options.column_names = ['d', 'e', 'f'] |
| reader = self.open_bytes(rows, read_options=read_options) |
| expected_schema = pa.schema([('d', pa.int64()), |
| ('e', pa.int64()), |
| ('f', pa.int64())]) |
| self.check_reader(reader, expected_schema, |
| [{'d': [1, 4], |
| 'e': [2, 5], |
| 'f': [3, 6]}]) |
| |
| # With include_columns |
| convert_options = ConvertOptions() |
| convert_options.include_columns = ['f', 'e'] |
| reader = self.open_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| expected_schema = pa.schema([('f', pa.int64()), |
| ('e', pa.int64())]) |
| self.check_reader(reader, expected_schema, |
| [{'e': [2, 5], |
| 'f': [3, 6]}]) |
| |
| # With column_types |
| convert_options.column_types = {'e': pa.string()} |
| reader = self.open_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| expected_schema = pa.schema([('f', pa.int64()), |
| ('e', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'e': ["2", "5"], |
| 'f': [3, 6]}]) |
| |
| # Missing columns in include_columns |
| convert_options.include_columns = ['g', 'f', 'e'] |
| with pytest.raises( |
| KeyError, |
| match="Column 'g' in include_columns does not exist"): |
| reader = self.open_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| |
| convert_options.include_missing_columns = True |
| reader = self.open_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| expected_schema = pa.schema([('g', pa.null()), |
| ('f', pa.int64()), |
| ('e', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'g': [None, None], |
| 'e': ["2", "5"], |
| 'f': [3, 6]}]) |
| |
| convert_options.column_types = {'e': pa.string(), 'g': pa.float64()} |
| reader = self.open_bytes(rows, read_options=read_options, |
| convert_options=convert_options) |
| expected_schema = pa.schema([('g', pa.float64()), |
| ('f', pa.int64()), |
| ('e', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'g': [None, None], |
| 'e': ["2", "5"], |
| 'f': [3, 6]}]) |
| |
| def test_encoding(self): |
| # latin-1 (invalid utf-8) |
| rows = b"a,b\nun,\xe9l\xe9phant" |
| read_options = ReadOptions() |
| reader = self.open_bytes(rows, read_options=read_options) |
| expected_schema = pa.schema([('a', pa.string()), |
| ('b', pa.binary())]) |
| self.check_reader(reader, expected_schema, |
| [{'a': ["un"], |
| 'b': [b"\xe9l\xe9phant"]}]) |
| |
| read_options.encoding = 'latin1' |
| reader = self.open_bytes(rows, read_options=read_options) |
| expected_schema = pa.schema([('a', pa.string()), |
| ('b', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'a': ["un"], |
| 'b': ["éléphant"]}]) |
| |
| # utf-16 |
| rows = (b'\xff\xfea\x00,\x00b\x00\n\x00u\x00n\x00,' |
| b'\x00\xe9\x00l\x00\xe9\x00p\x00h\x00a\x00n\x00t\x00') |
| read_options.encoding = 'utf16' |
| reader = self.open_bytes(rows, read_options=read_options) |
| expected_schema = pa.schema([('a', pa.string()), |
| ('b', pa.string())]) |
| self.check_reader(reader, expected_schema, |
| [{'a': ["un"], |
| 'b': ["éléphant"]}]) |
| |
| def test_small_random_csv(self): |
| csv, expected = make_random_csv(num_cols=2, num_rows=10) |
| reader = self.open_bytes(csv) |
| table = reader.read_all() |
| 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 = [19, 21, 23, 26, 37, 111] |
| csvs = [csv_base, csv_base.rstrip(b'\r\n')] |
| for csv in csvs: |
| for block_size in block_sizes: |
| # Need at least two lines for type inference |
| assert csv[:block_size].count(b'\n') >= 2 |
| read_options = ReadOptions(block_size=block_size) |
| reader = self.open_bytes( |
| csv, read_options=read_options) |
| table = reader.read_all() |
| assert table.schema == expected.schema |
| if not table.equals(expected): |
| # Better error output |
| assert table.to_pydict() == expected.to_pydict() |
| |
| def test_batch_lifetime(self): |
| gc.collect() |
| old_allocated = pa.total_allocated_bytes() |
| |
| # Memory occupation should not grow with CSV file size |
| def check_one_batch(reader, expected): |
| batch = reader.read_next_batch() |
| assert batch.to_pydict() == expected |
| |
| rows = b"10,11\n12,13\n14,15\n16,17\n" |
| read_options = ReadOptions() |
| read_options.column_names = ['a', 'b'] |
| read_options.block_size = 6 |
| reader = self.open_bytes(rows, read_options=read_options) |
| check_one_batch(reader, {'a': [10], 'b': [11]}) |
| allocated_after_first_batch = pa.total_allocated_bytes() |
| check_one_batch(reader, {'a': [12], 'b': [13]}) |
| assert pa.total_allocated_bytes() <= allocated_after_first_batch |
| check_one_batch(reader, {'a': [14], 'b': [15]}) |
| assert pa.total_allocated_bytes() <= allocated_after_first_batch |
| check_one_batch(reader, {'a': [16], 'b': [17]}) |
| assert pa.total_allocated_bytes() <= allocated_after_first_batch |
| with pytest.raises(StopIteration): |
| reader.read_next_batch() |
| assert pa.total_allocated_bytes() == old_allocated |
| reader = None |
| assert pa.total_allocated_bytes() == old_allocated |
| |
| def test_header_skip_rows(self): |
| super().test_header_skip_rows() |
| |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| # Skipping all rows immediately results in end of iteration |
| opts = ReadOptions() |
| opts.skip_rows = 4 |
| opts.column_names = ['ab', 'cd'] |
| reader = self.open_bytes(rows, read_options=opts) |
| with pytest.raises(StopIteration): |
| assert reader.read_next_batch() |
| |
| def test_skip_rows_after_names(self): |
| super().test_skip_rows_after_names() |
| |
| rows = b"ab,cd\nef,gh\nij,kl\nmn,op\n" |
| |
| # Skipping all rows immediately results in end of iteration |
| opts = ReadOptions() |
| opts.skip_rows_after_names = 3 |
| reader = self.open_bytes(rows, read_options=opts) |
| with pytest.raises(StopIteration): |
| assert reader.read_next_batch() |
| |
| # Skipping beyond all rows immediately results in end of iteration |
| opts.skip_rows_after_names = 99999 |
| reader = self.open_bytes(rows, read_options=opts) |
| with pytest.raises(StopIteration): |
| assert reader.read_next_batch() |
| |
| |
| class TestSerialStreamingCSVRead(BaseStreamingCSVRead, unittest.TestCase): |
| @property |
| def use_threads(self): |
| return False |
| |
| |
| class TestThreadedStreamingCSVRead(BaseStreamingCSVRead, unittest.TestCase): |
| @property |
| def use_threads(self): |
| return True |
| |
| |
| class BaseTestCompressedCSVRead: |
| |
| def setUp(self): |
| self.tmpdir = tempfile.mkdtemp(prefix='arrow-csv-test-') |
| |
| def tearDown(self): |
| shutil.rmtree(self.tmpdir) |
| |
| def read_csv(self, csv_path): |
| try: |
| return read_csv(csv_path) |
| except pa.ArrowNotImplementedError as e: |
| pytest.skip(str(e)) |
| |
| 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) |
| table = self.read_csv(csv_path) |
| table.validate(full=True) |
| 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) |
| |
| def test_concatenated(self): |
| # ARROW-5974 |
| csv_path = os.path.join(self.tmpdir, self.csv_filename) |
| with gzip.open(csv_path, 'wb', 3) as f: |
| f.write(b"ab,cd\nef,gh\n") |
| with gzip.open(csv_path, 'ab', 3) as f: |
| f.write(b"ij,kl\nmn,op\n") |
| table = self.read_csv(csv_path) |
| assert table.to_pydict() == { |
| 'ab': ['ef', 'ij', 'mn'], |
| 'cd': ['gh', 'kl', 'op'], |
| } |
| |
| |
| 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 |
| |
| |
| def test_write_read_round_trip(): |
| t = pa.Table.from_arrays([[1, 2, 3], ["a", "b", "c"]], ["c1", "c2"]) |
| record_batch = t.to_batches(max_chunksize=4)[0] |
| for data in [t, record_batch]: |
| # Test with header |
| buf = io.BytesIO() |
| write_csv(data, buf, WriteOptions(include_header=True)) |
| buf.seek(0) |
| assert t == read_csv(buf) |
| |
| # Test without header |
| buf = io.BytesIO() |
| write_csv(data, buf, WriteOptions(include_header=False)) |
| buf.seek(0) |
| |
| read_options = ReadOptions(column_names=t.column_names) |
| assert t == read_csv(buf, read_options=read_options) |
| |
| # Test with writer |
| for read_options, write_options in [ |
| (None, WriteOptions(include_header=True)), |
| (ReadOptions(column_names=t.column_names), |
| WriteOptions(include_header=False)), |
| ]: |
| buf = io.BytesIO() |
| with CSVWriter(buf, t.schema, write_options=write_options) as writer: |
| writer.write_table(t) |
| buf.seek(0) |
| assert t == read_csv(buf, read_options=read_options) |
| |
| buf = io.BytesIO() |
| with CSVWriter(buf, t.schema, write_options=write_options) as writer: |
| for batch in t.to_batches(max_chunksize=1): |
| writer.write_batch(batch) |
| buf.seek(0) |
| assert t == read_csv(buf, read_options=read_options) |
| |
| |
| def test_read_csv_reference_cycle(): |
| # ARROW-13187 |
| def inner(): |
| buf = io.BytesIO(b"a,b,c\n1,2,3\n4,5,6") |
| table = read_csv(buf) |
| return weakref.ref(table) |
| |
| with util.disabled_gc(): |
| wr = inner() |
| assert wr() is None |
