如何在Rust Polars中应用自定义函数并返回带命名字段的Struct

原代码及现有结果 #

以下是一个GroupBy聚合操作的简单示例，当前代码会将每组的min和max值以数组/向量形式作为单列返回：

#[pyfunction]
fn test_fn(pydf: PyDataFrame, colnm: &str, by_cols: Vec<&str>) -> PyResult<PyDataFrame> {
    let df: DataFrame = pydf.into();
    let res = df
        .lazy()
        .groupby(by_cols)
        .agg([col(colnm).apply(
            |s| {
                let v: Vec<f64> = vec![s.min().unwrap(), s.max().unwrap()];
                Ok(Some(Series::new("s", v)))
            },
            GetOutput::default(),
        )])
        .collect()
        .map_err(PyPolarsErr::from)?;
    Ok(PyDataFrame(res))
}

#[pymodule]
fn test_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(test_fn, m)?)?;
    Ok(())
}

执行Python调用代码后，返回的DataFrame中a列是包含两个元素的列表（对应每组的min和max值）：

import polars as pl
import test_module

df = pl.DataFrame(
    {"a": [1.0, 2.0, 3.0, 4.0, 5.0], "g1": [1, 1, 2, 2, 2], "g2": [1, 1, 1, 2, 2]}
)

>>> test_module.test_fn(df, "a", ["g1", "g2"])
shape: (3, 3)
┌─────┬─────┬────────────┐
│ g1  ┆ g2  ┆ a          │
│ --- ┆ --- ┆ ---        │
│ i64 ┆ i64 ┆ list[f64]  │
╞═════╪═════╪════════════╡
│ 1   ┆ 1   ┆ [1.0, 2.0] │
│ 2   ┆ 2   ┆ [4.0, 5.0] │
│ 2   ┆ 1   ┆ [3.0, 3.0] │
└─────┴─────┴────────────┘

需求：修改test_fn，使其返回结构体/字典而非向量，从而得到带命名字段的结果，期望两种输出格式之一。

方案一：返回Struct类型单列 #

利用Polars内置的struct!宏直接组合min和max值，无需使用性能较低的apply操作：

#[pyfunction]
fn test_fn(pydf: PyDataFrame, colnm: &str, by_cols: Vec<&str>) -> PyResult<PyDataFrame> {
    let df: DataFrame = pydf.into();
    let res = df
        .lazy()
        .groupby(by_cols)
        .agg([struct!(
            min = col(colnm).min(),
            max = col(colnm).max()
        ).alias(colnm)])
        .collect()
        .map_err(PyPolarsErr::from)?;
    Ok(PyDataFrame(res))
}

#[pymodule]
fn test_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(test_fn, m)?)?;
    Ok(())
}

执行后得到Struct类型的单列结果：

>>> test_module.test_fn(df, "a", ["g1", "g2"])
shape: (3, 3)
┌─────┬─────┬───────────┐
│ g1  ┆ g2  ┆ a         │
│ --- ┆ --- ┆ ---       │
│ i64 ┆ i64 ┆ struct[2] │
╞═════╪═════╪═══════════╡
│ 1   ┆ 1   ┆ {1.0,2.0} │
│ 2   ┆ 2   ┆ {4.0,5.0} │
│ 2   ┆ 1   ┆ {3.0,3.0} │
└─────┴─────┴───────────┘

方案二：拆分为独立的多列 #

直接在聚合阶段分别计算min和max并指定别名，性能最优且后续数据处理更便捷：

#[pyfunction]
fn test_fn(pydf: PyDataFrame, colnm: &str, by_cols: Vec<&str>) -> PyResult<PyDataFrame> {
    let df: DataFrame = pydf.into();
    let min_col = format!("{}_min", colnm);
    let max_col = format!("{}_max", colnm);
    
    let res = df
        .lazy()
        .groupby(by_cols)
        .agg([
            col(colnm).min().alias(&min_col),
            col(colnm).max().alias(&max_col)
        ])
        .collect()
        .map_err(PyPolarsErr::from)?;
    Ok(PyDataFrame(res))
}

#[pymodule]
fn test_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(test_fn, m)?)?;
    Ok(())
}

执行后得到拆分后的多列结果：

>>> test_module.test_fn(df, "a", ["g1", "g2"])
shape: (3, 4)
┌─────┬─────┬───────┬───────┐
│ g1  ┆ g2  ┆ a_min ┆ a_max │
│ --- ┆ --- ┆ ---   ┆ ---   │
│ i64 ┆ i64 ┆ f64   ┆ f64   │
╞═════╪═════╪═══════╪═══════╡
│ 2   ┆ 1   ┆ 3.0   ┆ 3.0   │
│ 2   ┆ 2   ┆ 4.0   ┆ 5.0   │
│ 1   ┆ 1   ┆ 1.0   ┆ 2.0   │
└─────┴─────┴───────┴───────┘

说明 #

• 优先使用Polars内置的向量化聚合函数，避免apply这类逐组标量操作，后者性能远低于前者。
• 方案二的可读性和后续数据处理便捷性更强，无需额外解析Struct结构。

内容的提问来源于stack exchange，提问作者lebesgue