如何用Rust实现Polars中DataFrame的分组遍历（附Python示例）

需求背景 #

对应Python中的分组遍历逻辑：

df = pl.DataFrame({"foo": ["a", "a", "b"], "bar": [1, 2, 3]})
for name, data in df.group_by("foo"):  
    print(name)
    print(data)

输出：

(a,)
shape: (2, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ a   ┆ 1   │
│ a   ┆ 2   │
└─────┴─────┘
(b,)
shape: (1, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ b   ┆ 3   │
└─────┴─────┘

Rust实现代码 #

在Rust的Polars中，你可以通过group_by配合into_iter遍历分组，代码如下：

use polars::prelude::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 创建目标DataFrame
    let df = DataFrame::new(vec![
        Series::new("foo", ["a", "a", "b"]),
        Series::new("bar", [1, 2, 3]),
    ])?;

    // 分组并遍历每个分组项
    for group in df.group_by(["foo"])? {
        let (keys, group_df) = group;
        
        // 打印分组键
        println!("{keys:?}");
        // 打印分组对应的数据框
        println!("{group_df}");
    }

    Ok(())
}

代码说明 #

• 导入Polars的prelude模块，包含常用的结构体与方法；
• 用Series::new构建列，再传入DataFrame::new创建数据框，逻辑和Python一致；
• group_by接收分组列的数组（单列也要用数组形式），返回Result类型，需用?处理错误；
• 遍历分组时，每个迭代项是元组，包含分组键keys和对应的数据框group_df；
• 分组键用{keys:?}格式化输出，语义和Python的(a,)一致，数据框的打印格式与Python版本完全匹配。

运行输出 #

["a"]
shape: (2, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ str ┆ i32 │
╞═════╪═════╡
│ a   ┆ 1   │
│ a   ┆ 2   │
└─────┴─────┘
["b"]
shape: (1, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ str ┆ i32 │
╞═════╪═════╡
│ b   ┆ 3   │
└─────┴─────┘

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