Polars列表列转换与展开技巧

本教程详细介绍了如何利用Polars高效地将包含列表的宽格式DataFrame转换为长格式，并同时将列表元素展开为独立的列。通过结合`unpivot`、`list.to_struct`和`unnest`等核心操作，读者将学会如何优雅地重塑数据，实现从原始的列名-列表值结构到Name-Value0/Value1/Value2等新列的转换，从而简化复杂的数据清洗和分析任务。

Polars中复杂列表列的重塑与展开

在数据处理中，我们经常会遇到包含列表类型数据的DataFrame，并且需要将其从宽格式转换为长格式，同时将列表中的元素展开成独立的列。Polars作为一款高性能的DataFrame库，提供了强大且富有表达力的API来应对这类挑战。本文将详细介绍如何通过一系列链式操作，实现对Polars DataFrame中列表列的特定方式转换。

初始数据结构与目标

假设我们有一个Polars DataFrame，其结构如下，其中"foo"和"bar"列都包含整数列表：

import polars as pl

df = pl.DataFrame({
    "foo": [[1, 2, 3], [7, 8, 9]],
    "bar": [[4, 5, 6], [1, 0, 1]]
})

print(df)
# 输出:
# shape: (2, 2)
# ┌───────────┬───────────┐
# │ foo       ┆ bar       │
# │ ---       ┆ ---       │
# │ list[i64] ┆ list[i64] │
# ╞═══════════╪═══════════╡
# │ [1, 2, 3] ┆ [4, 5, 6] │
# │ [7, 8, 9] ┆ [1, 0, 1] │
# └───────────┴───────────┘

我们的目标是将其转换为以下结构：

shape: (4, 4)
┌──────┬────────┬────────┬────────┐
│ Name ┆ Value0 ┆ Value1 ┆ Value2 │
│ ---  ┆ ---    ┆ ---    │ ---    │
│ str  ┆ i64    ┆ i64    ┆ i64    │
╞══════╪════════╪════════╪════════╡
│ foo  ┆ 1      ┆ 2      ┆ 3      │
│ foo  ┆ 7      ┆ 8      ┆ 9      │
│ bar  ┆ 4      ┆ 5      ┆ 6      │
│ bar  ┆ 0      ┆ 1      ┆ 1      │
└──────┴────────┴────────┴────────┘

这个转换涉及到两个主要步骤：首先，将列名（"foo", "bar"）转换为一个名为"Name"的新列，并将它们对应的值放入一个名为"value"的新列中，即从宽格式到长格式的转换。其次，将"value"列中的列表元素展开为独立的列（"Value0", "Value1", "Value2"）。

转换步骤详解

Polars提供了一系列强大的表达式和方法来高效地完成这个转换。我们将分步进行说明。

步骤一：使用 unpivot 进行列名转换与长格式化

unpivot 方法用于将DataFrame的指定列（或所有非id_vars列）从宽格式转换为长格式。它会将列名转换为一个新列的行值，并将这些列的原始值放入另一个新列中。

在这个例子中，我们没有指定id_vars，所以所有列（"foo", "bar"）都将被unpivot。variable_name="Name"参数指定了存储原始列名的新列的名称。默认情况下，值列的名称是value。

# 初始DataFrame
# df = pl.DataFrame({
#     "foo": [[1, 2, 3], [7, 8, 9]],
#     "bar": [[4, 5, 6], [1, 0, 1]]
# })

df_unpivoted = df.unpivot(variable_name="Name")
print(df_unpivoted)
# 输出:
# shape: (4, 2)
# ┌──────┬───────────┐
# │ Name ┆ value     │
# │ ---  ┆ ---       │
# │ str  ┆ list[i64] │
# ╞══════╪═══════════╡
# │ foo  ┆ [1, 2, 3] │
# │ bar  ┆ [4, 5, 6] │
# │ foo  ┆ [7, 8, 9] │
# │ bar  ┆ [1, 0, 1] │
# └──────┴───────────┘

现在，我们有了一个Name列（包含原始列名）和一个value列（包含原始列表数据）。

步骤二：使用 list.to_struct 将列表转换为结构体

在Polars中，要将列表的元素展开为单独的列，通常需要先将列表转换为结构体（Struct）。list.to_struct() 方法正是为此目的设计的。它将列表中的每个元素映射到结构体的一个字段。

fields 参数允许我们自定义结构体中字段的名称。在这里，我们使用一个lambda函数 lambda x: f"Value{x}" 来为每个列表元素生成字段名，例如"Value0", "Value1", "Value2"等。

df_struct = df_unpivoted.with_columns(
    pl.col("value").list.to_struct(fields=lambda x: f"Value{x}")
)
print(df_struct)
# 输出:
# shape: (4, 2)
# ┌──────┬───────────────────────────┐
# │ Name ┆ value                     │
# │ ---  ┆ ---                       │
# │ str  ┆ struct[3]                 │
# ╞══════╪═══════════════════════════╡
# │ foo  ┆ {1,2,3}                   │
# │ bar  ┆ {4,5,6}                   │
# │ foo  ┆ {7,8,9}                   │
# │ bar  ┆ {1,0,1}                   │
# └──────┴───────────────────────────┘

现在，value列已经从list[i64]类型变为了struct[3]类型，其中包含了三个字段，默认名称为field_0, field_1, field_2（或者根据我们lambda函数生成的名称）。

步骤三：使用 unnest 展开结构体列

最后一步是使用 unnest 方法将结构体列展开为独立的列。只需指定要unnest的列名，Polars就会自动将其内部的字段提升为DataFrame的顶层列。

df_final = df_struct.unnest("value")
print(df_final)
# 输出:
# shape: (4, 4)
# ┌──────┬────────┬────────┬────────┐
# │ Name ┆ Value0 ┆ Value1 ┆ Value2 │
# │ ---  ┆ ---    ┆ ---    ┆ ---    │
# │ str  ┆ i64    ┆ i64    ┆ i64    │
# ╞══════╪════════╪════════╪════════╡
# │ foo  ┆ 1      ┆ 2      ┆ 3      │
# │ bar  ┆ 4      ┆ 5      ┆ 6      │
# │ foo  ┆ 7      ┆ 8      ┆ 9      │
# │ bar  ┆ 1      ┆ 0      ┆ 1      │
# └──────┴────────┴────────┴────────┘

至此，我们成功地将原始DataFrame转换成了目标格式。

完整代码示例

将上述三个步骤链式组合起来，我们可以得到一个简洁高效的解决方案：

import polars as pl

# 原始DataFrame
df = pl.DataFrame({
    "foo": [[1, 2, 3], [7, 8, 9]],
    "bar": [[4, 5, 6], [1, 0, 1]]
})

# 链式操作实现转换
output_df = (
    df
    .unpivot(variable_name="Name")
    .with_columns(pl.col("value").list.to_struct(fields=lambda x: f"Value{x}"))
    .unnest("value")
)

print(output_df)

注意事项与总结

• 性能优势： Polars的表达式系统和惰性计算（当使用scan_csv等时）使得这些复杂的数据转换操作在内存和CPU效率上都表现出色。
• 灵活性： fields参数在list.to_struct中非常灵活，可以根据具体需求动态生成列名。如果列表长度不固定，list.to_struct会填充null值以保持结构体的一致性。
• 错误处理： 确保列表中的元素类型一致，否则在转换为结构体时可能会遇到类型不匹配的问题。
• 可读性： 链式操作虽然强大，但在处理非常复杂的转换时，适当拆分为多个步骤或添加注释可以提高代码的可读性。

通过掌握unpivot、list.to_struct和unnest这三个关键操作，您将能够高效地在Polars中处理和重塑包含列表的复杂DataFrame，从而解锁更多数据分析的可能性。