有没有办法应用一个UDF函数,在Rust极坐标中返回多个值?

3ks5zfa0  于 2023-03-08  发布在  其他
关注(0)|答案(2)|浏览(127)

我尝试使用polars将另一个库中的函数应用到输入的每一行,但找不到任何使用Expr应用函数的示例或测试,即使它只有一个返回值;所以我迷路了。
它接受一个包含两个浮点列的输入 Dataframe ,并尝试追加三个由以下形式的函数生成的列:

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

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

56lgkhnf

56lgkhnf1#

这里有不同的策略。你可以把返回值赋给不同的列。或者你可以把返回值赋给List<Float64>类型的单个列。我将展示这两种策略。

不同色谱柱

将它们分配到不同的列并不真正适合lazy API,因此我们在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(())
}
eqoofvh9

eqoofvh92#

列表列(使用ChunkedArray):

fn to_list() -> Result<(), Box<dyn Error>> {
    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(Some(
                         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::<ChunkedArray<ListType>>()
                    .into_series()))
            },
            [col("a"), col("b")],
            GetOutput::from_type(DataType::Float64),
        ).alias("new column")
        ])
        .collect()?;

    dbg!(df);

    Ok(())
}

/// Your function that takes 2 argument and returns 3
fn black_box(a: f64, b: f64) -> (f64, f64, f64) {
    (a+b, 5.4 * a - 2.1 * b, a*b)
}

to_list()?;

输出:

df = shape: (3, 1)
┌─────────────────┐
│ new column      │
│ ---             │
│ list[f64]       │
╞═════════════════╡
│ [2.0, 3.3, 1.0] │
│ [4.0, 6.6, 4.0] │
│ [6.0, 9.9, 9.0] │
└─────────────────┘

相关问题