在Python-Polars中处理列类型List[int]

Question

很多时候我发现自己处于这样的情况：我有一个 DataFrame 并且一列的类型为 List[int]。

例如，我有以下 DF：

df = pl.DataFrame(
    {"group": ["A", "A", "B", "B", "B", "B"], 
     "value": [[3, 2, 5], [2,2,2], [2,5,9,4], [5,4,7,5,1], [9,4,5], [2,2]]}
)

通常，我在这种情况下使用

explode

和

group_by

方法。
然而，当处理大量列时，代码可能会变得有些“脏”。

为了解决这个问题，我想使用map_elements方法：

(
    df
    .group_by('group')
    .agg(
        (pl.col('value').map_elements(lambda l: pl.concat(l)))
        )
    .with_columns(
        pl.col('value').map_elements(lambda l: pl.Series.median(l))
    )
)

不幸的是，这种方法牺牲了 Polars 提供的并行化优势。 而且它的执行非常耗费资源。如果我有数百万行，执行时间可能会从几秒延长到几分钟。

有更好的方法来使用 List[int] 吗？ 有没有好的方法来优化我的代码？

Answer 1

有一个 explode 表达式，也可以通过

.flatten()

别名使用。

(df.group_by('group')
   .agg(pl.col('value').flatten().median())
)

shape: (2, 2)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ f64   │
╞═══════╪═══════╡
│ B     ┆ 4.5   │
│ A     ┆ 2.0   │
└───────┴───────┘