Polars中不使用map_elements将列表封装为嵌套列表的方法

要解决when/then中普通行pl.List(pl.String)与处理后大行pl.List(pl.List(pl.String))的类型不匹配问题，核心是统一两个分支的返回类型：将普通行的列表包装为单元素的嵌套列表，让它和处理后的大行类型一致。

正确实现代码 #

import polars as pl

# 示例数据
df = pl.DataFrame({
    "explode_me": [
        ["a"] * 10,       # 普通行：无需拆分
        ["b"] * 10000,    # 大行：需要拆分
        ["c"] * 5         # 普通行：无需拆分
    ]
})

# 统一类型的处理逻辑
df = df.with_columns(
    pl.when(pl.col("explode_me").list.lengths() > 100)
    # 仅对大行使用map_elements拆分
    .then(pl.col("explode_me").map_elements(
        lambda x: [x[i:i+1000] for i in range(0, len(x), 1000)]
    ))
    # 普通行用list.wrap()包装成嵌套列表，统一类型
    .otherwise(pl.col("explode_me").list.wrap())
    .alias("explode_me")
)

# 两次explode得到最终拆分结果
result = df.explode("explode_me").explode("explode_me")
print(result.head())

关键说明 #

• pl.col("explode_me").list.wrap()会将原列表[a,b,c]转为[[a,b,c]]，把普通行的List(String)类型升级为List(List(String))，完美匹配大行处理后的类型。
• 后续通过两次explode：第一次展开嵌套列表的外层，第二次展开内层的子列表，最终得到单个元素的行，和直接explode普通行的效果一致。

更高效的替代方案（避免map_elements） #

如果想进一步优化性能，可改用Polars原生表达式拆分大行，无需依赖map_elements：

df = df.with_columns(
    pl.when(pl.col("explode_me").list.lengths() > 100)
    .then(
        pl.int_range(0, pl.col("explode_me").list.lengths(), 1000)
        .list.eval(pl.col("explode_me").list.slice(pl.element(), 1000))
    )
    .otherwise(pl.col("explode_me").list.wrap())
    .alias("explode_me")
)

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