pandas pd.concat横向拼接DataFrame时出现无限循环卡顿问题

问题现象 #

• 业务逻辑：遍历字典NEM中存储的所有DataFrame，为每个DF拼接wind_output中对应前缀列的行求和汇总列，单DF数据量约50万行
• 异常表现：移除pd.concat相关代码时，全循环运行耗时不足1秒；加入concat代码行后，单次循环迭代等待10分钟以上无法结束。同写法的concat语句在小数据量场景下可瞬时完成。

原问题中使用的代码如下：

from timeit import default_timer as timer
start = timer()
for s in NEM.keys():
    F = s[0]
    print(F)
    cols=NEM[s].shape[1]
    stateW = wind_output.loc[:,wind_output.columns.str.startswith(F)].copy()
    stateW["rsum"]=stateW.sum(axis=1)
    #NEM[s].insert(cols,"wsum","")
    p = NEM[s].copy()
    print("state rows are ",len(p))
    print("wind output rows are ",len(stateW))
    p["wsum"]=stateW["rsum"]
    j=pd.concat([p.copy(),stateW["rsum"]], axis=1)
    NEM[s]=j.copy()
end = timer()
print()
print("The time taken is ", end - start)

根因分析 #

90%以上的同类卡死问题，都源于pandas的索引对齐机制：
执行axis=1的按列concat操作时，pandas不会直接按物理位置拼接数据，而是默认按索引值做匹配对齐。如果两个待拼接对象的索引存在大量重复值、或索引值不完全一致，对齐过程会产生平方级的计算量，甚至触发笛卡尔积式的数据膨胀——原本50万行的数据，对齐后可能生成数亿行的临时结果，表现出来就是程序长时间无响应，类似陷入无限循环。
小数据量下即使触发索引对齐，计算量也在可接受范围内，所以同类写法在小场景下不会出现异常。

可以先执行两行代码快速验证根因：

print(p.index.equals(stateW["rsum"].index))
print(p.index.is_unique, stateW["rsum"].index.is_unique)

如果输出结果为False加两个False，就可以确认是索引不唯一、不对齐导致的卡死。

另外原代码存在逻辑冗余：执行concat之前已经通过p["wsum"]=stateW["rsum"]完成了新增列的赋值，后续concat属于重复操作，本身就会带来不必要的性能损耗。

修复方案 #

按优先级从高到低选择即可：

• 删除冗余逻辑，用无索引赋值替代concat
  这是性能最高、最不容易出问题的写法，直接跳过pandas的索引对齐逻辑，按行位置直接赋值：

  from timeit import default_timer as timer
  start = timer()
  for s in NEM.keys():
      F = s[0]
      print(F)
      stateW = wind_output.loc[:,wind_output.columns.str.startswith(F)]
      # 转numpy数组去除索引，直接按位置赋值
      NEM[s]["wsum"] = stateW.sum(axis=1).to_numpy()
  end = timer()
  print()
  print("The time taken is ", end - start)
  

• 必须使用concat时，提前重置索引
  如果确实需要用concat实现多对象按列拼接，拼接前先把所有对象的索引重置为从0开始的连续整数，强制按位置对齐：

  p_reset = p.reset_index(drop=True)
  rsum_reset = stateW["rsum"].reset_index(drop=True)
  j = pd.concat([p_reset, rsum_reset], axis=1)
  

• 增加前置校验
  拼接前先确认两个待拼接对象的行数完全一致，从源头避免数据膨胀问题。

避坑提示 #

• 所有axis=1的列拼接、列赋值操作，只要确认两份数据是按行顺序一一对应的，优先转numpy数组或者重置索引后再操作，不要直接携带原始索引触发自动对齐。
• 数据量超过10万行时，尽量避免不必要的.copy()操作，减少内存占用和拷贝耗时。

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