如何在Rust中将含字符串的Polars DataFrame转为ndarray（无需独热编码）

Polars 的 to_ndarray 方法仅支持数值类型列，直接转换含字符串的 DataFrame 会失败。要在不使用独热编码的情况下完成转换，需先对字符串列执行标签编码（将每个唯一字符串映射为唯一数值），再将处理后的全数值 DataFrame 转为 ndarray。

实现步骤 #

1. 提取字符串列的唯一值，建立字符串到数值的映射关系；
2. 将原 DataFrame 中的字符串列替换为对应数值；
3. 移除原字符串列，保留全数值列后转换为 ndarray。

示例代码 #

use polars::prelude::*;
use ndarray::Array2;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 构造含字符串列的示例 DataFrame
    let df = df!(
        "fruit" => ["apple", "banana", "apple", "orange", "banana"],
        "price" => [3.5, 2.8, 3.2, 4.0, 2.9]
    )?;

    // 为字符串列生成标签编码映射
    let unique_fruits = df.column("fruit")?.unique()?;
    let fruit_map: Vec<(&str, f64)> = unique_fruits
        .utf8()?
        .into_iter()
        .enumerate()
        .map(|(idx, fruit)| (fruit.unwrap(), idx as f64))
        .collect();

    // 将字符串列替换为编码后的数值列
    let df_encoded = df.with_column(
        col("fruit")
            .replace(&fruit_map, None)
            .alias("fruit_encoded")
    )?;

    // 移除原字符串列，保留全数值列
    let df_numeric = df_encoded.drop("fruit")?;

    // 转换为 ndarray
    let array: Array2<f64> = df_numeric.to_ndarray::<Float64Type>(Default::default())?;

    println!("转换后的 ndarray:\n{:?}", array);
    Ok(())
}

关键说明 #

• 标签编码为每个唯一字符串分配连续整数（转为 f64 以匹配 Float64Type），相同字符串对应相同数值，保证数据一致性；
• 多字符串列场景下，对每个列重复编码步骤即可；
• 若无需连续数值映射，也可使用字符串哈希值直接转换，但哈希存在冲突风险，标签编码可控性更强。

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