python - 为什么性能会随着数据帧的大小而下降？

Question

我正在使用一个相对较大的数据集(大约 500 万个观测值，由大约 5.5 万家公司组成)。

我需要为每个公司运行 60 个月滚动窗口的 OLS 回归。我注意到当我运行以下代码时，性能非常慢:

for idx, sub_df in master_df.groupby("firm_id"):
    # OLS code

然而，当我第一次将我的数据帧拆分为大约 5.5k dfs 然后迭代每个 dfs 时，性能得到了显着提高。

grouped_df = master_df.groupby("firm_id")
df_list = [group for group in grouped_df]

for df in df_list:
    my_df = df[1]
    # OLS code

我说的是在第一个版本中完成 1-2 周的时间(24/7)，而最高需要 8-9 小时。

谁能解释一下 为什么将主 df 拆分为 N 个较小的 df，然后迭代每个较小的 df 比在主 df 中迭代相同数量的组的性能更好？

非常感谢!

Answer 1

我无法重现您的观察。下面是一些生成数据然后分别对直接和间接方法计时的代码。在这两种情况下所花费的时间都非常相似。
您是否有可能在两次运行之间不小心按组键对数据框进行了排序？按组键排序会导致运行时间明显不同。
否则，我开始认为您的代码中可能存在其他一些差异。如果您能发布完整的代码，那就太好了。

import numpy as np
import pandas as pd
from datetime import datetime

def generate_data():
  ''' returns a Pandas DF with columns 'firm_id' and 'score' '''
  # configuration
  np.random.seed(22)
  num_groups = 50000 # number of distinct groups in the DF
  mean_group_length = 200 # how many records per group?
  cov_group_length = 0.10 # throw in some variability in the num records per group
  
  # simulate group lengths
  stdv_group_length = mean_group_length * cov_group_length
  group_lengths = np.random.normal(
    loc=mean_group_length,
    scale=stdv_group_length,
    size=(num_groups,)).astype(int)
  group_lengths[group_lengths <= 0] = mean_group_length
  
  # final length of DF
  total_length = sum(group_lengths)
  
  # compute entries for group key column
  firm_id_list = []
  for i, l in enumerate(group_lengths):
    firm_id_list.extend([(i + 1)] * l)
    
  # construct the DF; data column is 'score' populated with Numpy's U[0, 1)
  result_df = pd.DataFrame(data={
      'firm_id': firm_id_list,
      'score': np.random.rand(total_length)
    })
    
  # Optionally, shuffle or sort the DF by group keys
  
  # ALTERNATIVE 1: (badly) unsorted df
  result_df = result_df.sample(frac=1, random_state=13).reset_index(drop=True)
  # ALTERNATIVE 2: sort by group key
  # result_df.sort_values(by='firm_id', inplace=True)
  return result_df

def time_method(df, method):
  ''' time 'method' with 'df' as its argument '''
  t_start = datetime.now()
  method(df)
  t_final = datetime.now()
  delta_t = t_final - t_start
  print(f"Method '{method.__name__}' took {delta_t}.")
  return

def process_direct(df):
  ''' direct for-loop over groupby object '''
  for group, df in df.groupby('firm_id'):
    m = df.score.mean()
    s = df.score.std()
  return

def process_indirect(df):
  ''' indirect method: generate groups first as list and then loop over list '''
  grouped_df = df.groupby('firm_id')
  group_list = [pair for pair in grouped_df]
  for pair in group_list:
    m = pair[1].score.mean()
    s = pair[1].score.std()

df = generate_data()
time_method(df, process_direct)
time_method(df, process_indirect)