rosefun
rosefun

Reputation: 1857

python: How to improve the speed of merging two DataFrames?

I have two DataFrames called a and b. All columns of a should be matched with the column keyB of DataFrame b. I define a match function as follows to achieve it, but the speed of the code is low as the DataFrame of a and b actually have a larger shape. So now I want to improve the speed of matching two 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

Hopefully for advice of improving the performance of the code.

Upvotes: 1

Views: 157

Answers (1)

jezrael
jezrael

Reputation: 862641

You can use map by Series created by b in loop for append each Series to list and last concat all together:

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

Another solution with merge in list comprehension and 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

Upvotes: 1

Related Questions