Pandas中DataFrame值替换操作耗时过长的问题求助

问题背景 #

我编写了一个concatenate_fields函数，移除df.loc[df['ln_1_adr'] == df['ln_2_adr'], df['ln_1_adr']] = ''这一行时，函数仅需毫秒级即可运行。但这行代码是用来清除地址第一行的冗余信息的，偏偏它运行耗时极长，传入函数的是一个3000行8列的DataFrame。

原函数代码如下：

import pandas as pd

def concatenate_fields(df):
    df = pd.DataFrame(df).copy()
    for col in ['ln_1_adr','ln_2_adr','city_adr', 'state_adr', 'zip_5_adr','zip_4_adr']:
        df[col] = df[col].str.strip()
        if col in ('ln_1_adr','ln_2_adr'):
            df[col] = df[col].str.replace(r"-", " ", regex=False)
    df.loc[df['ln_1_adr'] == df['ln_2_adr'], df['ln_1_adr']] = '' # 这行代码耗时极长
    df["zip_plus_4"] = df.zip_5_adr.str.cat(df['zip_4_adr'], sep=' ', na_rep='')
    df["ADDRESS"] = df.ln_1_adr.str.cat(df[['ln_2_adr','city_adr', 'state_adr', 'zip_plus_4']], sep = ' ', na_rep='').str.strip()
                                            
    return df

优化方案 #

原代码中耗时的那行问题在于：使用df.loc[条件, df['ln_1_adr']]的写法会创建列的副本，带来额外的内存开销和计算延迟。改用pandas的向量化方法mask()可以直接对列进行高效处理，避免不必要的拷贝。

优化后的函数代码：

import pandas as pd

def concatenate_fields(df):
    df = pd.DataFrame(df).copy()
    for col in ['ln_1_adr','ln_2_adr','city_adr', 'state_adr', 'zip_5_adr','zip_4_adr']:
        df[col] = df[col].str.strip()
        if col in ('ln_1_adr','ln_2_adr'):
            df[col] = df[col].str.replace(r"-", " ", regex=False)
    # 替换原耗时行，使用向量化的mask方法
    df['ln_1_adr'] = df['ln_1_adr'].mask(df['ln_1_adr'] == df['ln_2_adr'], '')
    df["zip_plus_4"] = df.zip_5_adr.str.cat(df['zip_4_adr'], sep=' ', na_rep='')
    df["ADDRESS"] = df.ln_1_adr.str.cat(df[['ln_2_adr','city_adr', 'state_adr', 'zip_plus_4']], sep = ' ', na_rep='').str.strip()
                                            
    return df

补充说明 #

另外，也可以把原loc写法中的列索引改为字符串形式（即df.loc[df['ln_1_adr'] == df['ln_2_adr'], 'ln_1_adr'] = ''），这也能避免创建列副本的问题，提升运行效率。不过mask()方法的可读性和执行效率更优，推荐使用。

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