如何重塑/融合数据集中的多个列并将结果值拆分为 Polars 中不同数量的列？

Question

我正在尝试通过将多对列的值组合成每对的单个列来重组数据框。例如，我有一个名为

age_pre

的列和一个名为

age_post

的列，我想将这两列合并为一个名为

age

的列。与

exposure_pre

和

exposure_post

列相同。

id	age_pre	age_post	曝光前	曝光帖子	...
123	55	56	0.49	0.51	...
123	56	57	0.49	0.51	...
456	49	50	0.80	0.20	...
456	50	51	0.80	0.20	...

我想要的输出看起来像这样

id	年龄	曝光
123	55	0.49
123	56	0.51
123	56	0.49
123	57	0.51
456	49	0.8
456	50	0.2
456	50	0.8
456	51	0.2

我尝试创建多个融化的数据框，并使用

.hstack()

创建一个具有我想要的输出的数据框，这看起来效率很低。我如何使用 Polars 语法来做到这一点？

Answer 1

连接数据帧可以用

concat

、

vstack

或

extend

获得。 vstack 是一种非常便宜的连接方式。

下面是一个使用

vstack

的例子：

import polars as pl

df = pl.DataFrame({
    'id' : [123, 123, 456, 456],
    'age_pre' : [55, 56, 49, 50],
    'age_post' : [56, 57, 50, 51],
    'exposure_pre' : [0.49, 0.49, 0.80, 0.80],
    'exposure_post' : [0.51, 0.51, 0.20, 0.20]
})

column_rename_pre = {column.name : column.name.replace('_pre','') 
    for column in df if 'pre' in column.name}
column_rename_post = {column.name : column.name.replace('_post','') 
    for column in df if 'post' in column.name}

out = (
    df.select(pl.col('^(id|.*_pre)$')).rename(column_rename_pre)
    .vstack(df.select(pl.col('^(id|.*_post)$')).rename(column_rename_post))
)

# result

shape: (8, 3)
┌─────┬─────┬──────────┐
│ id  ┆ age ┆ exposure │
│ --- ┆ --- ┆ ---      │
│ i64 ┆ i64 ┆ f64      │
╞═════╪═════╪══════════╡
│ 123 ┆ 55  ┆ 0.49     │
│ 123 ┆ 56  ┆ 0.49     │
│ 456 ┆ 49  ┆ 0.8      │
│ 456 ┆ 50  ┆ 0.8      │
│ 123 ┆ 56  ┆ 0.51     │
│ 123 ┆ 57  ┆ 0.51     │
│ 456 ┆ 50  ┆ 0.2      │
│ 456 ┆ 51  ┆ 0.2      │
└─────┴─────┴──────────┘

有一个很好的 Stack Overflow 答案解释了 concat、vstack 和 extend 之间的区别：How can I append or concatenate two dataframes in python polars?

我在下面复制该答案的建议

追加数据有以下几种方式：

concat -> 连接所有给定的数据帧。这是一种数据帧的链表。如果你传递 rechunk=True，所有内存将被重新分配给连续的块。

vstack -> 通过递增引用计数将其他数据添加到 DataFrame。这是超级便宜。建议在多次 vstack 之后调用 rechunk。或者简单地使用 pl.concat.

extend 该操作复制数据。它尝试将数据从其他数据复制到 DataFrame。但是，如果 DataFrame 的引用计数大于 1。将分配一个新的内存缓冲区来保存两个 DataFrame。

Answer 2

您可以使用

pl.concat_list

然后

.explode

.

out = df.with_columns(
    pl.concat_list(["age_pre", "age_post"]).alias("age"),
    pl.concat_list(["exposure_pre", "exposure_post"]).alias("exposure"),
).select(['id', 'age', 'exposure']).explode(['age', 'exposure'])
    
print(out)
shape: (8, 3)
┌─────┬─────┬──────────┐
│ id  ┆ age ┆ exposure │
│ --- ┆ --- ┆ ---      │
│ i64 ┆ i64 ┆ f64      │
╞═════╪═════╪══════════╡
│ 123 ┆ 55  ┆ 0.49     │
│ 123 ┆ 56  ┆ 0.51     │
│ 123 ┆ 56  ┆ 0.49     │
│ 123 ┆ 57  ┆ 0.51     │
│ 456 ┆ 49  ┆ 0.8      │
│ 456 ┆ 50  ┆ 0.2      │
│ 456 ┆ 50  ┆ 0.8      │
│ 456 ┆ 51  ┆ 0.2      │
└─────┴─────┴──────────┘