Polars中如何展开所有列的单元素列表？现有实现报错求助

问题描述 #

现有如下Polars DataFrame：

import polars as pl

df = pl.DataFrame(dict(
  j=[[1], [2], [3]],
  k=[[1, 1], [2], [3]],
))

输出结构：

j (list[i64])  k (list[i64])
 [1]            [1, 1]
 [2]            [2]
 [3]            [3]
shape: (3, 2)

需求：将所有列中全为单元素的列表展开为基础数据类型，保留包含多元素列表的列结构，最终得到：

dfj = pl.DataFrame(dict(
  j=[1, 2, 3],
  k=[[1, 1], [2], [3]],
))

输出结构：

j (i64)  k (list[i64])
 1        [1, 1]
 2        [2]
 3        [3]
shape: (3, 2)

尝试以下代码后报错：

dfj = (df
  .with_columns(
    pl.when(pl.col(col).list.lengths().max() == 1)
    .then(pl.col(col).list.first())
    .otherwise(pl.col(col))
    for col in df.columns
  )
)

错误信息：

exceptions.ArrowErrorException: NotYetImplemented("Casting from Int64 to LargeList(Field { name: \"item\", data_type: Int64, is_nullable: true, metadata: {} }) not supported")

错误原因分析 #

pl.when/then/otherwise要求所有分支返回的数据类型必须一致：

• 对j列，then分支提取列表元素后返回Int64基础类型；
• otherwise分支返回原列的List[Int64]类型；
  Polars会尝试将Int64强制转换为List[Int64]以匹配类型，但Arrow目前不支持这种基础类型转列表类型的反向转换，因此抛出错误。

可行实现方法 #

通过select遍历列，针对每列单独判断并生成对应类型的列，避免类型冲突：

import polars as pl

df = pl.DataFrame(dict(
  j=[[1], [2], [3]],
  k=[[1, 1], [2], [3]],
))

dfj = df.select(
    [
        pl.col(col).list.first().alias(col) 
        if pl.col(col).list.lengths().max() == 1 
        else pl.col(col)
        for col in df.columns
    ]
)

print(dfj)

逻辑说明 #

1. 遍历DataFrame的每一列；
2. 检查该列所有列表的长度最大值是否为1（即所有元素都是单元素列表）；
3. 若是，提取列表的第一个元素并替换原列，自动转为基础数据类型；
4. 若否，直接保留原列的列表类型。

这种方式会让Polars自动推导每列的最终数据类型，不会出现强制类型转换的冲突问题。

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