Как добавить новую категорию в столбец Polars?

Похоже, категории, которые вы хотите использовать ("A", "B", "C" и "Unknown"), известны заранее. В этом случае следует использовать pl.Enum.

Рассмотрим следующий пример данных.

import polars as pl

df_train = pl.DataFrame({
    "cat": ["A", "B", "C"],
})

df_test = pl.DataFrame({
    "cat": ["A", "X", "C"],
})

На основе данных обучения можно создать тип данных pl.Enum, включающий дополнительную категорию "Unknown", следующим образом.

enum_dtype = pl.Enum(
    df_train.get_column("cat").unique().drop_nulls().append(pl.Series(["Unknown"]))
)

Примечание. .drop_nulls() используется (хотя данные в этом примере не содержат значений NULL) в качестве меры предосторожности, поскольку категории Enum не могут содержать значения NULL.

Теперь категориальный столбец в обучающем фрейме данных можно привести непосредственно к типу данных pl.Enum.

df_train.with_columns(pl.col("cat").cast(enum_dtype))

shape: (3, 1)
┌──────┐
│ cat  │
│ ---  │
│ enum │
╞══════╡
│ A    │
│ B    │
│ C    │
└──────┘

Для тестового фрейма данных это не удастся, поскольку "D" не является известной категорией в типе данных pl.Enum. Здесь использование strict=False можно использовать для приведения неизвестных категорий к None. Впоследствии пропущенные значения можно заменить категорией "Unknown".

df_test.with_columns(
    pl.col("cat").cast(enum_dtype, strict=False).fill_null("Unknown")
)

shape: (3, 1)
┌─────────┐
│ cat     │
│ ---     │
│ enum    │
╞═════════╡
│ A       │
│ Unknown │
│ C       │
└─────────┘

Редактировать. Если тестовые данные уже содержат пропущенные значения, которые следует сохранить как таковые, неизвестные значения можно явно заменить на "Unknown" перед (строгим) приведением.

df_test.with_columns(
    pl.when(
        pl.col("cat").is_not_null() & ~pl.col("cat").is_in(enum_dtype.categories)
    ).then(
        pl.lit("Unknown")
    ).otherwise(
        pl.col("cat")
    ).cast(enum_dtype)
)