高效查找pandas DataFrame中匹配的行（基于内容）

Question

我正在写一些测试，使用Pandas的DataFrame来存放一个很大的数据集，大约有60万行和10列。我从源数据中提取了10行随机数据（用Stata工具），现在我想写一个测试，看看这些行是否在我的测试数据集中。

这里有个简单的例子：

np.random.seed(2)
raw_data = pd.DataFrame(np.random.rand(5,3), columns=['one', 'two', 'three'])
random_sample = raw_data.ix[1]

在这个例子中，raw_data是：

而random_sample是为了确保能找到匹配的行，它是：

目前我写的代码是：

for idx, row in raw_data.iterrows():
    if random_sample.equals(row):
        print "match"
        break

这个方法可以用，但在处理这么大的数据集时，速度非常慢。有没有更有效的方法来检查整行数据是否在DataFrame中？

顺便提一下：我的例子还需要能够比较np.NaN的相等性，这就是我使用equals()方法的原因。

Answer 1

我想到的最好办法是采用过滤的方法，这种方法在处理大数据集时效果不错，可以减少很多比较的操作：

tmp = raw_data    
for idx, val in random_sample.iteritems():
    try:
        if np.isnan(val):
            continue
    except:
        pass
    tmp = tmp[tmp[idx] == val]
if len(tmp) == 1: print "match"

注意：对于上面的小数据集，这种方法实际上会慢一些。但在处理大数据集时，这种方法比基本的迭代快大约9倍。

Answer 2

equals 这个方法似乎不能自动处理多个数据，但我们可以手动进行相等比较：

>>> df = pd.DataFrame(np.random.rand(600000, 10))
>>> sample = df.iloc[-1]
>>> %timeit df[((df == sample) | (df.isnull() & sample.isnull())).all(1)]
1 loops, best of 3: 231 ms per loop
>>> df[((df == sample) | (df.isnull() & sample.isnull())).all(1)]
              0         1         2         3         4         5         6  \
599999  0.07832  0.064828  0.502513  0.851816  0.976464  0.761231  0.275242   

               7        8         9  
599999  0.426393  0.91632  0.569807  

对我来说，这种方法比逐行比较要快得多（后者需要超过30秒）。

不过，由于我们的数据行很多而列相对较少，我们可以逐列处理，这样通常可以大大减少需要查看的行数。例如，像这样：

def finder(df, row):
    for col in df:
        df = df.loc[(df[col] == row[col]) | (df[col].isnull() & pd.isnull(row[col]))]
    return df

给我的结果是：

>>> %timeit finder(df, sample)
10 loops, best of 3: 35.2 ms per loop

这个方法快了大约十倍，因为在处理完第一列后，只剩下一行需要查看。

（我记得以前有一个更简洁的方法来做到这一点，但现在想不起来了。）