如何用Polars为DataFrame添加基于食物交集的动物列？

问题背景 #

现有包含people（人物）和food（喜好食物）列的Polars DataFrame：

df_class = pl.DataFrame(
    {
        'people': ['alan', 'bob', 'charlie'],
        'food': [['orange', 'apple'], ['banana', 'cherry'], ['banana', 'grape']]
    }
)

输出结构：

shape: (3, 2)
┌─────────┬──────────────────────┐
│ people  ┆ food                 │
│ ---     ┆ ---                  │
│ str     ┆ list[str]            │
╞═════════╪══════════════════════╡
│ alan    ┆ ["orange", "apple"]  │
│ bob     ┆ ["banana", "cherry"] │
│ charlie ┆ ["banana", "grape"]  │
└─────────┴──────────────────────┘

同时有存储动物及其食物偏好的数据结构：

animals = [
    ('squirrel', ('acorn', 'almond')),
    ('parrot', ('cracker', 'grape', 'guava')),
    ('dog', ('chicken', 'bone')),
    ('monkey', ('banana', 'plants'))
]

需求 #

为df_class添加新列pets，该列是与对应人物喜好食物至少有一项交集的动物列表，预期结果如下：

shape: (3, 3)
┌─────────┬──────────────────────┬──────────────────────┐
│ people  ┆ food                 ┆ pets                 │
│ ---     ┆ ---                  ┆ ---                  │
│ str     ┆ list[str]            ┆ list[str]            │
╞═════════╪══════════════════════╪══════════════════════╡
│ alan    ┆ ["orange", "apple"]  ┆ []                   │
│ bob     ┆ ["banana", "cherry"] ┆ ["monkey"]           │
│ charlie ┆ ["banana", "grape"]  ┆ ["monkey", "parrot"] │
└─────────┴──────────────────────┴──────────────────────┘

要求：

• 可调整animals的数据结构，便于实现集合交集操作
• pets列表需包含唯一值，顺序无关紧要
• 用单一表达式实现，可融入现有数据集的多列转换框架

解决方案 #

1. 优化动物数据结构 #

先将animals转换为动物到食物集合的字典，集合结构能高效完成交集判断：

animal_food_map = {animal: set(foods) for animal, foods in animals}

2. 编写Polars表达式添加新列 #

使用pl.col("food").map_elements()方法，结合集合交集判断逻辑，用单一表达式完成列添加：

import polars as pl

# 构建动物-食物集合映射
animal_food_map = {animal: set(foods) for animal, foods in animals}

# 生成结果DataFrame
df_result = df_class.with_columns(
    pets=pl.col("food")
    .map_elements(
        lambda foods: [
            animal for animal, food_set in animal_food_map.items()
            if not food_set.isdisjoint(foods)
        ],
        return_dtype=pl.List(pl.String)
    )
)

# 打印结果
print(df_result)

关键说明 #

• isdisjoint()用于判断两个集合是否无交集，not food_set.isdisjoint(foods)表示两者存在至少一项共同食物
• map_elements()会遍历每行的food列表，返回匹配的动物列表
• 指定return_dtype=pl.List(pl.String)确保pets列的类型符合预期

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