数值类型分类变量保存至Parquet时丢失的原因及解决方法

问题复现 #

字符串分类列（正常情况） #

创建字符串类型的分类列并保存加载：

import pandas as pd
df = pd.DataFrame({"a":['1','2','3']}).astype("category")

查看分类信息：

>>> df['a'].cat.categories
Index(['1', '2', '3'], dtype='object')

保存后重新加载，分类类型保留：

>>> df.to_parquet("1.parquet")
>>> pd.read_parquet("1.parquet")['a']
0    1
1    2
2    3
Name: a, dtype: category
Categories (3, object): ['1', '2', '3']

数值分类列（异常情况） #

创建整数类型的分类列，保存加载后分类类型丢失：

>>> df = pd.DataFrame({"a":[1,2,3]}).astype("category")
>>> df['a']
0    1
1    2
2    3
Name: a, dtype: category
Categories (3, int64): [1, 2, 3]
>>> df.to_parquet("1.parquet")
>>> pd.read_parquet("1.parquet")['a']
0    1
1    2
2    3
Name: a, dtype: int64

核心疑问 #

1. 为何数值型分类列在保存/加载后会丢失分类类型？
2. 如何避免这种情况？

问题解析与解决方案 #

1. 丢失原因 #

这是pyarrow引擎（pandas保存Parquet的默认引擎）的判定逻辑导致的：当分类列的取值为连续整数，且与内部编码索引的映射关系无额外信息时，pyarrow会认为分类类型是冗余的，直接将列以原始数值类型存储，不再保留分类元数据。而字符串分类列不存在这种判定，因为字符串值和编码索引完全不同，所以会完整保留分类信息。

2. 解决方案 #

方法一：切换为fastparquet引擎

fastparquet对数值分类列的元数据保留更严谨，无需额外配置即可完整保存分类类型：

# 保存
df.to_parquet("1.parquet", engine="fastparquet")
# 加载
loaded_df = pd.read_parquet("1.parquet", engine="fastparquet")
print(loaded_df['a'].dtype)  # 输出 category

方法二：加载时强制指定列类型

如果必须使用pyarrow引擎，加载时明确指定目标列的 dtype 为category即可还原：

loaded_df = pd.read_parquet("1.parquet", dtype={"a": "category"})
print(loaded_df['a'].dtype)  # 输出 category

方法三：转换分类值为非数值类型（可选）

将数值分类值转为字符串类型的分类，让pyarrow无法判定为冗余，从而保留分类元数据：

# 转换分类值为字符串
df['a'] = df['a'].cat.astype(str).astype('category')
# 保存加载
df.to_parquet("1.parquet")
loaded_df = pd.read_parquet("1.parquet")
print(loaded_df['a'].dtype)  # 输出 category

注意：此方法会改变分类值的类型，仅适用于不依赖数值类型的场景。

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