在round函数中使用列值

Question

我有一个数据表：

shape: (5, 3)
┌───────────────┬──────────────┬────────────────────────┐
│ exchange_rate ┆ sig_figs_len ┆ reverse_rate_from_euro │
│ ---           ┆ ---          ┆ ---                    │
│ f64           ┆ u32          ┆ f64                    │
╞═══════════════╪══════════════╪════════════════════════╡
│ 6.4881        ┆ 5            ┆ 0.154128               │
│ 6.5196        ┆ 5            ┆ 0.153384               │
│ 6.4527        ┆ 5            ┆ 0.154974               │
│ 6.41          ┆ 3            ┆ 0.156006               │
│ 6.425         ┆ 4            ┆ 0.155642               │
└───────────────┴──────────────┴────────────────────────┘

我想把每个单元格里的reverse_rate_from_euro的值，按照对应单元格的sig_figs_len列里的值进行四舍五入。

我想出了一个解决办法，使用了apply/map_rows这个函数，但因为数据集比较大，使用apply或者纯Python的方法并不是最优的解决方案，所以我想找一个更好的办法。

这是我的代码片段：

df = df.with_columns(
    (df.map_rows(lambda df_: round(df_[-1], df_[-2])))
    .to_series()
    .alias("reverse_rate_to_euro_rounded_sig_figs")

有没有更好的方法，可以利用内置的Polars表达式API来实现？

结果集应该是这样的：

shape: (5, 4)
┌───────────────┬──────────────┬────────────────────────┬───────────────────────────────────┐
│ exchange_rate ┆ sig_figs_len ┆ reverse_rate_from_euro ┆ reverse_rate_to_euro_rounded_sig… │
│ ---           ┆ ---          ┆ ---                    ┆ ---                               │
│ f64           ┆ u32          ┆ f64                    ┆ f64                               │
╞═══════════════╪══════════════╪════════════════════════╪═══════════════════════════════════╡
│ 6.4881        ┆ 5            ┆ 0.154128               ┆ 0.15413                           │
│ 6.5196        ┆ 5            ┆ 0.153384               ┆ 0.15338                           │
│ 6.4527        ┆ 5            ┆ 0.154974               ┆ 0.15497                           │
│ 6.41          ┆ 3            ┆ 0.156006               ┆ 0.156                             │
│ 6.425         ┆ 4            ┆ 0.155642               ┆ 0.1556                            │
└───────────────┴──────────────┴────────────────────────┴───────────────────────────────────┘

任何建议都非常感谢！谢谢你的阅读！

Answer 1

你可以试试这样做：

df.with_columns(
    (
        pl.col('reverse_rate_from_euro') * 
        pl.lit(10).pow(pl.col('sig_figs_len'))
    ).round() * 
    pl.lit(0.1).pow(pl.col('sig_figs_len'))
)

┌───────────────┬──────────────┬────────────────────────┐
│ exchange_rate ┆ sig_figs_len ┆ reverse_rate_from_euro │
│ ---           ┆ ---          ┆ ---                    │
│ f64           ┆ i64          ┆ f64                    │
╞═══════════════╪══════════════╪════════════════════════╡
│ 6.4881        ┆ 3            ┆ 0.154                  │
│ 6.5196        ┆ 4            ┆ 0.1534                 │
└───────────────┴──────────────┴────────────────────────┘

另外，如果你知道不同的 sig_figs_len 值不多的话，你可以直接列举出 unique() 的值，然后用 coalesce() 和 when() 来创建结果：

df.with_columns(
    pl.coalesce(
        pl.when(pl.col("sig_figs_len") == x)
        .then(pl.col("reverse_rate_from_euro").round(x))
        for x in df['sig_figs_len'].unique()
    ).alias('reverse_rate_to_euro_rounded_sig_figs')
)

┌───────────────┬──────────────┬────────────────────────┬───────────────────────────────────┐
│ exchange_rate ┆ sig_figs_len ┆ reverse_rate_from_euro ┆ reverse_rate_to_euro_rounded_sig… │
│ ---           ┆ ---          ┆ ---                    ┆ ---                               │
│ f64           ┆ i64          ┆ f64                    ┆ f64                               │
╞═══════════════╪══════════════╪════════════════════════╪═══════════════════════════════════╡
│ 6.4881        ┆ 3            ┆ 0.154128               ┆ 0.154                             │
│ 6.5196        ┆ 4            ┆ 0.153384               ┆ 0.1534                            │
└───────────────┴──────────────┴────────────────────────┴───────────────────────────────────┘

其实，我们在这里做的就是创建一个 Iterable[IntoExpr]，你可以把它看作是一个列的列表，每一列只包含按照特定的 sig_figs_len 值四舍五入后的值：

df.with_columns(
    pl.when(pl.col("sig_figs_len") == x)
    .then(pl.col("reverse_rate_from_euro").round(x)).name.suffix(f"_{x}")
     for x in df['sig_figs_len'].unique()
))

┌───────────────┬──────────────┬─────────────────────┬─────────────────────┬─────────────────────┬─────────────────────┐
│ exchange_rate ┆ sig_figs_len ┆ reverse_rate_from_e ┆ reverse_rate_from_e ┆ reverse_rate_from_e ┆ reverse_rate_from_e │
│ ---           ┆ ---          ┆ uro                 ┆ uro_3               ┆ uro_4               ┆ uro_5               │
│ f64           ┆ i64          ┆ ---                 ┆ ---                 ┆ ---                 ┆ ---                 │
│               ┆              ┆ f64                 ┆ f64                 ┆ f64                 ┆ f64                 │
╞═══════════════╪══════════════╪═════════════════════╪═════════════════════╪═════════════════════╪═════════════════════╡
│ 6.4881        ┆ 5            ┆ 0.154128            ┆ null                ┆ null                ┆ 0.15413             │
│ 6.5196        ┆ 5            ┆ 0.153384            ┆ null                ┆ null                ┆ 0.15338             │
│ 6.4527        ┆ 5            ┆ 0.154974            ┆ null                ┆ null                ┆ 0.15497             │
│ 6.41          ┆ 3            ┆ 0.156006            ┆ 0.156               ┆ null                ┆ null                │
│ 6.425         ┆ 4            ┆ 0.155642            ┆ null                ┆ 0.1556              ┆ null                │
└───────────────┴──────────────┴─────────────────────┴─────────────────────┴─────────────────────┴─────────────────────┘

然后我们利用 pl.coalesce() 可以接受 Expr 的可迭代对象这一点，来获取第一个非空的列。