如何为Polars分类列添加新类别并处理未知类别？

核心思路 #

要实现你的需求，关键是对齐训练集和测试集的分类类别集合，同时区分空值和未知类别：

1. 先扩展训练集的分类列类别，加入Unknown；
2. 对测试集的非空值做判断，不在训练集原类别中的转为Unknown；
3. 空值要么保留，要么转为专属类别（按需选择）。

代码实现 #

1. 构造示例数据 #

先模拟你的训练集和测试集：

import polars as pl

# 训练集：包含A/B/C三类+空值，已转为分类列
train_df = pl.DataFrame({
    "cat_col": ["A", "B", "C", None, "A", "B"]
}).with_columns(
    pl.col("cat_col").cast(pl.Categorical)
)

# 测试集：包含已知类别、未知类别X/Y、空值
test_df = pl.DataFrame({
    "cat_col": ["A", "X", "C", None, "B", "Y"]
})

2. 处理训练集：添加Unknown类别 #

Polars的分类列通过cat.set_categories扩展类别，不会修改原有数据（包括空值）：

# 获取训练集原有类别，添加'Unknown'后更新分类列
train_categories = train_df["cat_col"].cat.categories() + ["Unknown"]
train_df = train_df.with_columns(
    pl.col("cat_col").cat.set_categories(train_categories, ordered=False)
)

3. 处理测试集：替换未知类别为Unknown，保留空值 #

用when/then/otherwise做分支判断，避免低效的map_elements，同时对齐训练集的分类类别：

# 提取训练集的有效类别（排除刚添加的'Unknown'）
train_valid_cats = set(train_categories) - {"Unknown"}

test_df = test_df.with_columns(
    pl.when(pl.col("cat_col").is_null())
    .then(pl.col("cat_col"))  # 保留空值
    .when(pl.col("cat_col").is_in(train_valid_cats))
    .then(pl.col("cat_col"))  # 保留已知类别
    .otherwise("Unknown")     # 未知类别转为'Unknown'
    .cast(pl.Categorical)
    .cat.set_categories(train_categories, ordered=False)  # 对齐训练集类别
)

可选：将空值转为专属类别 #

如果想把空值设为专属类别（比如NaN_Category），调整如下：

训练集处理

# 将训练集空值转为专属类别，再添加'Unknown'
train_df = train_df.with_columns(
    pl.col("cat_col").fill_null("NaN_Category")
)
train_categories = train_df["cat_col"].cat.categories() + ["Unknown"]
train_df = train_df.with_columns(
    pl.col("cat_col").cat.set_categories(train_categories, ordered=False)
)

测试集处理

train_valid_cats = set(train_categories) - {"Unknown"}

test_df = test_df.with_columns(
    pl.when(pl.col("cat_col").is_null())
    .then("NaN_Category")  # 空值转为专属类别
    .when(pl.col("cat_col").is_in(train_valid_cats))
    .then(pl.col("cat_col"))
    .otherwise("Unknown")
    .cast(pl.Categorical)
    .cat.set_categories(train_categories, ordered=False)
)

关键注意点 #

• 训练集和测试集的分类列必须使用完全一致的类别集合，否则后续建模会出现类别不匹配的错误；
• 优先使用when/then/otherwise而非map_elements，前者是向量化操作，处理大数据集效率更高；
• Polars的分类列原生支持空值，无需额外处理即可保留。

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