如何确保Polars创建list类型列而非object类型列

现有代码创建的paid列外观类似列表，但实际为object类型，无法通过.cast()方法转为列表列。如何确保创建的paid列为Polars的列表（List）类型？

原代码 #

import numpy as np
import polars as pl
import scipy.stats as 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"])
    .map_elements(lambda x: np.random.binomial(n=1, p=x['mu'], size=x['trials']))
    .alias('paid')
)

df_cluster.head()

解决方案 #

方法1：显式指定map_elements的返回类型 #

Polars的map_elements默认会将numpy数组识别为object类型，只需显式声明返回类型为列表，就能直接生成List列：

df_cluster = df_cluster.with_columns(
    pl.struct(["mu", "trials"])
    .map_elements(
        lambda x: np.random.binomial(n=1, p=x['mu'], size=x['trials']),
        return_dtype=pl.List(pl.Int32)  # 指定返回为整数列表
    )
    .alias('paid')
)

方法2：将numpy数组转为Python原生列表 #

在lambda函数中，把生成的numpy数组通过.tolist()转为Python列表，Polars会自动识别为List类型：

df_cluster = df_cluster.with_columns(
    pl.struct(["mu", "trials"])
    .map_elements(lambda x: np.random.binomial(n=1, p=x['mu'], size=x['trials']).tolist())
    .alias('paid')
)

验证结果 #

执行后可通过print(df_cluster.schema)查看列类型，paid会显示为List(Int32)，此时能正常使用Polars的列表操作（如.list.lengths()、.list.sum()等）。

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