python: Как улучшить скорость слияния двух DataFrames?

Question

У меня есть два DataFrames с именами a и b. Все столбцы a должны соответствовать столбцу keyB DataFrame b. Я определяю match функцию следующим образом, чтобы достичь ее, но скорость кода низкая, как Фреймы данных a и b на самом деле имеют большую форму. Так что теперь я хочу улучшить скорость сопоставления двух DataFrame.

import pandas as pd
import time

start=time.time()
a=pd.DataFrame({'key1':[1,5,1],'key2':[1,2,11]})
b=pd.DataFrame({'keyB':[1,2,3,4,5],'other':['q','q','w','w','r']})

def match(num,a,b,col):
    aFeat=a.iloc[num:num+1]
    bFeat=b[b['keyB'].isin([a[col].loc[num]])]
    aFeat.reset_index(drop=True,inplace=True)
    bFeat.reset_index(drop=True,inplace=True)
    new=pd.concat([aFeat,bFeat],axis=1)
    return new

newb=pd.DataFrame({})
for col in ['key1','key2']:
    newa=pd.DataFrame({})
    for num in range(len(a)):
        newa=pd.concat([newa,match(num,a,b,col)],axis=0)
    newa.reset_index(drop=True,inplace=True)
    del newa[col]
    newb.reset_index(drop=True,inplace=True)
    newb=pd.concat([newb,newa],axis=1)
    newb = newb.rename(columns={'keyB': 'keyB_'+col, 'other': 'other_'+col})

print(newb)
end=time.time()
print('time:',end-start)

Input:

a    key1  key2
0     1     1
1     5     2
2     1    11

b    keyB other
0     1     q
1     2     q
2     3     w
3     4     w
4     5     r

Output:

   key2  keyB_key1 other_key1  key1  keyB_key2 other_key2
0     1          1          q     1        1.0          q
1     2          5          r     5        2.0          q
2    11          1          q     1        NaN        NaN

Used time:

time: 0.015628576278686523

Надеемся на совет по улучшению производительности кода.

Answer 1

Вы можете использовать map от Series, созданного b в цикле, для добавления каждой серии к списку и последнему concat все вместе:

s = b.set_index('keyB')['other']
print (s)
keyB
1    q
2    q
3    w
4    w
5    r
Name: other, dtype: object

dfs = []
for col in ['key1','key2']:
    dfs.append(a[col])
    val = a[col].map(s).rename('other_' + col)
    dfs.append(pd.Series(np.where(val.notnull(), a[col], np.nan), name='keyB_' + col))
    dfs.append(val)

df = pd.concat(dfs, axis=1)
print (df)
   key1  keyB_key1 other_key1  key2  keyB_key2 other_key2
0     1        1.0          q     1        1.0          q
1     5        5.0          r     2        2.0          q
2     1        1.0          q    11        NaN        NaN

Другое решение с merge в понимании списка и concat:

dfs = [b.merge(a[[col]], left_on='keyB', right_on=col)
        .rename(columns={'keyB':'keyB_'+col,'other':'other_'+col}) for col in ['key1','key2']]
df = pd.concat(dfs, axis=1)
print (df)
   keyB_key1 other_key1  key1  keyB_key2 other_key2  key2
0          1          q     1        1.0          q   1.0
1          1          q     1        2.0          q   2.0
2          5          r     5        NaN        NaN   NaN