Polars中列表列的结构化转换与重塑技巧

本文详细介绍了如何在polars dataframe中将包含列表的列进行高效重塑。通过组合使用`unpivot`、`list.to_struct`和`unnest`等核心操作，教程演示了如何将宽格式的列表列转换为长格式，并动态地将列表元素扩展为独立的数值列，从而实现复杂的数据结构转换，提升数据处理的灵活性和效率。

在数据分析和处理中，我们经常会遇到需要将数据从一种结构转换到另一种结构的情况。特别是在处理包含列表（List）类型数据的列时，将其展开并重塑成更易于分析的表格形式是一个常见的需求。Polars作为一款高性能的DataFrame库，提供了强大且灵活的API来应对这类挑战。本教程将详细讲解如何利用Polars的unpivot、list.to_struct和unnest等操作，将一个包含列表列的DataFrame转换为指定的长格式，其中原始列名将成为一个新列的值，而列表中的元素则被展开成新的数值列。

初始数据结构

假设我们有一个Polars 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      │
└──────┴────────┴────────┴────────┘

转换步骤详解

要实现上述转换，我们需要分三步操作：

1. 解除透视（Unpivot）：将原始列名转换为一个新列的值。
2. 列表转结构体（List to Struct）：将包含列表的列转换为结构体（Struct）列，为下一步的展开做准备。
3. 展开结构体（Unnest）：将结构体列展开成多个独立的列。

下面我们将详细介绍每一步的操作。

步骤一：解除透视 (unpivot)

unpivot操作（也常被称为“melt”或“stack”）用于将DataFrame从宽格式转换为长格式。它会将一个或多个指定列的名称和值转换为新的两列：一列包含原始列名（通常称为“变量”列），另一列包含原始列的值（通常称为“值”列）。

在本例中，我们将foo和bar两列解除透视。variable_name参数用于指定存储原始列名的新列的名称，value_name参数用于指定存储原始列值的新列的名称。

df_unpivoted = df.unpivot(variable_name="Name", value_name="value")
print("\n解除透视后的DataFrame:")
print(df_unpivoted)

输出如下：

解除透视后的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] │
└──────┴───────────┘

现在，原始的foo和bar列名已合并到Name列中，而它们对应的列表值则合并到value列中。

步骤二：列表转结构体 (list.to_struct)

unnest操作只能作用于结构体（Struct）列。因此，在展开value列中的列表之前，我们需要先将其转换为一个结构体列。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("\n列表转换为结构体后的DataFrame:")
print(df_struct)

输出如下：

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

可以看到，value列现在已经从list[i64]类型变成了struct[i64, i64, i64]类型，其内部包含了三个匿名字段，对应着原始列表的元素。

步骤三：展开结构体 (unnest)

最后一步是使用unnest操作将结构体列展开为多个独立的列。我们指定要展开的列名为value。

df_final = df_struct.unnest("value")
print("\n最终转换后的DataFrame:")
print(df_final)

输出如下：

最终转换后的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转换成了目标格式。

完整代码示例

为了提高代码的可读性和执行效率，通常会将这些操作链式调用：

import polars as pl
df = pl.DataFrame({
"foo": [[1, 2, 3], [7, 8, 9]],
"bar": [[4, 5, 6], [1, 0, 1]]
})
transformed_df = (
df
.unpivot(variable_name="Name") # 默认 value_name 为 "value"
.with_columns(pl.col("value").list.to_struct(fields=lambda x : f"Value{x}"))
.unnest("value")
)
print("最终转换后的DataFrame (链式调用):")
print(transformed_df)

注意事项与总结

• 列名冲突：在使用unpivot时，如果原始DataFrame中已经存在名为Name或value的列，需要通过variable_name和value_name参数指定不同的名称，以避免冲突。
• 列表长度一致性：list.to_struct操作要求列表中所有子列表的长度一致。如果列表长度不一致，可能会导致错误或填充null值。在实际应用中，可能需要先对列表进行填充或截断操作。
• 性能：Polars的表达式API和其底层Rust实现使得这些链式操作在处理大数据时依然保持高效。
• 灵活性：fields参数的lambda函数提供了极大的灵活性，可以根据需要动态生成各种列名。

通过本教程，您应该已经掌握了在Polars中进行复杂数据重塑的关键技巧，特别是如何处理和展开包含列表的列。这些操作在数据预处理、特征工程和报告生成等场景中都非常实用。