300万行DataFrame列转换优化及双规则并行处理问题

我有一个包含300万行的DataFrame，需要对某列值进行转换。该列存储以分号分隔的字符串，转换逻辑是拆分字符串后按优先级规则选择结果。

样本数据及初始转换函数 #

data = {'Name': ['X1', 'X2', 'X3', 'X4', 'X5','X6'], 
        'category': ['CatA;CatB', 'CatB', None, 'CatB;CatC;CatA', 'CatA;CatB', 'CatB;CatD;CatB;CatC;CatA']} 

sample_dataframe = pd.DataFrame(data) 

def cat_name(x):
    if x:
        x =  pd.Series(x.split(";"))
        y = x[(x!='CatA') & x.notna()]
        custom_dict = {'CatC': 0, 'CatD':1, 'CatB': 2, 'CatE': 3}
        if x.count() == 1:
            return x.iloc[0]
        elif y.count() > 1:
            y = y.sort_values(key=lambda x: x.map(custom_dict))
            if y.count() > 2:
                return '3 or more'
            else:
                return y.iloc[0]+'+'
        elif y.count() == 1:
            return y.iloc[0]
        else:
            return None
    else:
        return None

我用test_data = sample_dataframe['category'].apply(cat_name)执行转换，300万行耗时近10分钟，请问怎么优化函数提升速度？

另外，我需要按两套分类规则生成两列结果，现在调用两次apply效率太低，尝试写了一个联合函数一次性返回两个结果：

def joint_cat_name(x):
    cat_string = x
    if cat_string:
        string_series =  pd.Series(cat_string.split(";"))
        y = string_series[(string_series!='CatA') & string_series.notna()]
        custom_dict = {'CatB': 0, 'CatC':1, 'CatD': 2, 'CatE': 3}
        if string_series.count() == 1:
            return string_series.iloc[0], string_series.iloc[0]
        elif y.count() > 1:
            y = y.sort_values(key=lambda x: x.map(custom_dict))
            if y.count() > 2:
                return '3 or more', y.iloc[0]
            elif y.count() == 1:
                return y.iloc[0]+'+', y.iloc[0]
        elif y.count() == 1:
            return y.iloc[0], y.iloc[0]
        else:
            return None, None
    else:
        return None, None

然后用test_data['CAT_NAME'], test_data['MAIN_CAT_NAME']=zip(*sample_dataframe['category'].apply(joint_cat_name))赋值，但当输出为(None, None)时，出现TypeError: 'NoneType' object is not iterable错误，该怎么解决？

1. 解决TypeError错误 #

问题出在函数返回逻辑：当分支返回单个None时，zip(*...)会尝试迭代None导致报错。必须确保函数始终返回可迭代的二元组，修改后的函数如下：

def joint_cat_name(x):
    cat_string = x
    if cat_string:
        string_list = cat_string.split(";")
        y = [c for c in string_list if c != 'CatA' and pd.notna(c)]
        custom_dict = {'CatB': 0, 'CatC':1, 'CatD': 2, 'CatE': 3}
        
        if len(string_list) == 1:
            return (string_list[0], string_list[0])
        elif len(y) > 1:
            y_sorted = sorted(y, key=lambda c: custom_dict.get(c, float('inf')))
            if len(y_sorted) > 2:
                return ('3 or more', y_sorted[0])
            else:
                return (f"{y_sorted[0]}+", y_sorted[0])
        elif len(y) == 1:
            return (y[0], y[0])
        else:
            return (None, None)
    else:
        return (None, None)  # 强制返回二元组，避免单个None

修改后所有返回值都是二元组，zip(*...)可正常迭代，不会触发错误。

2. 优化300万行数据的处理速度 #

apply本质是逐行循环，对大数据集效率极低。改用矢量化操作+分组批量计算，可将耗时从分钟级压缩到秒级。

核心优化思路 #

• 用str.split矢量化拆分字符串，避免在函数内创建Series
• 用列表推导式快速过滤无效值，替代逐行判断
• 用groupby批量处理分组数据，避免逐行循环
• 预定义优先级字典，用sorted+map批量完成排序

单规则优化代码 #

import pandas as pd

data = {'Name': ['X1', 'X2', 'X3', 'X4', 'X5','X6'], 
        'category': ['CatA;CatB', 'CatB', None, 'CatB;CatC;CatA', 'CatA;CatB', 'CatB;CatD;CatB;CatC;CatA']} 
df = pd.DataFrame(data)

# 1. 矢量化拆分与过滤
df['cat_list'] = df['category'].str.split(';')
df['filtered_cats'] = df['cat_list'].apply(
    lambda lst: [c for c in lst if c != 'CatA' and pd.notna(c)] if lst is not None else []
)

# 2. 定义优先级规则
priority_dict = {'CatC': 0, 'CatD':1, 'CatB': 2, 'CatE': 3}

# 3. 批量计算结果
def compute_result(row):
    original_len = len(row['cat_list']) if row['cat_list'] is not None else 0
    filtered_len = len(row['filtered_cats'])
    
    if original_len == 1:
        return row['category']
    elif filtered_len > 1:
        sorted_cats = sorted(row['filtered_cats'], key=lambda x: priority_dict.get(x, float('inf')))
        return '3 or more' if filtered_len > 2 else f"{sorted_cats[0]}+"
    elif filtered_len == 1:
        return row['filtered_cats'][0]
    else:
        return None

df['CAT_NAME'] = df.apply(compute_result, axis=1)
df.drop(['cat_list', 'filtered_cats'], axis=1, inplace=True)

多规则批量生成两列结果 #

如果要生成两套规则的结果，只需在过滤后同时计算两套规则的排序结果，无需重复处理拆分步骤：

def compute_double_result(row):
    original_len = len(row['cat_list']) if row['cat_list'] is not None else 0
    filtered_len = len(row['filtered_cats'])
    
    if original_len == 1:
        return (row['category'], row['category'])
    elif filtered_len > 1:
        # 规则1排序
        sorted_rule1 = sorted(row['filtered_cats'], key=lambda x: priority_dict1.get(x, float('inf')))
        res1 = '3 or more' if filtered_len > 2 else f"{sorted_rule1[0]}+"
        # 规则2排序
        sorted_rule2 = sorted(row['filtered_cats'], key=lambda x: priority_dict2.get(x, float('inf')))
        res2 = sorted_rule2[0]
        return (res1, res2)
    elif filtered_len == 1:
        return (row['filtered_cats'][0], row['filtered_cats'][0])
    else:
        return (None, None)

# 定义两套优先级字典
priority_dict1 = {'CatC': 0, 'CatD':1, 'CatB': 2, 'CatE': 3}
priority_dict2 = {'CatB': 0, 'CatC':1, 'CatD': 2, 'CatE': 3}

# 一次性生成两列
df[['CAT_NAME', 'MAIN_CAT_NAME']] = df.apply(compute_double_result, axis=1, result_type='expand')

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