Polars 列表类型为逗号分隔字符串

Question

我有一个 df，我想对其进行分组并写入 csv 格式。但是，其中一列的列表类型阻止将 df 写入 csv。

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

当尝试将上述数据帧写入 csv 时，出现错误：

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

如果尝试将列表类型转换为 pl.Utf8 会导致错误

(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')

如果我尝试在 group_by 上下文中分解列表：

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

输出不是想要的：

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

对我来说 group_by 然后写入 csv 最方便的方法是什么？

以 csv 形式写入的所需输出：

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

Answer 1

最近有一次关于为什么会出现这种情况的讨论。

可以使用

._s.get_fmt()

 来“字符串化”列表：

print(
   df
    .groupby(by="Column A").agg(pl.col("Column B").unique())
    .with_columns(
       pl.col("Column B").map(lambda row: 
          [row._s.get_fmt(n, 0) for n in range(row.len())]
       ).flatten())
    .write_csv(),
    end=""
)

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

另一种方法是使用

str()

，如 

@FObersteiner 建议的那样。

print(
   df.groupby("Column A").agg(
      pl.col("Column B")
        .unique()
        .apply(lambda col: str(col.to_list()))
   ).write_csv(), 
   end=""
)

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

“字符串化”列表的主要问题是 - 当您读回 CSV 数据时 - 您不再具有

list[]

 类型。

import io

pl.read_csv(io.StringIO(
   'Column A,Column B\nVariable 4,"[""CD""]"\n'
   'Variable 1,"[""AB""]"\nVariable 2,"[""AB"", ""CD""]"\n'
   'Variable 3,"[""CD"", ""AB""]"\n'
))

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

这就是建议使用替代格式的原因。