如何在Polars数据框链式操作中，以列作为Numpy随机函数输入新增列？

我有一个包含mean列的数据框，希望将该列作为随机数生成器的输入。在R语言中，借助dplyr管道可以轻松实现：

library(dplyr)

tibble(alpha = rnorm(1000),
       beta = rnorm(1000)) %>%
  mutate(mean = alpha + beta) %>%
  bind_cols(random_output = rnorm(n = nrow(.), mean = .$mean, sd = 1))
#> # A tibble: 1,000 × 4
#>     alpha   beta    mean random_output
#>     <dbl>  <dbl>   <dbl>         <dbl>
#>  1  0.231 -0.243 -0.0125         0.551
#>  2  0.213  0.647  0.861          0.668
#>  3  0.824 -0.353  0.471          0.852
#>  4  0.665 -0.916 -0.252         -1.81 
#>  5 -0.850  0.384 -0.465         -3.90 
#>  6  0.721  0.679  1.40           2.54 
#>  7  1.46   0.857  2.32           2.14 
#>  8 -0.242 -0.431 -0.673         -0.820
#>  9  0.234  0.188  0.422         -0.662
#> 10 -0.494 -2.15  -2.65          -3.01 
#> # ℹ 990 more rows

在Python中，我目前需要先创建中间对象，将其作为np.random.normal()的输入，再绑定到原数据框，这种方式较为繁琐。请问能否在Polars的链式操作/流水线中直接添加random_output列？

当前Python实现代码如下：

import polars as pl
import numpy as np

# 创建数据框
df = (
    pl.DataFrame(
        {
            "alpha": np.random.standard_normal(1000),
            "beta": np.random.standard_normal(1000)
        }
    )
    .with_columns(
        (pl.col("alpha") + pl.col("beta")).alias("mean")
    )
    
)

# 创建中间对象
sim_vals = np.random.normal(df.get_column("mean"))

# 将模拟值绑定到原数据框
(
    df.with_columns(random_output = pl.lit(sim_vals))
)
#> shape: (1_000, 4)
┌───────────┬───────────┬───────────┬───────────────┐
│ alpha     ┆ beta      ┆ mean      ┆ random_output │
│ ---       ┆ ---       ┆ ---       ┆ ---           │
│ f64       ┆ f64       ┆ f64       ┆ f64           │
╞═══════════╪═══════════╪═══════════╪═══════════════╡
│ -1.380249 ┆ 1.531959  ┆ 0.15171   ┆ 0.938207      │
│ -0.332023 ┆ -0.108255 ┆ -0.440277 ┆ 0.081628      │
│ -0.718319 ┆ -0.612187 ┆ -1.330506 ┆ -1.286229     │
│ 0.22067   ┆ -0.497258 ┆ -0.276588 ┆ 0.908147      │
│ …         ┆ …         ┆ …         ┆ …             │
│ 0.299117  ┆ -0.371846 ┆ -0.072729 ┆ 0.592632      │
│ 0.789633  ┆ 0.95712   ┆ 1.746753  ┆ 2.954801      │
│ -0.264415 ┆ -0.761634 ┆ -1.026049 ┆ -1.369753     │
│ 1.893911  ┆ 1.554736  ┆ 3.448647  ┆ 5.192537      │
└───────────┴───────────┴───────────┴───────────────┘

优化后的Polars链式实现 #

完全可以在Polars的链式操作中直接完成，利用Polars内置的pl.random.normal()函数，它支持将列表达式作为mean参数传入，直接生成对应每行均值的随机数：

import polars as pl
import numpy as np

df = (
    pl.DataFrame(
        {
            "alpha": np.random.standard_normal(1000),
            "beta": np.random.standard_normal(1000)
        }
    )
    .with_columns(mean = pl.col("alpha") + pl.col("beta"))
    .with_columns(random_output = pl.random.normal(mean=pl.col("mean"), std=1))
)

print(df.head())

执行后会直接得到包含random_output列的完整数据框，无需任何中间对象，实现和R语言dplyr管道一样简洁的链式操作。

内容的提问来源于stack exchange，提问作者Mark Rieke