rustpython-polarsrust-polarspolars

How to get all group in polars by rust?


In python, just like this

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

output:

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

How to do same in rust?


Solution

  • This works a bit different in Rust, because a lazy DataFrame / LazyGroupBy is used here.
    You can do something like this using the apply function:

    use polars::prelude::*;
    
    fn main() -> PolarsResult<()> {
        let df = df!(
            "foo" => &["a", "a", "b"],
            "bar" => &[1, 2, 3]
        )?;
    
        let groups = df.group_by(["foo"])?;
        groups.apply(|group| {
            let group_name =
                group.column("foo")?.str()?.get(0).ok_or_else(|| {
                    PolarsError::ComputeError("Failed to retrieve group name.".into())
                })?;
            println!("{}", group_name);
            println!("{:?}", group);
            Ok(group)
        })?;
    
        Ok(())
    }