Polars列表转宽格式实践指南

本文旨在指导用户如何在Polars中高效地将包含列表的DataFrame列转换为更规整的宽格式。通过结合`unpivot`、`list.to_struct`和`unnest`等核心操作，我们能够将原始列名转换为新的标识列，并将列表中的元素展开为独立的数值列，从而实现复杂数据的结构化重塑，提升数据分析的便利性与效率。

在数据处理和分析中，我们经常会遇到需要对DataFrame的列进行复杂重塑的情况，特别是当列中包含列表（List）类型的数据时。本教程将详细介绍如何利用Polars的强大功能，将一个包含列表列的DataFrame转换为一种更易于分析的宽格式，其中原始列名被转换为新的标识列，而列表中的元素则被展开为独立的数值列。

1. 问题场景与初始数据

假设我们有一个Polars DataFrame，其中包含多列，每列的值都是一个整数列表。我们的目标是将其转换为一种格式，使得：

• 原始的列名（如 "foo", "bar"）成为一个名为 "Name" 的新列的值。
• 每个列表中的元素被展开为新的列，例如 "Value0", "Value1", "Value2" 等。

以下是我们的初始DataFrame示例：

import polars as pl

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

print("原始 DataFrame:")
print(df)

输出：

原始 DataFrame:
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  ┆ 1      ┆ 0      ┆ 1      │
└──────┴────────┴────────┴────────┘

2. 解决方案步骤详解

要实现上述转换，我们需要依次使用Polars的三个关键操作：unpivot、list.to_struct 和 unnest。

步骤 1: 使用 unpivot 将列名转换为值

unpivot 操作可以将DataFrame的多个列“旋转”成长格式。它会将指定的列名转换为一个新列的值，同时将这些列的原始值放入另一个新列中。

在这个场景中，我们将所有的列（"foo", "bar"）都进行 unpivot 操作。variable_name 参数用于指定存储原始列名的新列的名称，而 value_name 则指定存储原始值的新列的名称（默认为 "value"）。

# 步骤 1: unpivot
df_unpivoted = df.unpivot(variable_name="Name")

print("\n步骤 1 - unpivot 后的 DataFrame:")
print(df_unpivoted)

输出：

步骤 1 - unpivot 后的 DataFrame:
shape: (4, 2)
┌──────┬───────────┐
│ Name ┆ value     │
│ ---  ┆ ---       │
│ str  ┆ list[i64] │
╞══════╪═══════════╡
│ foo  ┆ [1, 2, 3] │
│ foo  ┆ [7, 8, 9] │
│ bar  ┆ [4, 5, 6] │
│ bar  ┆ [1, 0, 1] │
└──────┴───────────┘

现在，我们有了一个 Name 列，其中包含原始列名，以及一个 value 列，其中包含对应的列表数据。

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

接下来，我们需要将 value 列中的每个列表转换为一个结构体（Struct），这样列表中的每个元素都可以被命名为一个字段。list.to_struct() 方法非常适合此目的。fields 参数允许我们定义结构体中各个字段的名称。这里我们使用一个 lambda 函数来动态生成字段名，例如 "Value0", "Value1", "Value2"。

# 步骤 2: 将列表转换为结构体
df_struct = df_unpivoted.with_columns(
    pl.col("value").list.to_struct(fields=lambda x: f"Value{x}")
)

print("\n步骤 2 - 转换为结构体后的 DataFrame:")
print(df_struct)

输出：

步骤 2 - 转换为结构体后的 DataFrame:
shape: (4, 2)
┌──────┬───────────────────────────┐
│ Name ┆ value                     │
│ ---  ┆ ---                       │
│ str  ┆ struct[3]                 │
╞══════╪═══════════════════════════╡
│ foo  ┆ {1,2,3}                   │
│ foo  ┆ {7,8,9}                   │
│ bar  ┆ {4,5,6}                   │
│ bar  ┆ {1,0,1}                   │
└──────┴───────────────────────────┘

现在，value 列的类型已经变成了 struct，并且内部包含了三个字段。

步骤 3: 使用 unnest 展开结构体列

最后一步是使用 unnest 操作。unnest 会将一个结构体列展开，使其内部的每个字段都成为DataFrame中的一个独立列。这正是我们想要实现的目标。

# 步骤 3: 展开结构体列
df_final = df_struct.unnest("value")

print("\n步骤 3 - unnest 后的最终 DataFrame:")
print(df_final)

输出：

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

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

3. 完整代码示例

为了提高效率和代码的可读性，我们可以将这些操作通过方法链（method chaining）的方式组合在一起：

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") # 将列名转换为 "Name" 列的值
    .with_columns(pl.col("value").list.to_struct(fields=lambda x: f"Value{x}")) # 将列表转换为结构体
    .unnest("value") # 展开结构体列
)

print("\n最终转换结果:")
print(output_df)

4. 注意事项与总结

• 列名的灵活性: unpivot 默认会处理所有非 id_vars 的列。如果只想处理部分列，可以使用 columns 参数指定。
• 列表长度一致性: list.to_struct 假定列表中的所有子列表长度相同。如果长度不一致，unnest 后的结果可能会包含空值（null），或者需要额外的处理来填充或对齐。
• 性能优势: Polars的这些操作都是高度优化的，并且利用了其底层的Rust引擎和表达式系统，因此即使处理大规模数据也能保持出色的性能。
• 可读性: 链式操作使得数据转换的逻辑清晰明了，易于理解和维护。

通过本教程，您应该已经掌握了在Polars中将包含列表的列转换为宽格式的有效方法。这种技术在处理嵌套数据或需要将“长”格式数据转换为“宽”格式以进行进一步分析时非常有用。