使用 pyarrow 读取 CSV 文件时，如何为所有列指定 dtype？

Question

我想用 pyarrow 读取一个大的 CSV 文件。我所有的列都是 float64 的。但 pyarrow 似乎正在推断 int64。

如何为所有列指定数据类型？

import gcsfs
import pyarrow.dataset as ds

fs = gcsfs.GCSFileSystem(project='my-google-cloud-project')

my_dataset = ds.dataset("bucket/foo/bar.csv", format="csv", filesystem=fs)

my_dataset.to_table()

其产生：

ArrowInvalid                              Traceback (most recent call last)
........py in <module>
----> 65 my_dataset.to_table()

File /opt/conda/envs/py39/lib/python3.9/site-packages/pyarrow/_dataset.pyx:491, in pyarrow._dataset.Dataset.to_table()

File /opt/conda/envs/py39/lib/python3.9/site-packages/pyarrow/_dataset.pyx:3235, in pyarrow._dataset.Scanner.to_table()

File /opt/conda/envs/py39/lib/python3.9/site-packages/pyarrow/error.pxi:143, in pyarrow.lib.pyarrow_internal_check_status()

File /opt/conda/envs/py39/lib/python3.9/site-packages/pyarrow/error.pxi:99, in pyarrow.lib.check_status()

ArrowInvalid: In CSV column #172: Row #28: CSV conversion error to int64: invalid value '6.58841482364418'

Answer 1

Pyarrow 的数据集模块以块的形式读取 CSV 文件（我认为默认值为 1MB），并并行处理这些块。这使得列推断有点棘手，它通过使用第一个块来推断数据类型来处理这个问题。因此，当文件的第一个块的列看起来是整数但在未来的块中该列具有十进制值时，您收到的错误非常常见。

如果您事先知道列名，那么您可以指定列的数据类型：

import pyarrow as pa
import pyarrow.csv as csv
import pyarrow.dataset as ds

column_types = {'a': pa.float64(), 'b': pa.float64(), 'c': pa.float64()}
convert_options = csv.ConvertOptions(column_types=column_types)
custom_csv_format = ds.CsvFileFormat(convert_options=convert_options)
dataset = ds.dataset('/tmp/foo.csv', format=custom_csv_format)

如果您不知道列名，那么事情会有点棘手。然而，听起来所有列都是float64。在这种情况下，由于您只有一个文件，您可能可以执行以下操作作为解决方法：

dataset = ds.dataset('/tmp/foo.csv', format='csv')
column_types = {}
for field in dataset.schema:
  column_types[field.name] = pa.float64()
# Now use column_types as above

这有效，因为我们调用 pa.dataset(...) 两次，并且它会有一小部分的开销。这是因为每次我们调用 pa.dataset(...) pyarrow 都会打开数据集中第一个文件的第一个块来确定模式（这就是为什么我们可以使用 dataset .schema）

如果您有多个具有不同列的文件，那么此方法将不起作用。在这种情况下，我建议向 Arrow user@ 邮件列表发送邮件，我们可以就解决问题的不同方法进行更一般的讨论。