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

Question

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

age_pre

的列和一个名为

age_post

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

age

的列。与

exposure_pre

和

exposure_post

列相同。

df = pl.from_repr("""
┌─────┬─────────┬──────────┬──────────────┬───────────────┐
│ id  ┆ age_pre ┆ age_post ┆ exposure_pre ┆ exposure_post │
│ --- ┆ ---     ┆ ---      ┆ ---          ┆ ---           │
│ i64 ┆ i64     ┆ i64      ┆ f64          ┆ f64           │
╞═════╪═════════╪══════════╪══════════════╪═══════════════╡
│ 123 ┆ 55      ┆ 56       ┆ 0.49         ┆ 0.51          │
│ 123 ┆ 56      ┆ 57       ┆ 0.49         ┆ 0.51          │
│ 456 ┆ 49      ┆ 50       ┆ 0.8          ┆ 0.2           │
│ 456 ┆ 50      ┆ 51       ┆ 0.8          ┆ 0.2           │
└─────┴─────────┴──────────┴──────────────┴───────────────┘
""")

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

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      │
└─────┴─────┴──────────┘

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

.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 -> 连接所有给定的 DataFrame。这是一种 DataFrame 的链接列表。如果您传递 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      │
└─────┴─────┴──────────┘