Method test_column_options

python/pyarrow/tests/test_csv.py:1793–1856

(self)

            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()),