Pandas `.loc` 多次赋值导致合理的减速

Question

我有一些代码，如果我不进行多重分配，而不是跨多行分配，例如

快速：

onset = pitch_df.loc[idx, 'onset_time']
dur = pitch_df.loc[idx, 'duration']

慢：

onset, dur = pitch_df.loc[idx, ['onset_time', 'duration']]

这是否有明显的原因，或者是一种更“熊猫”的方式来做我正在做的事情。我想在这里分配以使我的代码更具可读性（即我不想到处写.loc[...]）。

这是一个最小的工作示例（此处加速 4 倍）：

import pandas as pd
import numpy as np
from timeit import timeit

df = pd.DataFrame(
    {'onset_time': [0, 0, 1, 2, 3, 4], 
     'pitch': [61, 60, 60, 61, 60, 60],
     'duration': [4, 1, 1, 0.5, 0.5, 2]}
).sort_values(['onset_time', 'pitch']).reset_index(drop=True)

def foo():
    for pitch, pitch_df in df.groupby('pitch'):
        for iloc in range(len(pitch_df)):
            idx = pitch_df.index[iloc]
            onset = pitch_df.loc[idx, 'onset_time']
            dur = pitch_df.loc[idx, 'duration']
            note_off = onset + dur

def bar():
    for pitch, pitch_df in df.groupby('pitch'):
        for iloc in range(len(pitch_df)):
            idx = pitch_df.index[iloc]
            onset, dur = pitch_df.loc[idx, ['onset_time', 'duration']]
            note_off = onset + dur

print(f'foo time: {timeit(foo, number=100)}')
print(f'bar time: {timeit(bar, number=100)}')

下面的图片便于阅读。

Answer 1

正如 Poolka 在对您的问题的评论中提到的那样，如果您想要标量访问，.at 的开销较小。我不是 python 专家，但这里有一个可能对你有用的解决方案：

def foo2():
    for pitch, pitch_df in df.groupby('pitch'):
        for iloc in range(len(pitch_df)):
            idx = pitch_df.index[iloc]
            onset, dur = (pitch_df.at[idx, x] for x in ('onset_time', 'duration'))
            note_off = onset + dur

foo time: 0.12590176300000167
bar time: 0.47044453300077294
foo2 time: 0.12269815599938738