基于Series参考规则为Dataframe添加batch分组列的技术需求

现有DataFrame df1 存储模型与部件关联信息，Series s1 定义不同模型的分组规则：模型"A"需按model字段分组，模型"B"需按part字段分组。需要为df1新增batch列，根据规则标记每条数据的批次编号。

完整实现代码 #

import pandas as pd
import numpy as np

# 原始数据
df1 = pd.DataFrame({'model':['A','A','A','A','A','B','B','B','B','B'],
                'part':['part_1','part_1','part_2','part_2','part_3','part_4','part_4','part_5','part_5','part_5']})
s1 = pd.Series({'A':'model', 'B':'part'})

# 将规则Series转为字典，便于快速查找
rule_map = s1.to_dict()

# 初始化batch列
df1['batch'] = np.nan

# 遍历每个模型对应的分组规则，生成批次编号
for model_val, group_col in rule_map.items():
    # 筛选当前模型的所有行
    mask = df1['model'] == model_val
    # 对指定分组字段生成连续的批次编号（相同值对应相同编号）
    df1.loc[mask, 'batch'] = df1.loc[mask, group_col].rank(method='dense').astype(int)

# 查看结果
print(df1)

运行结果 #

model    part  batch
0     A  part_1      1
1     A  part_1      1
2     A  part_2      1
3     A  part_2      1
4     A  part_3      1
5     B  part_4      1
6     B  part_4      1
7     B  part_5      2
8     B  part_5      2
9     B  part_5      2

关键逻辑说明 #

• rule_map = s1.to_dict()：将规则Series转换为字典，实现模型到分组字段的快速映射，后续新增模型规则时只需修改s1即可。
• rank(method='dense')：生成连续无间隔的批次编号，确保相同分组值对应相同编号，避免出现跳跃式编号。
• 遍历规则字典：批量处理不同模型的分组逻辑，代码复用性和扩展性更强。

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