熊猫合并与布尔索引

Question

I'm using pandas in Python 3.4 to identify matches between two data frames. Matches are based on strict equality except for the last column, where close matches (+/- 5) are fine.

One data frame contains many rows, and the second is just a single row in this case. The desired result is a data frame containing a subset of the first data frame which match the row, as mentioned.

I went with the concrete solution of boolean indexing first, but this took a while to chug through all of the data, so I tried out the pandas merge function. However, my implementation of merge is even slower on my test data. It runs between 2 and 4 times slower than the boolean indexing.

Here is a test run:

import pandas as pd
import random
import time

def make_lsts(lst, num, num_choices):
    choices = list(range(0,num_choices))
    [lst.append(random.choice(choices)) for i in range(0,num)]
    return lst

def old_way(test, data):
    t1 = time.time()
    tmp = data[(data.col_1 == test.col_1[0]) &
              (data.col_2 == test.col_2[0]) &
              (data.col_3 == test.col_3[0]) &
              (data.col_4 == test.col_4[0]) &
              (data.col_5 == test.col_5[0]) &
              (data.col_6 == test.col_6[0]) &
              (data.col_7 == test.col_7[0]) &
              (data.col_8 >= (test.col_8[0]-5)) &
              (data.col_8 <= (test.col_8[0]+5))]
    t2 = time.time()
    print('old time:', t2-t1)

def new_way(test, data):
    t1 = time.time()
    tmp = pd.merge(test, data, how='inner', sort=False, copy=False,
                   on=['col_1', 'col_2', 'col_3', 'col_4', 'col_5', 'col_6', 'col_7'])
    tmp = tmp[(tmp.col_8_y >= (test.col_8[0] - 5)) & (tmp.col_8_y <= (test.col_8[0] + 5))]
    t2 = time.time()
    print('new time:', t2-t1)

if __name__ == '__main__':
    t1 = time.time()
    data = pd.DataFrame({'col_1':make_lsts([], 4000000, 7),
                         'col_2':make_lsts([], 4000000, 3),
                         'col_3':make_lsts([], 4000000, 3),
                         'col_4':make_lsts([], 4000000, 5),
                         'col_5':make_lsts([], 4000000, 4),
                         'col_6':make_lsts([], 4000000, 4),
                         'col_7':make_lsts([], 4000000, 2),
                         'col_8':make_lsts([], 4000000, 20)})

    test = pd.DataFrame({'col_1':[1], 'col_2':[1], 'col_3':[1], 'col_4':[4], 'col_5':[0], 'col_6':[1], 'col_7':[0], 'col_8':[12]})
    t2 = time.time()
    old_way(test, data)
    new_way(test, data)
    print('time building data:', t2-t1)

On my most recent run I see the following:

 # old time: 0.2209608554840088
 # new time: 0.9070699214935303
 # time building data: 75.05818915367126

Note that even the new method with the merge function uses boolean indexing on the last column dealing with the range of values, but I thought the merge might be able to do the heavy lifting in the problem. This is clearly not the case since the merge on the first columns takes up almost all of the time used in the new method.

Is it possible to optimize my implementation of the merge function? (Coming from R and data.table, I spent 30 minutes unsuccessfully searching for a way to set the key in a pandas data frame.) Is this just a problem that merge isn't good at handling? Why is boolean indexing faster than merge in this example?

I don't fully understand the memory backend of these approaches, so any insight is appreciated.

Answer 1

While you can merge on any set of columns, the performance of the merge is going to be best when you are merging on indexes.

If you replace

tmp = pd.merge(test, data, how='inner', sort=False, copy=False,
               on=['col_1', 'col_2', 'col_3', 'col_4', 'col_5', 'col_6', 'col_7'])

with

cols = ['col_%i' % (i+1) for i in xrange(7)]
test.set_index(cols, inplace=True)
data.set_index(cols, inplace=True)
tmp = pd.merge(test, data, how='inner', left_index=True, right_index=True)
test.reset_index(inplace=True)
data.reset_index(inplace=True)

Does that run faster? I haven't tested it, but I think that should help...

By indexing the columns you want to merge, the DataFrame will organize the data under the hood in such a way that it knows where to finds values much more quickly than if the data is simply in ordinary columns.