使用timeit测试纯Pandas代码时多次运行结果一致的问题

问题概述 #

编写了纯Pandas实现的pure_pandas函数，通过compare_queries调用timeit工具分别执行10次和50次性能测试，发现两次测试结果完全相同；但拆分函数实现时能得到不同结果。同时当前代码存在大量重复逻辑，需要在不重复代码的前提下解决问题，验证同一查询的性能差异。

现有代码 #

纯Pandas实现函数 #

def pure_pandas(index):
    if index == 0:
        result = nycflights_flights_df.merge(nycflights_planes_df, how='left',\
    left_on='tailnum', right_on='tailnum')\
    .drop(['type', 'manufacturer','model', 'engines', 'engine'], axis=1)\
    .rename(columns = {'year_x':'year','year_y':'plane_year', 'seats':'plane_seats', \
    'speed':'plane_speed'})

    elif index == 1:
        result = nycflights_flights_df[['carrier', 'tailnum']].drop_duplicates()\
    .merge(nycflights_planes_df, how='inner', left_on='tailnum', right_on='tailnum')\
    .merge(nycflights_airlines_df, how='inner', left_on='carrier', right_on='carrier')\
    .assign(carrier = lambda x: x.pop('carrier'))\
    .sort_values(by=['tailnum','carrier'], ascending=[True,False])\
    .reset_index(drop=True)

    else:
        result = nycflights_flights_df[nycflights_flights_df['origin'] == 'EWR']\
    .merge(nycflights_weather_df[nycflights_weather_df['origin'] == 'EWR']\
    .groupby(['year','month','day'])[['temp', 'humid']]\
    .mean().reset_index()\
    .rename(columns = {'temp':'atemp','humid':'ahumid'}),\
    how='left', on=['year', 'month', 'day'])

    return result

性能测试代码 #

Query_Type_list = ['Group By & Sort', 'Join', 'Join',\
                   'Join & Group By']

def compare_queries(Query_Type_list, queries, no_of_times):
    
    if no_of_times == 10:
        print('Inside:',no_of_times)
        for i, qtl in zip(range(3),Query_Type_list):
            task = 'Task' + str(i+1)
            index.append(task)
            query_type_ls.append(qtl)

            pure_pandas_time = timeit.timeit(lambda: pure_pandas(i), number=no_of_times)
            pandas_query_time.append(pure_pandas_time)
            # print(pure_pandas_time)

        timeit_df = pd.DataFrame(list(zip(index,Query_Type_list,pandas_query_time)),\
                                 columns=['Query No','Query Type','Pure Pandas Time'])
        timeit_df['RUN TIMES'] = '10T'
        timeit_df['Key'] = timeit_df['RUN TIMES'] + timeit_df['Query No']
        return timeit_df
        del timeit_df
    
    elif no_of_times == 50:
        print('Inside:',no_of_times)
        for i, qtl in zip(range(3),Query_Type_list):
            task = 'Task' + str(i+1)
            index.append(task)
            query_type_ls.append(qtl)

            pure_pandas_time = timeit.timeit(lambda: pure_pandas(i), number=no_of_times)
            pandas_query_time.append(pure_pandas_time)
            # print(pure_pandas_time)

        timeit_df = pd.DataFrame(list(zip(index,Query_Type_list,pandas_query_time)),\
                                 columns=['Query No','Query Type','Pure Pandas Time'])
        timeit_df['RUN TIMES'] = '50T'
        timeit_df['Key'] = timeit_df['RUN TIMES'] + timeit_df['Query No']
        return timeit_df
        del timeit_df
    else:
        print('Do Nothing:')

# 执行10次和50次测试的调用：
timeit_df_10 = compare_queries(Query_Type_list, 10)

timeit_df_50 = compare_queries(Query_Type_list, 50)

问题根源 #

1. 参数传递错误：compare_queries定义了三个参数，但调用时只传了两个，导致no_of_times参数未被正确赋值，实际未进入10/50次的分支逻辑（或逻辑执行错误）。
2. 全局变量污染：index、query_type_ls、pandas_query_time为全局列表，多次调用函数时会持续追加数据，导致结果重复。
3. 闭包变量捕获问题：循环中定义的lambda: pure_pandas(i)会共享同一个i变量，实际执行时可能都使用循环结束后的i值，导致测试的是同一个查询。
4. 代码重复：10次和50次测试的逻辑几乎完全相同，仅RUN TIMES标识不同，未做抽象。

优化后的解决方案 #

优化后的测试函数 #

import timeit
import pandas as pd

Query_Type_list = ['Group By & Sort', 'Join', 'Join', 'Join & Group By']

def compare_queries(query_types, no_of_times):
    # 初始化内部临时列表，避免全局变量污染
    index = []
    query_type_ls = []
    pandas_query_time = []
    
    run_times_tag = f"{no_of_times}T"
    
    print(f'Inside: {no_of_times}')
    for i, qtl in zip(range(3), query_types):
        task = f'Task{i+1}'
        index.append(task)
        query_type_ls.append(qtl)
        
        # 使用默认参数捕获当前循环的i值，避免闭包问题
        pure_pandas_time = timeit.timeit(lambda i=i: pure_pandas(i), number=no_of_times)
        pandas_query_time.append(pure_pandas_time)
    
    # 统一生成结果DataFrame
    timeit_df = pd.DataFrame({
        'Query No': index,
        'Query Type': query_type_ls,
        'Pure Pandas Time': pandas_query_time
    })
    timeit_df['RUN TIMES'] = run_times_tag
    timeit_df['Key'] = timeit_df['RUN TIMES'] + timeit_df['Query No']
    
    return timeit_df

# 正确调用测试函数
timeit_df_10 = compare_queries(Query_Type_list, 10)
timeit_df_50 = compare_queries(Query_Type_list, 50)

优化点说明 #

1. 修正参数传递：移除未使用的queries参数，调整函数参数为(query_types, no_of_times)，调用时直接传递次数。
2. 内部临时变量：在函数内部初始化存储结果的列表，每次调用都会重新生成，避免全局变量的数据累计。
3. 修复闭包问题：使用lambda i=i: pure_pandas(i)，通过默认参数捕获当前循环的i值，确保每个lambda对应正确的查询索引。
4. 代码抽象：将重复的DataFrame生成逻辑统一处理，仅根据no_of_times生成对应的标识，消除代码重复。

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