如何将含浮点numpy数组列的Dask DataFrame写入Parquet文件？

没错，你猜的完全对——当列的 dtype 是 object 时，Parquet 序列化器会默认把它当作字符串处理，而 numpy 数组显然不是字符串，这就导致了序列化失败。下面给你两种靠谱的解决思路，亲测有效：

方法一：显式指定数组类型并转换列 #

Dask 其实支持直接处理 numpy 数组类型的列，只是默认会把数组识别为 object。你可以先确认列中所有数组的形状一致（比如都是长度为N的一维浮点数组），然后用 astype 把列转换为正确的 numpy 数组 dtype：

import dask.dataframe as dd
import numpy as np

# 假设你的Dask DataFrame是df，'vec'列是一维float64数组
# 指定目标dtype：这里以长度5的float64数组为例，根据实际情况调整
array_dtype = np.dtype('float64', ndim=1)

# 转换列类型
df['vec'] = df['vec'].astype(array_dtype)

# 写入Parquet，注意指定engine='pyarrow'（fastparquet对数组类型支持有限）
df.to_parquet('your_file.parquet', engine='pyarrow')

方法二：使用PyArrow的FixedShapeTensorType（更推荐） #

Parquet 原生支持张量类型，通过 PyArrow 的 FixedShapeTensorType 可以更规范地处理固定形状的数组列，这种方式兼容性更好，也更符合Parquet的设计：

import dask.dataframe as dd
import numpy as np
import pyarrow as pa

# 假设数组是长度为5的一维float64数组，定义张量类型
tensor_type = pa.FixedShapeTensorType(pa.float64(), shape=(5,))

# 把pandas列转换为PyArrow的张量数组，适配Dask的分区处理
df['vec'] = df['vec'].map_partitions(
    lambda s: pa.array(s.to_list(), type=tensor_type),
    meta=('vec', tensor_type)
)

# 写入Parquet
df.to_parquet('your_file.parquet', engine='pyarrow')

关键注意事项 #

• 统一数组形状：不管用哪种方法，必须确保列中所有numpy数组的形状完全一致，Parquet作为列存格式，需要固定的列结构，形状不一致会直接导致序列化失败。
• 指定PyArrow引擎：一定要使用 engine='pyarrow'，fastparquet对复杂数组类型的支持远不如pyarrow完善，容易出现兼容性问题。
• 可变长度数组的特殊处理：如果你的数组是可变长度的，Parquet可能不是最优选择，这时候可以考虑把数组拆分成多个列，或者使用HDF5等支持可变长度数组的格式；但如果是固定长度的浮点数组，上面的方法完全可以解决问题。

内容的提问来源于stack exchange，提问作者junichiro