将 pyo3-polars 与“group_by”一起使用的最佳实践

问：

我目前正在试验 pyo3-polars 来优化数据聚合。从更抽象的意义上讲，我想到的是以下结构：

df.group_by(c.col1).agg(c.col2.foo())

这里，表示产生结果的非平凡函数。为了提供更具体的解释，可以使用以下伪代码来说明该函数：foof64foo

agg = 0.0
for i in 0..len(col2):
    for j in i+1..len(col2):
        func(agg, col2[i], col2[j])
return agg

我拥有的当前实现旨在返回：Series

#[polars_expr(output_type=Int64)]
fn int_agg(inputs: &[Series]) -> PolarsResult<Series> {
    let ca = inputs[0].i64()?;

    let mut v = Vec::new();
    let mut tot: i64 = 0;

    for i in 0..(ca.len()) { // Fake approximation of actual task
        unsafe {
            tot += ca.value_unchecked(i); 
        }
    }
    v.push(tot);
    let out: Int64Chunked = v.into_iter().map(|i| Some(i + 2)).collect_ca(ca.name());
    Ok(out.into_series())
}

但是，对于这是否是完成此特定任务的最佳方法，我持保留态度。存储库中显示的所有示例都返回强制创建包含单个元素的 Series。在 python 会话的后期，由于这种人为约束，我不得不串行调用以获取元素。有没有更好的替代方案可以直接返回单个元素？PolarsResult<Series>foo().first()

此外，我可能对更全面的解决方案感兴趣，其中聚合函数接受多个列，例如 .df.group_by("col1").agg(foo("col2", "col3"))