如何使用 pyarrow.read_csv 添加/更改列名？

Question

我目前正在尝试将一个没有任何标题的大 csv 文件 (50GB+) 导入到 pyarrow 表中，总体目标是将此文件导出为 Parquet 格式，并进一步在 Pandas 或 Dask DataFrame 中处理它。如何在 pyarrow 中为 csv 文件指定列名和列 dtype？

我已经考虑将标头附加到 csv 文件中。这会强制完全重写文件，这看起来像是不必要的开销。据我所知，pyarrow 提供了模式来定义特定列的 dtypes，但是文档在将 csv 文件转换为箭头表时缺少这样做的具体示例。

想象一下，这个 csv 文件只是一个简单的例子，即两列“A”和“B”。 我当前的代码如下所示：

import numpy as np
import pandas as pd
import pyarrow as pa
df_with_header = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})

print(df_with_header)
df_with_header.to_csv("data.csv", header=False, index=False)

df_without_header = pd.read_csv('data.csv', header=None)
print(df_without_header)
opts = pa.csv.ConvertOptions(column_types={'A': 'int8',
                                            'B': 'int8'})

table = pa.csv.read_csv(input_file = "data.csv", convert_options = opts)
print(table)

如果我打印出最终表格，它不会更改列的名称。

pyarrow.Table
1: int64
3: int64

我现在如何更改加载的列名和数据类型？是否还有可能例如传入包含名称及其数据类型的字典？

Answer 1

您可以为列指定类型覆盖：

    fp = io.BytesIO(b'one,two,three\n1,2,3\n4,5,6')
    fp.seek(0)
    table = csv.read_csv(
        fp,
        convert_options=csv.ConvertOptions(
            column_types={
                'one': pa.int8(),
                'two': pa.int8(),
                'three': pa.int8(),
            }
        ))

但在您的情况下，您没有标题，据我所知，箭头不支持此用例：

    fp = io.BytesIO(b'1,2,3\n4,5,6')
    fp.seek(0)
    table = csv.read_csv(
        fp,
        parse_options=csv.ParseOptions(header_rows=0)
    )

这引发了：

pyarrow.lib.ArrowInvalid: header_rows == 0 needs explicit column names

代码在这里：https://github.com/apache/arrow/blob/3cf8f355e1268dd8761b99719ab09cc20d372185/cpp/src/arrow/csv/reader.cc#L138

这类似于这个问题apache arrow - reading csv file

下个版本应该会修复它：https://github.com/apache/arrow/pull/4898