如何在polars数据帧列中获取一个项，并将其放回不同位置的同一列中

Question

对polar和rust来说仍然是新事物。。。这里有一个小问题：

如何在特定位置访问DataFrame中的值。

如何在特定位置覆盖DataFrame中的值。

这是一个非工作代码：

use polars::prelude::*;
fn main() {
let df = df! [
"STOCK"         => ["TSLA", "META", "AA",],
"STRIKES"       => [10, 20, 5],
]
.unwrap();
println!("dft{:?}", df);
// Take TSLA's STRIKE (10)
let tsla_strike = df
.lazy()
.filter((col("STOCK") == lit("TSLA")))
.with_column(col("STRIKES"))
.first()
.collect();
let o_i32 = GetOutput::from_type(DataType::Int32);
// Overwrite AA's STRIKE with tsla_strike (5 ==> 10)
let df = df
.lazy()
.filter((col("STOCK") == lit("AA")).into())
.with_column(col("STRIKES").map(|x| tsla_strike,o_i32))
.collect()
.unwrap();
println!("dft{:?}", df);
}

这是我喜欢得到的结果：

RESULT:
df      shape: (3, 2)
┌───────┬─────────┐
│ STOCK ┆ STRIKES │
│ ---   ┆ ---     │
│ str   ┆ i32     │
╞═══════╪═════════╡
│ TSLA  ┆ 10      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ META  ┆ 20      │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ AA    ┆ 10      │
└───────┴─────────┘

一种反模式的方法是遍历DF，然后同时构建具有所需值的新DF。

Answer 1

您可以使用when -> then -> otherwise构造。当STOCK=="AA"取STRIKE，其中取STOCK=="TSLA"，否则取STRIKE。此构造是向量化的且快速的(它不对单个元素进行操作(。

let df2 = df
.lazy()
.clone()
.select([
col("STOCK"),
when(col("STOCK").eq(lit("AA")))
.then(col("STRIKES").filter(col("STOCK").eq(lit("TSLA"))))
.otherwise(col("STRIKES"))
])
.collect()?;

如果您有很多映射要做，另一个选项是映射数据帧，并左键连接替换值。

let mapping = df! [
"ORIGINAL_STOCK"    => ["TSLA", "AA"],
"REPLACEMENT_STOCK" => ["AA", "META"]
]?;
let df2 = df
.clone()
.lazy()
.join(mapping.clone().lazy(), [col("STOCK")], [col("ORIGINAL_STOCK")], JoinType::Left)
.join(df.clone().lazy(), [col("REPLACEMENT_STOCK")], [col("STOCK")], JoinType::Left)
.select([
col("STOCK"),
when(col("STRIKES_right").is_not_null())
.then(col("STRIKES_right"))
.otherwise(col("STRIKES"))
.alias("STRIKES")
])
.collect()?;