如何在Polars中实现类似numpy.random.choice的行级概率采样

在NumPy中，我们可以通过np.random.choice指定概率分布采样单个值：

import numpy as np

# 可选结果
values = [0, 1, 2, 3, 4]
# 对应概率
probabilities = [0.15, 0.30, 0.25, 0.20, 0.10]

# 按概率采样
sampled_value = np.random.choice(values, p=probabilities)
print(sampled_value)

但如果是Polars DataFrame，每行都有对应不同结果的概率列（如下示例），需要为每行生成一个按该行概率采样的结果列sampled_goal，可以用以下两种方式实现：

示例DataFrame #

import polars as pl

df = pl.DataFrame({
    "id": [1, 2, 3],
    "prob_0": [0.1, 0.2, 0.15],
    "prob_1": [0.3, 0.25, 0.35],
    "prob_2": [0.25, 0.3, 0.25],
    "prob_3": [0.2, 0.15, 0.15],
    "prob_4": [0.15, 0.1, 0.1]
})

方法1：逐行处理（简单直观） #

通过pl.struct打包每行的概率列，再用map_elements调用np.random.choice逐行采样：

def sample_row(probs):
    return np.random.choice([0,1,2,3,4], p=list(probs.values()))

df = df.with_columns(
    sampled_goal=pl.struct(pl.col("prob_*")).map_elements(sample_row, return_dtype=pl.Int64)
)

print(df)

注意：这种方法逻辑简单，但map_elements是逐行执行，大数据集下效率较低。

方法2：向量化处理（高效） #

利用Polars的向量化操作，结合累积概率和随机数实现高效采样：

values = [0,1,2,3,4]
prob_cols = [f"prob_{v}" for v in values]

df = df.with_columns(
    # 为每行生成一个0-1之间的随机数
    rand=pl.float_rand(seed=42),
    # 将概率列转为列表并计算累积和
    cum_probs=pl.concat_list(prob_cols).list.cum_sum(),
    # 生成对应结果值的列表（每行重复一次）
    value_list=pl.lit(values).repeat(pl.count())
).with_columns(
    # 找到第一个累积和大于随机数的索引，取对应结果值
    sampled_goal=pl.col("cum_probs")
    .list.eval(pl.element() > pl.col("rand"))
    .list.arg_true().first()
    .list.get(pl.col("value_list"))
).drop(["rand", "cum_probs", "value_list"])  # 删除临时列

print(df)

这种方法大部分操作是向量化执行，适合处理大规模数据集。

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