如何确保Polars创建的是列表类型而不是对象类型的列

Question

下面的代码将创建一个名为paid 的列，它看起来像一个列表，但它是一个对象，因此实际上作为列毫无用处。如何确保创建的列是列表列而不是对象列，因为创建后

.cast()

无法应用于对象列。

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()

Answer 1

在您的

.map_elements

调用中，让

numpy

将值作为 Python

list

返回，而不是作为

np.array

。在之前的返回类型中，Polars 0.20.25 可以正确投射到

list[i64]

。

>>> 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')
... )
>>> df_cluster.head()
shape: (5, 4)
┌────────────┬──────────┬────────┬─────────────┐
│ cluster_id ┆ mu       ┆ trials ┆ paid        │
│ ---        ┆ ---      ┆ ---    ┆ ---         │
│ i64        ┆ f64      ┆ i32    ┆ list[i64]   │
╞════════════╪══════════╪════════╪═════════════╡
│ 1          ┆ 0.52546  ┆ 14     ┆ [1, 0, … 1] │
│ 2          ┆ 0.584446 ┆ 22     ┆ [1, 0, … 0] │
│ 3          ┆ 0.590928 ┆ 16     ┆ [0, 1, … 0] │
│ 4          ┆ 0.506396 ┆ 19     ┆ [1, 1, … 0] │
│ 5          ┆ 0.564219 ┆ 29     ┆ [0, 0, … 0] │
└────────────┴──────────┴────────┴─────────────┘