rust 谁能提供一个与polars相关的例子as_struct().apply()

bq8i3lrv  于 2024-01-08  发布在  其他
关注(0)|答案(1)|浏览(142)

我尝试将数组列添加到现有的数据框中。
| 1 | 3 |
| 4 | 4 |
和输出
| 1 | 3 |【??】|
| 4 |【??】| [?, ?, ?, ?] |
数组的值将由一些自定义函数填充。
我试图实现这样的东西。

let df = df![
  "a" => [1, 2],
  "b" => [3, 4]
]?;
        
let lf: LazyFrame = df.lazy().select(
  as_struct(&[col("a"), col("b")]).apply(
    somefn,
    GetOutput::from_type(DataType::List(Box::new(DataType::Float32))),
));

字符串
我不知道如何实现这个somefn

mqkwyuun

mqkwyuun1#

用户指南中有一些为apply函数编写自定义函数的示例,如下所示:https://pola-rs.github.io/polars/user-guide/expressions/user-defined-functions/
忽略标题暗示它只适用于Python的事实,也有Rust示例。
从该页面:

let out = df
    .lazy()
    .select([
        // pack to struct to get access to multiple fields in a custom `apply/map`
        as_struct(vec![col("keys"), col("values")])
            // we will compute the len(a) + b
            .apply(
                |s| {
                    // downcast to struct
                    let ca = s.struct_()?;

                    // get the fields as Series
                    let s_a = &ca.fields()[0];
                    let s_b = &ca.fields()[1];

                    // downcast the `Series` to their known type
                    let ca_a = s_a.utf8()?;
                    let ca_b = s_b.i32()?;

                    // iterate both `ChunkedArrays`
                    let out: Int32Chunked = ca_a
                        .into_iter()
                        .zip(ca_b)
                        .map(|(opt_a, opt_b)| match (opt_a, opt_b) {
                            (Some(a), Some(b)) => Some(a.len() as i32 + b),
                            _ => None,
                        })
                        .collect();

                    Ok(Some(out.into_series()))
                },
                GetOutput::from_type(DataType::Int32),
            )
            .alias("solution_apply"),
        (col("keys").str().count_matches(lit("."), true) + col("values"))
            .alias("solution_expr"),
    ])
    .collect()?;
println!("{}", out);

字符串

相关问题