Rust: 使用结构体向量在polars中创建DataFrame

Question

问题

我想从mysql数据库中读取数据到polars数据框架中。我正在使用sqlx。

sqlx生成一个结构体向量，例如下面的Vec<Country>：

来自sqlx 文档：

// 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

如何使用Vec<Country>生成一个polars dataframe

来自polars 文档:

 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();

可能的解决方案

我能想到的唯一解决方案是，如果我可以为Country结构体中的每一列/数据创建一个序列，并使用这些单独的序列创建一个数据框。

我不知道如何将Vec<Country>分解成Vec<country>和Vec<count>。

- Hamza Zubair

Answer 1

你可以使用构建器(builder)或者迭代器(iterators)来实现。从迭代器中收集通常很快，但在这种情况下需要将Vec<Country>循环两次，因此你应该进行基准测试。

下面是两个解决方案的示例函数。

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(),
    )
}

顺便说一下，如果你想要最大化性能，我真的推荐使用connector-x。它集成了polars和arrow，并且具有疯狂的性能。

- ritchie46

Answer 2

我不了解这些DataFrame或Series的结构，但如果您想拆分向量结构，则可以编写一个小型reducer，迭代向量内容并将其推入多个新向量中。

一个简化的游乐场片段： 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;
}

你基本上需要取一个结构体的向量，然后迭代它们并将它们一个接一个地折叠到提供的新（空）向量中。最后，返回结果（一个包含2个向量的元组）。现在，你可以根据自己的需求对它们进行操作。

- Kendas

---

性能注意事项，使用Vec::with_capacity(v.len())初始化空向量会更快。 - Jmb

我也问了这个问题，以便更直接地从SQL读取数据到Polars。 - Hamza Zubair