使用Parquet文件时如何保留Pandas的category数据类型？

问题现象 #

将包含pd.NA的数值型序列转为Category类型后，写入Parquet文件再读取，该列类型会从Category变为float64；而字符串型的Category列则能正常保留类型。

原因分析 #

• 数值型序列中的pd.NA在pandas中会被默认处理为浮点型NaN，将这类序列转为Category时，底层存储的是浮点型值而非整数，Parquet无法识别这种混合类型的Category元数据，读取时会还原为浮点类型。
• 字符串型序列的pd.NA属于字符串缺失值，转为Category后底层类型统一，因此能被Parquet正确保留。

解决方案 #

1. 先将数值型序列转为Nullable Integer类型（Int64，首字母大写），再转为Category，确保缺失值的类型统一为整数类型的缺失值而非浮点NaN。
2. 指定engine='pyarrow'写入Parquet，pyarrow对Nullable类型和Category的支持更完善。

修改后代码示例 #

import pandas as pd 

df = pd.DataFrame({
    'a': [pd.NA, 'a', 'b', 'c'], 
    'b': [1,2,3,pd.NA]
})

# 先转为Nullable Integer类型，再转Category
df['b'] = df['b'].astype("Int64").astype("category")
df['a'] = df['a'].astype("category")

print("dtype before parquet write/read: ", df.dtypes, sep='\n')
# 使用pyarrow引擎写入
df.to_parquet('trial.parquet.gzip', compression='gzip', engine='pyarrow')

df = pd.read_parquet('trial.parquet.gzip', engine='pyarrow')
print("dtype after parquet write/read:", df.dtypes, sep='\n')

输出结果 #

dtype before parquet write/read: 
a    category
b    category
dtype: object
dtype after parquet write/read:
a    category
b    category
dtype: object

补充说明 #

• Nullable Integer类型（Int64）可以在保留整数类型的同时支持缺失值，避免转为浮点型NaN。
• pyarrow引擎相比默认的fastparquet，对pandas的Nullable类型和Category类型的兼容性更好，能更准确地保留元数据。

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