如何避免pandas to_parquet将列表类型转为numpy数组？

我有一个包含list类型列的DataFrame，将其写入parquet文件后再读取，原本的list类型会变为numpy array类型。

最小可复现代码 #

import os
import pandas as pd

home_directory = os.path.join(rf"C:\Users\{os.getlogin()}")
df = pd.DataFrame({'my_column': [['a', 'b', 'c'], ['c', 'd', 'e']]})
print(type(df['my_column'].iloc[0]))
df.to_parquet(os.path.join(home_directory,'data'), engine='pyarrow')
df_from_parquet = pd.read_parquet(os.path.join(home_directory,'data'), engine='pyarrow')
print(type(df_from_parquet['my_column'].iloc[0]))

输出结果 #

<class 'list'>
<class 'numpy.ndarray'>

解决方法 #

方法1：读取后手动转换列类型 #

直接对读取后的列使用apply(list)将numpy数组转回list：

df_from_parquet['my_column'] = df_from_parquet['my_column'].apply(list)
print(type(df_from_parquet['my_column'].iloc[0]))  # 输出 <class 'list'>

方法2：写入时指定Arrow Schema #

通过pyarrow定义明确的列表类型Schema，确保写入和读取时保持list结构：

import pyarrow as pa
import pyarrow.parquet as pq

# 定义包含列表类型的Schema
schema = pa.schema([
    ('my_column', pa.list_(pa.string()))
])

# 转换为Arrow Table并写入Parquet
table = pa.Table.from_pandas(df, schema=schema)
pq.write_table(table, os.path.join(home_directory,'data'))

# 读取并转回DataFrame
df_from_parquet = pq.read_table(os.path.join(home_directory,'data')).to_pandas()
print(type(df_from_parquet['my_column'].iloc[0]))  # 输出 <class 'list'>

方法3：读取时指定dtype_backend参数（Pandas 2.0+） #

在read_parquet中设置dtype_backend='numpy_nullable'，会将Arrow列表类型映射为Pandas原生的List dtype，避免转为numpy数组：

df_from_parquet = pd.read_parquet(
    os.path.join(home_directory,'data'), 
    engine='pyarrow', 
    dtype_backend='numpy_nullable'
)
print(type(df_from_parquet['my_column'].iloc[0]))  # 输出 <class 'list'>

内容的提问来源于stack exchange，提问作者E. Sommer