如何在Polars中将字符串列编码为整数？

给定的Polars DataFrame如下：

df = (
    pl.DataFrame({"animal": ['elephant', 'dog', 'cat', 'mouse'],
    "country": ['Mexico', 'Denmark', 'Mexico', 'France'], 
    "cost": [1000.0, 20.0, 10.0, 120.0]})
)

要实现字符串列的整数编码，有以下几种简洁方案：

方法1：利用Categorical类型自动编码 #

Polars的Categorical类型会为每个唯一字符串分配内置整数ID，通过to_physical()可直接提取该ID，编码从0开始：

df_encoded = df.with_columns(
    animal_encoded=pl.col("animal").cast(pl.Categorical).to_physical(),
    country_encoded=pl.col("country").cast(pl.Categorical).to_physical()
)

print(df_encoded)

输出结果：

shape: (4, 5)
┌───────────┬──────────┬───────┬────────────────┬──────────────────┐
│ animal    ┆ country  ┆ cost  ┆ animal_encoded ┆ country_encoded │
│ ---       ┆ ---      ┆ ---   ┆ ---            ┆ ---              │
│ str       ┆ str      ┆ f64   ┆ u32            ┆ u32              │
╞═══════════╪══════════╪═══════╪════════════════╪══════════════════╡
│ elephant  ┆ Mexico   ┆ 1000.0┆ 0              ┆ 0                │
│ dog       ┆ Denmark  ┆ 20.0  ┆ 1              ┆ 1                │
│ cat       ┆ Mexico   ┆ 10.0  ┆ 2              ┆ 0                │
│ mouse     ┆ France   ┆ 120.0 ┆ 3              ┆ 2                │
└───────────┴──────────┴───────┴────────────────┴──────────────────┘

方法2：用rank生成从1开始的编码 #

如果需要编码从1开始，可使用rank("dense")生成连续整数：

df_encoded = df.with_columns(
    animal_encoded=pl.col("animal").rank("dense").cast(pl.Int32),
    country_encoded=pl.col("country").rank("dense").cast(pl.Int32)
)

print(df_encoded)

输出结果：

shape: (4, 5)
┌───────────┬──────────┬───────┬────────────────┬──────────────────┐
│ animal    ┆ country  ┆ cost  ┆ animal_encoded ┆ country_encoded │
│ ---       ┆ ---      ┆ ---   ┆ ---            ┆ ---              │
│ str       ┆ str      ┆ f64   ┆ i32            ┆ i32              │
╞═══════════╪══════════╪═══════╪════════════════╪══════════════════╡
│ elephant  ┆ Mexico   ┆ 1000.0┆ 1              ┆ 2                │
│ dog       ┆ Denmark  ┆ 20.0  ┆ 2              ┆ 1                │
│ cat       ┆ Mexico   ┆ 10.0  ┆ 3              ┆ 2                │
│ mouse     ┆ France   ┆ 120.0 ┆ 4              ┆ 3                │
└───────────┴──────────┴───────┴────────────────┴──────────────────┘

方法3：手动创建映射表后关联 #

若需要明确管理编码映射关系，可先生成唯一值的编码表，再通过join关联到原DataFrame：

# 生成各列的编码映射
animal_map = df.select("animal").unique().with_row_index("animal_encoded")
country_map = df.select("country").unique().with_row_index("country_encoded")

# 关联映射表得到编码列
df_encoded = df.join(animal_map, on="animal").join(country_map, on="country")

print(df_encoded)

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