Polars List Type to Comma Separated String - csv

Question

I have a df that I'd like to groupby and write to csv format. However, one of the columns has a list type that prevents writing the df to 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"]})
Which I want to group by as below:
df.groupby(by="Column A").agg(pl.col("Column B").unique())
Output:
shape: (4, 2)
┌────────────┬──────────────┐
│ Column A ┆ Column B │
│ --- ┆ --- │
│ str ┆ list[str] │
╞════════════╪══════════════╡
│ Variable 3 ┆ ["AB", "CD"] │
│ Variable 1 ┆ ["AB"] │
│ Variable 4 ┆ ["CD"] │
│ Variable 2 ┆ ["CD", "AB"] │
└────────────┴──────────────┘
When trying to write the above dataframe to csv it comes up with an error: "ComputeError: CSV format does not support nested data. Consider using a different data format. Got: 'list[str]'"
If trying to convert the list type to pl.Utf8 it leads to an error
(df
.groupby(by="Column A").agg(pl.col("Column B").unique())
.with_columns(pl.col("Column B").cast(pl.Utf8))
)
Output: "ComputeError: Cannot cast list type"
If I try to explode the list in the groupby context:
df.groupby(by="Column A").agg(pl.col("Column B").unique().explode())
The output is not desired:
shape: (4, 2)
┌────────────┬─────────────────────┐
│ Column A ┆ Column B │
│ --- ┆ --- │
│ str ┆ list[str] │
╞════════════╪═════════════════════╡
│ Variable 1 ┆ ["A", "B"] │
│ Variable 3 ┆ ["A", "B", ... "D"] │
│ Variable 2 ┆ ["A", "B", ... "B"] │
│ Variable 4 ┆ ["A", "B", ... "D"] │
└────────────┴─────────────────────┘
What would be the most convenient way for me to groupby and then write to csv?
Desired output written in 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

There was a recent discussion about why this is the case.
It is possible to use ._s.get_fmt() to "stringify" the lists:
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""]"
Another way is using str() as #FObersteiner has suggested.
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']
The main probem with "stringifying" lists is - when you read the CSV data back in - you no longer have a list[] type.
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"] │
└────────────┴──────────────┘
This is the reason for the recommendation of using an alternative format.