如何将Polars分组后含list类型列的DataFrame导出为CSV？

问题场景 #

使用Polars对DataFrame分组后，某列变为list[str]类型，直接导出CSV时因CSV不支持嵌套数据报错；尝试直接将列表列转为字符串类型失败；使用explode会保留列表结构，不符合导出需求。

原始数据代码 #

import polars as pl

df = pl.DataFrame({
    "Column A": ["Variable 1", "Variable 2", "Variable 2", "Variable 3", "Variable 3", "Variable 4"],
    "Column B": ["AB", "AB", "CD", "AB", "CD", "CD"]
})

分组代码及结果 #

df.group_by("Column A").agg(pl.col("Column B").unique())

分组后输出：

shape: (4, 2)
┌────────────┬──────────────┐
│ Column A   ┆ Column B     │
│ ---        ┆ ---          │
│ str        ┆ list[str]    │
╞════════════╪══════════════╡
│ Variable 4 ┆ ["CD"]       │
│ Variable 2 ┆ ["CD", "AB"] │
│ Variable 3 ┆ ["AB", "CD"] │
│ Variable 1 ┆ ["AB"]       │
└────────────┴──────────────┘

遇到的错误 #

1. 直接导出CSV报错：

ComputeError: CSV format does not support nested data. Consider using a different data format. Got: 'list[str]'

2. 尝试直接转换列表为字符串报错：

(df
   .group_by("Column A").agg(pl.col("Column B").unique())
   .with_columns(pl.col("Column B").cast(pl.String))
)

InvalidOperationError: cannot cast List type (inner: 'String', to: 'String')

3. 使用explode不符合预期：

df.group_by("Column A").agg(pl.col("Column B").unique().explode())

结果仍保留列表结构，无法导出CSV。

便捷解决方案 #

核心思路是将list[str]列转换为字符串格式，让CSV可以识别。以下提供两种方法：

方法1：快速转换（适合小数据量） #

使用map_elements(str)将列表直接转为Python字符串格式：

# 分组并转换列表列为字符串
result_df = df.group_by("Column A").agg(pl.col("Column B").unique()) \
              .with_columns(pl.col("Column B").map_elements(str))

# 导出CSV
result_df.write_csv("grouped_output.csv")

导出的CSV内容示例：

Column A,Column B
Variable 4,"['CD']"
Variable 2,"['CD', 'AB']"
Variable 3,"['AB', 'CD']"
Variable 1,"['AB']"

方法2：自定义格式（适合大数据量/特定格式需求） #

使用Polars内置列表函数构建符合期望的字符串格式（如双引号包裹的列表），避免逐行处理的性能损耗：

result_df = df.group_by("Column A").agg(pl.col("Column B").unique()) \
              .with_columns(
                  pl.col("Column B")
                  .list.join('", "')  # 用", "连接列表元素
                  .map_elements(lambda x: f'["{x}"]', return_dtype=pl.String)  # 包裹成列表格式字符串
              )

result_df.write_csv("grouped_output_custom.csv")

导出的CSV内容示例：

Column A,Column B
Variable 4,"[""CD""]"
Variable 2,"[""CD"", ""AB""]"
Variable 3,"[""AB"", ""CD""]"
Variable 1,"[""AB""]"

说明 #

• 直接cast(pl.String)失败是因为Polars不支持嵌套类型直接转为基础字符串类型，必须显式将列表内容转为字符串表示。
• explode的作用是将列表展开为多行，但会破坏分组后的一行对应一个分组的结构，因此不适用。

内容的提问来源于stack exchange，提问作者miroslaavi