Search code examples
dataframerustpython-polarsrust-polars

How to create polars DataFrame with Vec<Vec<f64>> as a Series


I desire a DataFrame like so:

timestamp  |          bids         |         asks         |  ticker
-------------------------------------------------------------------
1598215600 |  [[10, 20], [15, 30]] | [[20, 10], [25, 20]] |  "AAPL"
1598222400 |  [[11, 25], [16, 35]] | [[22, 15], [28, 25]] |  "MSFT"
1598229200 |  [[12, 30], [18, 40]] | [[24, 20], [30, 30]] |  "GOOG"

The bids Series has a Vec<Vec> structure, which in plain words is a vector that holds a pair (another vector) of the price and amount (two values).

What is the required rust code to create this? If possible answer in rust, but python works too I guess I can recreate it.


Solution

  • I'm new to rust so it's possible this is not optimal.

    From looking around it seems like ChunkedArray may be the way to go?

    use polars::prelude::*;
    
    fn build_column(rows: &Vec<[[i64; 2]; 2]>) -> Series {
        ListChunked::from_iter(rows.into_iter().map(|row| {
            ListChunked::from_iter(
                row.into_iter()
                    .map(|values| Int64Chunked::from_slice("", values).into_series()),
            )
            .into_series()
        }))
        .into_series()
    }
    
    fn main() -> PolarsResult<()> {
        let asks = vec![
            [[20, 10], [25, 20]],
            [[22, 15], [28, 25]],
            [[24, 20], [30, 30]],
        ];
    
        let bids = vec![
            [[10, 20], [15, 30]],
            [[11, 25], [16, 35]],
            [[12, 30], [18, 40]],
        ];
    
        let df = df!(
                "timestamp" => [1598215600, 1598222400, 1598229200],
                "asks" => build_column(&asks),
                "bids" => build_column(&bids),
                "ticker" => ["AAPL", "MSFT", "GOOG"]
        );
    
        println!("{:?}", df);
    
        Ok(())
    }
    
    Ok(shape: (3, 4)
    ┌────────────┬──────────────────────┬──────────────────────┬────────┐
    │ timestamp  ┆ asks                 ┆ bids                 ┆ ticker │
    │ ---        ┆ ---                  ┆ ---                  ┆ ---    │
    │ i32        ┆ list[list[i64]]      ┆ list[list[i64]]      ┆ str    │
    ╞════════════╪══════════════════════╪══════════════════════╪════════╡
    │ 1598215600 ┆ [[20, 10], [25, 20]] ┆ [[10, 20], [15, 30]] ┆ AAPL   │
    │ 1598222400 ┆ [[22, 15], [28, 25]] ┆ [[11, 25], [16, 35]] ┆ MSFT   │
    │ 1598229200 ┆ [[24, 20], [30, 30]] ┆ [[12, 30], [18, 40]] ┆ GOOG   │
    └────────────┴──────────────────────┴──────────────────────┴────────┘)