CODEWITHSUNDEEP
sqldataframerustrust-sqlxrust-polars
Hamza Zubair
Hamza Zubair

Reputation: 1410

Rust: DataFrame in polars using a vector of structs

Problem

I want to read in data into polars dataframe from mysql database. I am using sqlx.

sqlx generates a vector of Structs for example: Vec<Country> below:

From sqlx Docs:

// no traits are needed
struct Country { country: String, count: i64 }

let countries = sqlx::query_as!(Country,
        "
SELECT country, COUNT(*) as count
FROM users
GROUP BY country
WHERE organization = ?
        ",
        organization
    )
    .fetch_all(&pool) // -> Vec<Country>
    .await?;

// countries[0].country
// countries[0].count

How do i use this Vec<Country> to generate a polars dataframe

From polars Docs:


 use polars_core::prelude::*;
 let s0 = Series::new("a", &[1i64, 2, 3]);
 let s1 = Series::new("b", &[1i64, 1, 1]);
 let s2 = Series::new("c", &[2i64, 2, 2]);
 let list = Series::new("foo", &[s0, s1, s2]);

 let s0 = Series::new("B", [1, 2, 3]);
 let s1 = Series::new("C", [1, 1, 1]);
 let df = DataFrame::new(vec![list, s0, s1]).unwrap();

Probable Solution

The only solution i can think of is, if i can create a series for every column/Data inside the Country struct and use those individual series to create a dataframe.

I have no idea how to break down a Vec<Country> into Vec<country> and Vec<count>

Upvotes: 3

Views: 3527

Answers (2)

ritchie46
ritchie46

Reputation: 14660

You could use the builders for that or collect from iterators. Collecting from iterators is often fast, but in this case it requires you to loop the Vec<Country> twice, so you should benchmark.

Below is an example function for both the solutions shown.

use polars::prelude::*;

struct Country {
    country: String,
    count: i64,
}

fn example_1(values: &[Country]) -> (Series, Series) {
    let ca_country: Utf8Chunked = values.iter().map(|v| &*v.country).collect();
    let ca_count: NoNull<Int64Chunked> = values.iter().map(|v| v.count).collect();
    let mut s_country: Series = ca_country.into();
    let mut s_count: Series = ca_count.into_inner().into();
    s_country.rename("country");
    s_count.rename("country");
    (s_count, s_country)
}

fn example_2(values: &[Country]) -> (Series, Series) {
    let mut country_builder = Utf8ChunkedBuilder::new("country", values.len(), values.len() * 5);
    let mut count_builder = PrimitiveChunkedBuilder::<Int64Type>::new("count", values.len());

    values.iter().for_each(|v| {
        country_builder.append_value(&v.country);
        count_builder.append_value(v.count)
    });

    (
        count_builder.finish().into(),
        country_builder.finish().into(),
    )
}

Btw, if you want maximum performance, I really recommend connector-x. It has got polars and arrow integration and has got insane performance.

Upvotes: 5

Kendas
Kendas

Reputation: 2243

I do not know about the structure of these DataFrames or Series but if you want to split out a vector of structs, you can write a small reducer that iterates over the vector's contents and pushes them into multiple new vectors.

A simplified playground snippet: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=f698866751214fc3f44c348b7c4f80c5

struct A(u8, i8);

fn main() {
    let v = vec![A(1, 4), A(2, 6), A(3, 5)];

    let result = v.into_iter()
        .fold((vec![], vec![]), |(mut u, mut i), item| {
            u.push(item.0);
            i.push(item.1);
            (u, i)
        });
    
    dbg!(result);

    // `result` is just a tuple of vectors
    // let (unsigneds, signeds): (Vec<u8>, Vec<i8>) = result;
}

You basically take a vector of structs, you iterate over them and fold them into the provided new (empty) vectors one by one. At the end, the result is returned (a tuple of 2 vectors). You can do what you want with these now.

Upvotes: 2

Related Questions