如何在Polars中用np.random.binomial生成变长列表列？

问题原因 #

使用pl.struct().apply()生成变长数组时，Polars无法自动推断返回值类型（默认期望标量），导致无法正确生成列表列。需要显式指定返回类型，或使用Polars原生向量化操作实现更高效的处理。

方法一：修复apply写法（显式指定返回类型） #

在apply中添加return_dtype参数，明确告知Polars返回整数列表类型：

import numpy as np
import polars as pl
from scipy import stats

CLUSTERS = 200 
MEAN_TRIALS = 20
MU = 0.5
SIGMA = 0.1

# 初始数据构造
df_cluster = pl.DataFrame({'cluster_id': range(1, CLUSTERS+1)}) 
df_cluster = df_cluster.with_columns(
    mu = stats.truncnorm(a=0, b=1, loc=MU, scale=SIGMA).rvs(size=CLUSTERS),
    trials = np.random.poisson(lam=MEAN_TRIALS, size=CLUSTERS)
)

# 生成变长列表列
df_cluster = df_cluster.with_columns(
    pl.struct(["mu", "trials"])
    .apply(lambda x: np.random.binomial(n=1, p=x['mu'], size=x['trials']), return_dtype=pl.List(pl.Int32))
    .alias('paid')
)

方法二：Polars原生向量化操作（推荐，性能更优） #

避免逐行apply，使用Polars内置的list.generate实现完全向量化的列表生成，处理大数据量时效率更高：

df_cluster = df_cluster.with_columns(
    pl.list.generate(
        lambda _: pl.int_rvs(0, 1, p=pl.col("mu")),
        length=pl.col("trials")
    ).alias("paid")
)

这里pl.int_rvs(0, 1, p=mu)会根据每个cluster的mu值生成0或1，list.generate按trials指定的长度生成对应列表。

验证结果 #

执行后通过print(df_cluster.head())查看结果，paid列将是每个cluster_id对应的、长度等于trials的0-1序列列表。

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