使用自定义函数分组汇总pandas DataFrame出现错误输出的排查

问题背景 #

现有如下Pandas DataFrame，记录了4个人喜欢的水果：

import numpy as np
import pandas as pd

df = pd.DataFrame({
    "name": ["danny", "danny", "danny", "monica", "monica", "monica", "fred", "fred", "sam", "sam"],
    "fruit": ["apricot", "apple", "orange", "apricot", "banana", "watermelon", "apple", "apricot", "apricot", "peach"]
})

期望输出每个用户是否同时喜欢apricot和apple的布尔值表格：

name  fruit
0   danny   True
1  monica  False
2    fred   True
3     sam  False

尝试与错误 #

定义了检查目标字符串是否全部存在的函数，单独调用功能正常：

def is_needle_in_haystack(needle, haystack):
  return all(x in haystack for x in needle)

# 测试正常
is_needle_in_haystack(["zebra", "lion"], ["whale", "lion", "dog"])  # 返回False
is_needle_in_haystack(["rabbit", "cat"], ["hamster", "cat", "monkey", "rabbit"])  # 返回True

但将其用于groupby的agg操作时，得到全False的错误结果：

target_fruits = ["apricot", "apple"]
df.groupby(df["name"]).agg({"fruit": lambda x: is_needle_in_haystack(target_fruits, x)})

输出：

fruit
name         
danny   False
fred    False
monica  False
sam     False

错误原因 #

核心问题是：groupby后传给lambda的x是Pandas Series对象，而非普通列表。
在Python中，对Series执行x in haystack时，判断的是x是否存在于Series的索引中，而不是Series的值里。比如针对danny的fruit Series，索引是[0,1,2]，判断"apricot" in x时，实际是检查"apricot"是否是索引的元素，显然不成立，因此每个判断都返回False，最终all()的结果自然全为False。

解决方法 #

方法1：修改函数参数为列表 #

在调用函数时，将Series转为普通列表即可，修改agg的lambda表达式：

target_fruits = ["apricot", "apple"]
result = df.groupby("name").agg({"fruit": lambda x: is_needle_in_haystack(target_fruits, x.tolist())}).reset_index()

或者修改函数内部，自动处理Series类型：

def is_needle_in_haystack(needle, haystack):
    # 将Series转为列表
    if hasattr(haystack, 'tolist'):
        haystack = haystack.tolist()
    return all(x in haystack for x in needle)

# 原调用代码即可正常工作
result = df.groupby("name").agg({"fruit": lambda x: is_needle_in_haystack(target_fruits, x)}).reset_index()

方法2：使用Pandas原生更高效的写法 #

无需自定义函数，直接利用Pandas的方法实现：

target_fruits = ["apricot", "apple"]
# 按分组检查目标水果是否全部是当前组的子集
result = df.groupby('name')['fruit'].apply(lambda x: set(target_fruits).issubset(x)).reset_index(name='fruit')

或者用values直接访问Series的值：

result = df.groupby('name')['fruit'].agg(lambda x: all(f in x.values for f in target_fruits)).reset_index()

以上方法都能得到符合预期的结果。

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