如何高效为百万级Pandas数据框按权限范围划分用户类别

需求说明 #

现有包含UserName和Permissions列的Pandas数据框，需新增User Class列，分类规则如下：

• 若用户所有权限值<10 → 标记为Admin
• 若用户所有权限值>=10 → 标记为User
• 若用户同时存在两类权限 → 标记为Admin/User

原采用循环分组方式处理80万条数据，效率偏低，需更优实现方案。

示例数据 #

原始数据 #

UserName    Permissions
John Doe             02
John Doe             11
 Example             09
 Example             08
   User3             11

预期结果 #

UserName    Permissions    User Class
John Doe             02    Admin/User
John Doe             11    Admin/User
 Example             09         Admin
 Example             08         Admin
   User3             11          User

原低效代码 #

for UserName, User_df in df.groupby(by='UserName'):
    LT10 = (User_df['Permissions'] < 10).any()
    GTE10 = (User_df['Permissions'] >= 10).any()
    if (LT10 & GTE10):
        UserClass = 'Admin/User'
    elif LT10:
        UserClass = 'Admin'
    elif GTE10:
        UserClass = 'User'
    df.at[User_df.index, 'User Class'] = UserClass

优化方案 #

利用Pandas的向量化分组操作替代Python循环，大幅提升处理效率：

步骤1：统一权限列数据类型 #

先确保Permissions为数值类型（示例中为字符串格式，需转换）：

df['Permissions'] = df['Permissions'].astype(int)

步骤2：批量生成分类标记 #

通过groupby.transform批量计算每个用户的权限特征，再用np.select完成分类：

import numpy as np

# 计算每个用户是否存在<10的权限（结果逐行映射回原数据）
has_lt10 = df.groupby('UserName')['Permissions'].transform(lambda x: (x < 10).any())
# 计算每个用户是否存在>=10的权限（结果逐行映射回原数据）
has_gte10 = df.groupby('UserName')['Permissions'].transform(lambda x: (x >= 10).any())

# 根据条件批量生成User Class
df['User Class'] = np.select(
    condlist=[has_lt10 & has_gte10, has_lt10, has_gte10],
    choicelist=['Admin/User', 'Admin', 'User']
)

优化原理 #

• 原代码的Python循环+df.at逐行赋值属于串行操作，在百万级数据下会产生极大的IO开销
• 优化方案采用Pandas原生的向量化分组API，所有计算由底层C实现，避免了Python循环的性能损耗，处理80万条数据的速度可提升数十倍

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