line_profiler无法细分pandas assign内各赋值行耗时的问题咨询

问题现象 #

使用line_profiler的%lprun魔法命令开展Python代码行级性能剖析时，跨越多行的pandas.DataFrame.assign多列赋值场景会出现统计结果不准确的问题：工具将assign方法的整体执行耗时几乎全部统计到最后一个参数所在行，其余列赋值行统计到的耗时占比均为0，无法拆分得到每个列赋值逻辑（含常量赋值、lambda表达式计算逻辑）的真实单独耗时，无法定位assign内部的实际性能瓶颈。

复现代码与运行结果 #

测试代码：

import pandas as pd

def test_lprun():
    data = {'Name':['Tom', 'Brad', 'Kyle', 'Jerry'],
        'Age':[20, 21, 19, 18],
        'Height' : [6.1, 5.9, 6.0, 6.1]
        }
    df = pd.DataFrame(data)
    df=df.assign(A=123,
                 B=lambda x:x.Age+x.Height,
                 C=lambda x:x.Name.str.upper(),
                 D=lambda x:x.Name.str.lower()
                 )
    return df

IPython中执行%lprun -f test_lprun test_lprun()的输出：

Timer unit: 1e-07 s

Total time: 0.0044901 s
File: <ipython-input-7-eaf21639fb5f>
Function: test_lprun at line 1

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def test_lprun():
     2         1         21.0     21.0      0.0      data = {'Name':['Tom', 'Brad', 'Kyle', 'Jerry'],
     3         1         13.0     13.0      0.0          'Age':[20, 21, 19, 18],
     4         1         15.0     15.0      0.0          'Height' : [6.1, 5.9, 6.0, 6.1]
     5                                                   }
     6         1       8651.0   8651.0     19.3      df = pd.DataFrame(data)
     7         1         19.0     19.0      0.0      df=df.assign(A=123,
     8         1         11.0     11.0      0.0                   B=lambda x:x.Age+x.Height,
     9         1         10.0     10.0      0.0                   C=lambda x:x.Name.str.upper(),
    10         1      36147.0  36147.0     80.5                   D=lambda x:x.Name.str.lower()
    11                                                            )
    12         1         14.0     14.0      0.0      return df

期望效果 #

line_profiler可像统计普通独立代码行一样，准确拆分统计assign方法内每个列赋值行的命中次数、实际耗时、耗时占比数据（如各赋值行分别对应10%、30%的时间占比），满足细粒度性能定位需求。

问题根因 #

这个统计偏差不是工具bug，是Python语法执行逻辑和line_profiler的统计规则共同导致的：

• line_profiler的统计粒度是Python解释器执行的独立语句/表达式行，只有触发实际计算的行才会累计耗时
• 跨多行书写的df.assign(...)本质是单个函数调用表达式，换行仅为代码排版格式：前几行的参数（如lambda定义、常量）仅在构造传入assign的参数对象，这个过程几乎不耗时；所有列的实际计算逻辑都在assign函数内部执行，函数返回前的全部耗时，会被Python解释器标记到整个调用表达式的最后一行，也就出现了最后一个赋值行占绝大多数耗时的统计结果。

可行解决方案 #

• 方案1：拆分assign为链式单参数调用
  把多列合并的assign拆成多次链式调用，每次assign仅传入一个列参数，每个列的计算逻辑单独占一行，line_profiler即可分别统计每一次调用的耗时，写法如下：

  def test_lprun_fix1():
      data = {'Name':['Tom', 'Brad', 'Kyle', 'Jerry'],
          'Age':[20, 21, 19, 18],
          'Height' : [6.1, 5.9, 6.0, 6.1]
          }
      df = pd.DataFrame(data)
      df = (df
            .assign(A=123)
            .assign(B=lambda x:x.Age+x.Height)
            .assign(C=lambda x:x.Name.str.upper())
            .assign(D=lambda x:x.Name.str.lower())
            )
      return df
  

  这种写法无额外依赖，统计逻辑和原生行级统计完全一致，是最通用的解法。

• 方案2：抽离列计算逻辑为独立函数加入剖析范围
  如果不想拆分assign调用，可以把每个列的计算逻辑抽成单独的具名函数，剖析时把这些函数也加入%lprun的统计列表，即可拿到每个计算逻辑的单独耗时：

  def calc_b(x):
      return x.Age + x.Height
  def calc_c(x):
      return x.Name.str.upper()
  def calc_d(x):
      return x.Name.str.lower()
  
  def test_lprun_fix2():
      data = {'Name':['Tom', 'Brad', 'Kyle', 'Jerry'],
          'Age':[20, 21, 19, 18],
          'Height' : [6.1, 5.9, 6.0, 6.1]
          }
      df = pd.DataFrame(data)
      df = df.assign(
          A=123,
          B=calc_b,
          C=calc_c,
          D=calc_d
      )
      return df
  

  执行剖析时把所有计算函数加入-f参数即可：%lprun -f test_lprun_fix2 -f calc_b -f calc_c -f calc_d test_lprun_fix2()。

• 方案3：临时替换为普通列赋值排查性能
  如果只是临时排查性能瓶颈，可以直接把assign的多列赋值改成普通的df['col'] = ...写法，每一行赋值都是独立Python语句，line_profiler可直接统计每行耗时，排查完成后再改回assign写法即可。

  注意：普通列赋值是在原DataFrame上原地修改，assign会返回新的DataFrame副本，这个差异不影响单段计算逻辑的耗时统计结果。

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