# 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
