分组切片DataFrame后列异常，求绘制Seaborn热力图的解决方法

问题背景 #

我有一个记录内核模块延迟数据的DataFrame，各模块的time数据长度在1000~3000之间，需要将每个模块的数据切片为0-999的等长数据（共1000条）。

原始DataFrame结构如下：

time, module_name, latency
0, module1, 268
1, module1, 300
...
999, module1, 300
0, module2, 234
1, module2, 345
...
3000, module2, 345

尝试的代码与问题 #

我用以下代码对每个模块切片：

trace1000 = df1.groupby('module_name').apply(lambda x: x.iloc[0:999])

但得到的DataFrame出现重复的module_name列和无名索引列，结构如下：

module_name, , module_name, latency
module1, 2000, module1, 268
module1, 2001, module1, 300
...
module2, 9085, module1, 234
module2, 9086, module1, 345
...

尝试用heat_df = trace1000[["module_name","latency"]]选择列也失败了。

目标需求 #

最终要绘制Seaborn热力图：

• x轴：time（范围0-999，对应1000个点）
• y轴：module_name（共30个模块）
• 热力值：latency（范围100~900）

解决方案 #

• 修复分组切片后的索引与列问题
  使用groupby的apply后会生成多层索引，且可能保留原列，解决方法是在切片后重置索引并丢弃原索引：

  # 切片每个模块的前1000条数据（0-999），并重置索引
  trace1000 = df1.groupby('module_name').apply(lambda x: x.iloc[:1000].reset_index(drop=True))
  # 移除多层索引，将groupby的索引转为列并重置整体索引
  trace1000 = trace1000.reset_index(level=0).reset_index(drop=True)
  

  或者更简洁的方式，使用groupby结合head（如果每个模块的time是连续从0开始的话）：

  # 直接取每个模块的前1000条数据
  trace1000 = df1.groupby('module_name').head(1000)
  

• 整理数据为热力图所需格式
  Seaborn的热力图需要宽表格式（行是模块，列是time，值是latency），所以需要用pivot转换：

  import seaborn as sns
  import matplotlib.pyplot as plt
  
  # 转换为宽表
  heatmap_data = trace1000.pivot(index='module_name', columns='time', values='latency')
  
  # 绘制热力图
  plt.figure(figsize=(12, 8))
  sns.heatmap(heatmap_data, cmap='viridis', vmin=100, vmax=900)
  plt.xlabel('Time')
  plt.ylabel('Module Name')
  plt.title('Kernel Module Latency Heatmap')
  plt.show()
  

• 关键说明

  • groupby.apply后产生的多层索引是导致列重复和选择失败的原因，通过reset_index可以清理索引结构。
  • pivot是将长表转换为热力图所需宽表的关键步骤，确保每个模块对应一行，每个time对应一列，latency填充对应位置。

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