如何在Polars DataFrame中快速筛选所有pl.List类型的列？

你可以通过两种简洁的方式快速筛选出DataFrame中所有pl.List类型的列：

方法一：使用Polars选择器（推荐） #

Polars的选择器API支持直接通过数据类型匹配列，针对pl.List类型，只需使用pl.selectors.dtype(pl.List)即可：

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    "id": [1, 2, 3],
    "names": [["Alice", "Bob"], ["Charlie"], ["Dave", "Eve"]],
    "scores": [[90, 85], [78], [92, 88]],
    "age": [25, 30, 35]
})

# 选择所有List类型列
list_cols_df = df.select(pl.selectors.dtype(pl.List))
print(list_cols_df)

输出结果：

shape: (3, 2)
┌──────────────────┬─────────────┐
│ names            ┆ scores      │
│ ---              ┆ ---         │
│ list[str]        ┆ list[i64]   │
╞══════════════════╪═════════════╡
│ ["Alice", "Bob"] ┆ [90, 85]    │
│ ["Charlie"]      ┆ [78]        │
│ ["Dave", "Eve"]  ┆ [92, 88]    │
└──────────────────┴─────────────┘

方法二：手动筛选列名 #

如果你偏好手动处理，可以遍历DataFrame的schema，筛选出类型为pl.List的列名后再选择：

# 筛选List类型列的名称
list_col_names = [col for col, dtype in df.schema.items() if dtype == pl.List]
# 选择目标列
list_cols_df = df.select(list_col_names)

这个方法和上面的选择器方式效果完全一致，适合需要自定义筛选逻辑的场景。

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