时间戳

首页 > 解决方案 > 如何在 rust 中向极坐标数据框添加条件计算 col？

rust - 如何在 rust 中向极坐标数据框添加条件计算 col？

问题描述

df1 有['a', 'b', 'c']3 个 cols，我想得到一个 df2 有 4 个 cols 的['a', 'b', 'c', 'd']. d 是这样计算的：

if a>5 {
  d = b + c 
} else if a<-5 {
  d = c - b + a
}  else {
  d = 3.0 * a
}

我怎么能用生锈的极性来做到这一点？也许既渴望又懒惰。

标签： rustrust-polars

解决方案

您可以使用when -> then -> when -> then -> otherwise表达式。请注意，您可以when, then无限扩展，就像else if分支一样。

下面是一个例子：

use polars::df;
use polars::prelude::*;

fn main() -> Result<()> {
    let df = df![
        "a" => [2, 9, 2, 5],
        "b" => [1, 2, 3, 4],
        "c" => [4, 4, 8, 4],
    ]?;

    let out = df
        .lazy()
        .select([
            col("*"),
            when(col("a").gt(lit(5)))
                .then(col("b") + col("c"))
                .when(col("a").lt(lit(5)))
                .then(col("c") - col("b") + col("a"))
                .otherwise(lit(3) * col("a"))
                .alias("d"),
        ])
        .collect()?;

    println!("{}", out);

    Ok(())
}

这输出：

shape: (4, 4)
┌─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   │
│ --- ┆ --- ┆ --- ┆ --- │
│ i32 ┆ i32 ┆ i32 ┆ i32 │
╞═════╪═════╪═════╪═════╡
│ 2   ┆ 1   ┆ 4   ┆ 5   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 9   ┆ 2   ┆ 4   ┆ 6   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2   ┆ 3   ┆ 8   ┆ 7   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 5   ┆ 4   ┆ 4   ┆ 15  │
└─────┴─────┴─────┴─────┘

推荐阅读