如何向量化Pandas DataFrame中基于自定义索引的列赋值？

• 现有约30万行学生课程考试数据，后续将增长至3亿行
• 需要基于学生+课程的组合分组，根据组内exam_status的多个值，按照优先级（Fail > Absent > Pass/Skip）设置course_status列
• 当前迭代或lambda方法在Mac上运行超5分钟，无法适配数据增长

数据示例 #

d = [{'name':'Student1', 'course':'Math', 'exam_status':'Pass'},
     {'name':'Student1', 'course':'Math', 'exam_status':'Fail'},
     {'name':'Student1', 'course':'Math', 'exam_status':'Pass'},
     {'name':'Student1', 'course':'Chemistry', 'exam_status':'Pass'},
     {'name':'Student1', 'course':'Chemistry', 'exam_status':'Pass'},
     {'name':'Student1', 'course':'Chemistry', 'exam_status':'Pass'},
     {'name':'Student1', 'course':'Art', 'exam_status':'Absent'},
     {'name':'Student1', 'course':'Art', 'exam_status':'Absent'},
     {'name':'Student1', 'course':'Art', 'exam_status':'Fail'},
     {'name':'Student1', 'course':'Computer Science', 'exam_status':'Fail'},
     {'name':'Student1', 'course':'Computer Science', 'exam_status':'Fail'},
     {'name':'Student1', 'course':'Computer Science', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Math', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Math', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Math', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Chemistry', 'exam_status':'Fail'},
     {'name':'Student2', 'course':'Chemistry', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Chemistry', 'exam_status':'Absent'},
     {'name':'Student2', 'course':'Art', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Art', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Art', 'exam_status':'Fail'},
     {'name':'Student2', 'course':'Computer Science', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Computer Science', 'exam_status':'Pass'},
     {'name':'Student2', 'course':'Computer Science', 'exam_status':'Pass'},]

df = pd.DataFrame(d)
df['course_status'] = ''
# 创建学生+课程的组合索引
df['name_course'] = df[['name', 'course']].apply(tuple, axis=1)
df = df.set_index(['name_course'], drop=True)

现有低效实现 #

方法1：遍历唯一索引 #

for i in df.index.unique():
    status_list = list(df.loc[[i]]['exam_status'])
    if 'Fail' in status_list: df.loc[[i], 'course_status']  = 'Fail'
    elif 'Absent' in status_list: df.loc[[i], 'course_status']  = 'Absent'
    elif 'Pass' or 'Skip' in status_list: df.loc[[i], 'course_status']  = 'Pass'
print(df.to_markdown())

方法2：Lambda分组处理 #

# 定义状态判断函数
def set_course_status(i):
    status_list = list(df.loc[[i]]['exam_status'])
    if 'Fail' in status_list: return 'Fail'
    elif 'Absent' in status_list: return 'Absent'
    elif 'Pass' in status_list: return 'Pass'
df['course_status'] = df.groupby(df.index).apply(lambda x : set_course_status(x.name))

高效向量化实现方案 #

方案思路 #

利用Pandas底层C实现的向量化操作，避免Python循环和lambda的性能损耗：

1. 给exam_status设置优先级权重（Fail:0, Absent:1, Pass/Skip:2），权重越小优先级越高
2. 按学生+课程分组，提取组内优先级最高的状态作为course_status

代码实现（简洁版） #

# 定义状态优先级映射（优先级从高到低：Fail > Absent > Pass/Skip）
status_order = {'Fail': 0, 'Absent': 1, 'Pass': 2, 'Skip': 2}

# 将状态转换为优先级数值
df['status_priority'] = df['exam_status'].map(status_order)

# 按分组取优先级最小的对应状态
df['course_status'] = df.groupby(df.index)['exam_status'].transform(
    lambda x: x[status_order[x] == status_order[x].min()].iloc[0]
)

# 删除临时列
df.drop('status_priority', axis=1, inplace=True)

代码实现（极致性能版） #

完全规避lambda，纯向量化操作，适配超大规模数据：

status_order = {'Fail': 0, 'Absent': 1, 'Pass': 2, 'Skip': 2}
# 创建优先级列
df['priority'] = df['exam_status'].map(status_order)

# 计算每个分组的最小优先级
min_priority = df.groupby(df.index)['priority'].transform('min')

# 筛选出每个分组中优先级最高的行，提取对应状态
course_status_map = df[df['priority'] == min_priority].drop_duplicates(df.index)['exam_status']
df['course_status'] = df.index.map(course_status_map)

# 删除临时列
df.drop('priority', axis=1, inplace=True)

性能说明 #

• 向量化操作比Python循环快100-1000倍，3亿行数据可在数分钟内完成处理（取决于硬件配置）
• 完全适配数据量增长后的性能需求

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