如何在Polars DataFrame中填充结构体元素并访问它们

你想直接在Polars里填充和使用自定义结构体，替代繁琐的JSON转换方法，同时需要了解如何取出存入的结构体，以下是具体实现方案：

前提说明 #

Polars的列类型基于自身的DataType体系，不支持直接存储任意Rust自定义结构体，但可以通过两种更直接的方式实现需求：

方案一：映射为Polars原生Struct类型（推荐） #

将自定义结构体的字段拆解为Polars Struct列的子字段，完全适配Polars的向量化性能，是最优方案。

1. 给自定义类型实现转换逻辑 #

先让你的枚举和结构体能转换成Polars支持的类型：

use polars::prelude::*;

#[derive(Debug, Clone)]
enum FriendType {
    Goodfriend,
    Badfriend,
}

// 把FriendType转成字符串，方便存入Polars
impl From<FriendType> for String {
    fn from(ft: FriendType) -> Self {
        match ft {
            FriendType::Goodfriend => "Goodfriend".to_string(),
            FriendType::Badfriend => "Badfriend".to_string(),
        }
    }
}

#[derive(Debug, Clone)]
struct MyStruct {
    pub person: Option<String>,
    pub relation: FriendType,
}

impl MyStruct {
    fn new(person: &str, relation: FriendType) -> Self {
        MyStruct {
            person: Some(person.to_string()),
            relation,
        }
    }

    // 把MyStruct转换成Polars的StructValue
    fn to_struct_value(&self) -> StructValue {
        let person_field = self.person.as_ref().map(|s| s.as_str()).into();
        let relation_field: String = self.relation.clone().into();
        
        StructValue::new(vec![
            ("person", person_field),
            ("relation", relation_field.into()),
        ])
    }
}

2. 构建Struct列并加入DataFrame #

将MyStruct实例转换为StructValue迭代器，再构建成Series加入DataFrame：

fn main() {
    // 创建结构体实例列表
    let struct_items = vec![
        MyStruct::new("Peter", FriendType::Badfriend),
        MyStruct::new("Mary", FriendType::Goodfriend),
        MyStruct::new("Elon", FriendType::Goodfriend),
        MyStruct::new("Joe", FriendType::Badfriend),
    ];

    // 转换为Polars StructValue集合
    let struct_values: Vec<_> = struct_items.iter().map(|s| s.to_struct_value()).collect();
    
    // 构建Struct类型的Series
    let struct_series = Series::new("MyStruct_col", struct_values);

    // 组装DataFrame
    let df = df![
        "Integer_col" => [1,2,3,4],
        "Float_col" => [1.1,0.3,9.6,4.2],
        "String_col" => ["apple", "pear", "lemon", "plum"],
        "values_nulls" => [Some(1), None, Some(3), Some(8)],
        "MyStruct_col" => struct_series,
    ].unwrap();

    println!("生成的DataFrame:\n{}", df);
}

3. 取出Struct列并还原为MyStruct #

遍历Struct列，提取子字段后还原成原始结构体：

// 取出Struct列
let struct_col = df.column("MyStruct_col").unwrap();
let struct_chunked = struct_col.struct_().unwrap();

for struct_val in struct_chunked.iter() {
    if let Some(sv) = struct_val {
        // 提取子字段
        let person: Option<&str> = sv.field_by_name("person").unwrap().try_extract().unwrap();
        let relation_str: &str = sv.field_by_name("relation").unwrap().try_extract().unwrap();

        // 还原FriendType枚举
        let relation = match relation_str {
            "Goodfriend" => FriendType::Goodfriend,
            "Badfriend" => FriendType::Badfriend,
            _ => panic!("未知的关系类型"),
        };

        // 还原MyStruct
        let restored_struct = MyStruct {
            person: person.map(|s| s.to_string()),
            relation,
        };

        println!("还原后的结构体: {:?}", restored_struct);
    }
}

方案二：使用Polars Object类型（仅适合小数据量） #

Polars支持Object类型存储任意实现Any trait的对象，但会失去向量化性能优势，不推荐大数据场景使用。

1. 给结构体实现必要Trait #

use polars::prelude::*;
use std::any::Any;

#[derive(Debug, Clone)]
enum FriendType {
    Goodfriend,
    Badfriend,
}

#[derive(Debug, Clone)]
struct MyStruct {
    pub person: Option<String>,
    pub relation: FriendType,
}

impl MyStruct {
    fn new(person: &str, relation: FriendType) -> Self {
        MyStruct {
            person: Some(person.to_string()),
            relation,
        }
    }
}

// 实现Any trait
impl Any for MyStruct {
    fn type_id(&self) -> std::any::TypeId {
        std::any::TypeId::of::<Self>()
    }
}

// 实现Polars的ObjectSafe trait
unsafe impl ObjectSafe for MyStruct {}

2. 构建Object列并加入DataFrame #

fn main() {
    let struct_items = vec![
        MyStruct::new("Peter", FriendType::Badfriend),
        MyStruct::new("Mary", FriendType::Goodfriend),
        MyStruct::new("Elon", FriendType::Goodfriend),
        MyStruct::new("Joe", FriendType::Badfriend),
    ];

    // 构建Object类型的Series
    let object_series = Series::new("MyStruct_col", struct_items);

    let df = df![
        "Integer_col" => [1,2,3,4],
        "Float_col" => [1.1,0.3,9.6,4.2],
        "String_col" => ["apple", "pear", "lemon", "plum"],
        "values_nulls" => [Some(1), None, Some(3), Some(8)],
        "MyStruct_col" => object_series,
    ].unwrap();
}

3. 取出Object列中的结构体 #

// 取出Object列
let object_col = df.column("MyStruct_col").unwrap();
let object_chunked = object_col.object::<MyStruct>().unwrap();

for item in object_chunked.iter() {
    if let Some(s) = item {
        println!("取出的结构体: {:?}", s);
    }
}

关键注意点 #

• 方案一是首选，完全适配Polars的性能优化，支持所有Polars的向量化操作；
• 方案二仅适合小数据场景，因为Object类型无法利用Polars的内部优化，查询效率低；
• 自定义枚举必须转换为Polars支持的基础类型（如字符串、整数）才能存入Struct列，取出时再还原。

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