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

由于有很多冗余数据，我想做的是更新大数据帧（df127255 行），但只需要运行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 表的['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.