Search code examples
dataframerustserieslazy-evaluationrust-polars

Add column/series to dataframe in polars rust


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.


Solution

  • 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  │
    └─────┴─────┴───────┴─────┴───────┘