rust-polars - 有没有办法应用在 Rust 极坐标中返回多个值的 UDF 函数？

Question

我正在尝试使用极坐标将另一个库中的函数应用于输入的每一行。我找不到任何使用 Expr 应用函数的示例或测试，即使它有一个返回值；所以我迷路了。

它采用具有两个浮点列的输入数据框，并尝试附加由具有此形式的函数生成的三列：

fn f(a: f64, b: f64) -> (f64, f64, f64);

有没有一种简单的方法可以做到这一点？

Answer 1

这里有不同的策略。您可以将返回值分配给不同的列。或者，您可以将返回值分配给 type 的单个列List<Float64>。我会向他们展示。

不同的列

将它们分配给不同的列并不适合lazyAPI，所以我们在eager.

/// Your function that takes 2 argument and returns 3
fn black_box(_a: f64, _b: f64) -> (f64, f64, f64) {
    (1.0, 2.0, 3.0)
}

fn to_different_columns() -> Result<()> {
    let df = df![
        "a" => [1.0, 2.0, 3.0],
        "b" => [1.0, 2.0, 3.0]
    ]?;

    let mut out_1 = vec![];
    let mut out_2 = vec![];
    let mut out_3 = vec![];

    df.column("a")?
        .f64()?
        .into_no_null_iter()
        .zip(df.column("b")?.f64()?.into_no_null_iter())
        .for_each(|(a, b)| {
            let (out_val1, out_val2, out_val3) = black_box(a, b);
            out_1.push(out_val1);
            out_2.push(out_val2);
            out_3.push(out_val3);
        });

    let out1 = Series::from_vec("out1", out_1);
    let out2 = Series::from_vec("out2", out_2);
    let out3 = Series::from_vec("out3", out_3);
    let df = DataFrame::new(vec![out1, out2, out3]);

    Ok(())
}

列表栏

如果我们决定退回单曲Series，我们最好使用polars lazy

fn to_list() -> Result<()> {
    let df = df![
        "a" => [1.0, 2.0, 3.0],
        "b" => [1.0, 2.0, 3.0]
    ]?;

    let df = df
        .lazy()
        .select([map_multiple(
            |columns| {
                Ok(columns[0]
                    .f64()?
                    .into_no_null_iter()
                    .zip(columns[1].f64()?.into_no_null_iter())
                    .map(|(a, b)| {
                        let out = black_box(a, b);
                        Series::new("", [out.0, out.1, out.2])
                    })
                    .collect::<ListChunked>()
                    .into_series())
            },
            [col("a"), col("b")],
            GetOutput::from_type(DataType::List(Box::new(DataType::Float64))),
        )])
        .collect()?;

    dbg!(df);

    Ok(())
}