Add column/series to dataframe in polars rust
Asked Answered
N

1

6

I had a hard time finding the answer for such simple question. I was stuck at trying to use "append", "extend" or other methods. Finally I found/realized that the with_column method is the way to go in polars.

I figure that I should put out my solution here for others who get stuck on the same problem.

Nadabb answered 28/1, 2023 at 8:51 Comment(0)
N
10
use polars::prelude::*;
fn main() {
    let a = Series::new("A", vec![1, 2, 3]);
    let b = Series::new("B", vec!["a", "b", "c"]);

    let mut df = DataFrame::new(vec![a, b]).unwrap();
    println!("{:?}", df);

    let c = Series::new("C", vec![true, false, false]);
    df.with_column(c).unwrap();
    println!("{:?}", df);

    let d = Series::new("D", vec![1.0, 2.0, 3.0]);
    let e = Series::new("E", vec![false, true, true]);

    // Also works with lazy and multiple series at once
    let df_lazy = df
        .lazy()
        .with_columns([d.lit(), e.lit()])
        .collect()
        .unwrap();
    println!("{:?}", df_lazy);
}

Output

┌─────┬─────┐
│ A   ┆ B   │
│ --- ┆ --- │
│ i32 ┆ str │
╞═════╪═════╡
│ 1   ┆ a   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 2   ┆ b   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 3   ┆ c   │
└─────┴─────┘
shape: (3, 3)
┌─────┬─────┬───────┐
│ A   ┆ B   ┆ C     │
│ --- ┆ --- ┆ ---   │
│ i32 ┆ str ┆ bool  │
╞═════╪═════╪═══════╡
│ 1   ┆ a   ┆ true  │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2   ┆ b   ┆ false │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 3   ┆ c   ┆ false │
└─────┴─────┴───────┘
shape: (3, 5)
┌─────┬─────┬───────┬─────┬───────┐
│ A   ┆ B   ┆ C     ┆ D   ┆ E     │
│ --- ┆ --- ┆ ---   ┆ --- ┆ ---   │
│ i32 ┆ str ┆ bool  ┆ f64 ┆ bool  │
╞═════╪═════╪═══════╪═════╪═══════╡
│ 1   ┆ a   ┆ true  ┆ 1.0 ┆ false │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2   ┆ b   ┆ false ┆ 2.0 ┆ true  │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 3   ┆ c   ┆ false ┆ 3.0 ┆ true  │
└─────┴─────┴───────┴─────┴───────┘

Nadabb answered 28/1, 2023 at 8:51 Comment(4)
I didn't need the .lit() but I'm not sure if that was because I was running it in Python or because it was an eager dataframe :shrug:Lepage
The syntax is somewhat different in python. I've not tested myself but try pl.lit(<data>). see docs.pola.rs/py-polars/html/reference/expressions/api/…Nadabb
I meant that I just did data = data.with_columns(fill_data). Not data = data.with_columns(pl.lit(fill_data)).Lepage
Oh sorry, I misread. Well then it's even simpler, thats great!Nadabb

© 2022 - 2024 — McMap. All rights reserved.