使用Dask读取指定PyArrow数据类型的CSV文件时遇到问题

我完全理解你的困扰——明明已经指定了engine='pyarrow'和dtype_backend='pyarrow'，但直接传PyArrow的pa.DataType给dtype参数却不生效，还要手动转成numpy类型确实麻烦。其实Dask针对这种场景是有内置解决方案的，不需要自己做类型映射，下面给你两种可行的方法：

方法一：使用PyArrow Schema对象指定类型（推荐） #

当使用pyarrow引擎时，Dask的read_csv支持直接传入PyArrow的Schema对象作为schema参数，这会直接让PyArrow读取器使用你定义的类型，完全符合你的需求。修改后的代码如下：

import dask.dataframe as dd
import pyarrow as pa

# 定义PyArrow Schema，明确每个列的类型
pa_schema = pa.schema([
    ('A', pa.string()),
    ('B', pa.float64())
])

file_path = './test.csv'

# 用schema参数替代原来的dtype，配合pyarrow引擎和dtype_backend
df = dd.read_csv(
    urlpath=file_path,
    sep=';',
    engine='pyarrow',
    dtype_backend='pyarrow',
    schema=pa_schema
)

df = df.sort_values('A')
df.to_csv(file_path, index=False, single_file=True, sep=';')

方法二：使用Pandas的PyArrow-backed数据类型 #

如果你更习惯用字典形式指定列类型，可以使用Pandas基于PyArrow实现的扩展数据类型，这样传给dtype参数就能生效，代码如下：

import dask.dataframe as dd
import pandas as pd
import pyarrow as pa

# 使用Pandas的PyArrow-backed dtype
pd_pa_dtypes = {
    'A': pd.StringDtype(storage='pyarrow'),
    'B': pd.Float64Dtype(storage='pyarrow')
}

file_path = './test.csv'

df = dd.read_csv(
    urlpath=file_path,
    sep=';',
    dtype=pd_pa_dtypes,
    engine='pyarrow',
    dtype_backend='pyarrow'
)

df = df.sort_values('A')
df.to_csv(file_path, index=False, single_file=True, sep=';')

为什么原来的方法不生效？ #

这里需要说明一下：Dask的read_csv的dtype参数目前并不直接接受PyArrow的pa.DataType对象——即使你指定了pyarrow引擎。dtype_backend='pyarrow'的作用是让Dask将读取后的数据转换为基于PyArrow存储的Pandas数据类型，而不是让dtype参数直接接受PyArrow类型。

这两种方法都适配你使用的版本（pyarrow==14.0.2、dask==2023.12.1），不需要再手动做类型映射，是更原生的解决方式。

备注：内容来源于stack exchange，提问作者newbnoob