Pandas，基於列值的唯一子集追加列

Question

我有一個包含多行的數據框。 我正在使用從自定義函數生成的數據附加一列，如下所示：

import numpy

df['new_column'] = numpy.vectorize(fx)(df['col_a'], df['col_b'])
# takes 180964.377 ms

它工作正常，我想做的是加快速度。 實際上只有一小部分col_a和col_b的獨特組合。 許多迭代是多余的。 我在想也許pandas會自己解決這個問題，但我認為事實並非如此。 考慮一下：

print len(df.index) #prints 127255
df_unique = df.copy().drop_duplicates(['col_a', 'col_b'])
print len(df_unique.index) #prints 9834

我還通過運行以下命令使自己相信可能的加速：

df_unique['new_column'] = numpy.vectorize(fx)(df_unique['col_a'], df_unique['col_b'])
# takes 14611.357 ms

由於有很多冗余數據，我想要做的是更新大數據幀（ df 127255 行），但只需要運行fx函數的次數最少（9834 次）。 這是因為col_a和col_b的所有重復行。 當然，這意味着df中會有多行對col_a和col_b具有相同的值，但是沒關系， df的其他列不同並使每一行都是唯一的。

在我創建一個普通的 for 循環迭代來循環遍歷df_unique數據幀並對df進行條件更新之前，我想問一下是否有一種更“pythonic”的簡潔方式來進行這種更新。 非常感謝。

** 更新 **

我創建了上面提到的簡單 for 循環，如下所示：

df = ...
df_unique = df.copy().drop_duplicates(['col_a', 'col_b'])
df_unique['new_column'] = np.vectorize(fx)(df_unique['col_a'], df_unique['col_b'])
for index, row in df_unique.iterrows():         
    df.loc[(df['col_a'] == row['col_a']) & (df['col_b'] == row['col_b']),'new_column'] = row['new_column']
# takes 165971.890

所以使用這個 for 循環可能會有輕微的性能提升，但不是我所期望的。

供參考

這是fx函數。 它查詢一個 mysql 數據庫。

def fx(d):
    exp_date = datetime.strptime(d.col_a, '%m/%d/%Y')
    if exp_date.weekday() == 5:
        exp_date -= timedelta(days=1)

    p = pandas.read_sql("select stat from table where a = '%s' and b_date = '%s';" % (d.col_a,exp_date.strftime('%Y-%m-%d')),engine)
    if len(p.index) == 0:
        return None
    else:
        return p.iloc[0].close

Answer 1

更新：

如果您可以設法將屬於table table 的三列['stat','a','b_date']到tab DF 中，那么您可以像這樣合並它：

tab = pd.read_sql('select stat,a,b_date from table', engine)
df.merge(tab, left_on=[...], right_on=[...], how='left')

舊答案：

您可以將預先計算的df_unique DF 與原始df DF 合並/加入：

df['new_column'] = df.merge(df_unique, on=['col_a','col_b'], how='left')['new_column']

Answer 2

MaxU 的答案可能已經是您想要的。 但我將展示另一種可能更快一點的方法（我沒有測量）。

我假設：

1. df[['col_a', 'col_b']]已排序，以便所有相同的條目都在連續的行中（這很重要）

2. df有一個唯一索引（如果沒有，您可以創建一些臨時唯一索引）。

我會用事實df_unique.index的一個子集df.index 。

# (keep='first' is actually default)
df_unique = df[['col_a', 'col_b']].drop_duplicates(keep='first').copy()

# You may try .apply instead of np.vectorize (I think it may be faster):
df_unique['result'] = df_unique.apply(fx, axis=1)

# Main part:
df['result'] = df_unique['result']                     # uses 2.
df['result'].fillna(method='ffill', inplace=True)      # uses 1.