首页 > 解决方案 > 如何有效地匹配关于多列的两个熊猫数据框?

问题描述

我有一个有 8M 行和 7 列的数据框 df1。其中一列 ('ID') 是样本 ID,A 列是二进制变量,其他 5 列是浮点值。

df1 = pd.DataFrame(np.random.random_sample((df_1_size, 5)), columns=list('BCDEF'))
df1['A'] = np.random.randint(1,3,size=df_1_size)
df1['ID'] = random.sample(range(0, df_1_size), df_1_size)

df2 具有 300K 和相同的列:

df2 = pd.DataFrame(np.random.random_sample((df_2_size, 5)), columns=list('BCDEF'))
df2['A'] = np.random.randint(1,3,size=df_2_size)
df2['ID'] = random.sample(range(0, df_2_size), df_2_size)

我需要匹配两个数据帧并从 df1 中找到 300K*k 个样本,这些样本与 df2 中关于 A、B、C 和 D 列的样本相似,无需替换。

我编写了一个代码,其中对于 df1 中的每个样本(例如 s_i),我首先根据二进制列 A 列过滤 df2,然后使用然后找到 s_i 与 df2 中 B、C 和 D 列的所有样本之间的余弦相似度,scipy.spatial.distance.cdist然后对距离进行排序并选择最接近的(或 top-k 最接近的)并将 df2 中的选定样本标记为已使用。但是,这个算法需要 300K * 8M t 才能完成,实际上它已经运行了两天,仍然没有完成。这是我的(不是那么有效)脚本:

df_1_size = 8000000
df_2_size = 300000
topk=1
df1 = pd.DataFrame(np.random.random_sample((df_1_size, 5)), columns=list('BCDEF'))
df1['A'] = np.random.randint(1,3,size=df_1_size)
df1['ID'] = random.sample(range(0, df_1_size), df_1_size)

df2 = pd.DataFrame(np.random.random_sample((df_2_size, 5)), columns=list('BCDEF'))
df2['A'] = np.random.randint(1,3,size=df_2_size)
df2['ID'] = random.sample(range(0, df_2_size), df_2_size)

match_based_on = ['B', 'C', 'D']

df1['MATCHED'] = 0
for index, row in df2.iterrows():
    df1_filtered = df1[(df1['A'] == row['A']) & (df1['MATCHED'] == 0)]
    if len(df1_filtered) == 0:
        similarities = cdist(np.reshape(row[match_based_on].values, (1,len(match_based_on))), np.reshape(df1[df1['MATCHED'] == 0][match_based_on].values, (-1,len(match_based_on))), metric='cosine')
        matched_samples = df1[df1['MATCHED'] == 0].iloc[similarities[0].argsort()[:topk]]
        df1.loc[df1['ID'].isin(matched_samples['ID'].values),'MATCHED'] = 1
    else:
        similarities = cdist(np.reshape(row[match_based_on].values, (1,len(match_based_on))), np.reshape(df1_filtered[match_based_on].values, (-1,len(match_based_on))), metric='cosine')
        matched_samples = df1_filtered.iloc[similarities[0].argsort()[:topk]]
        df1.loc[df1['ID'].isin(matched_samples['ID'].values),'MATCHED'] = 1
pdb.set_trace()
df1[df1['MATCHED']==1].loc[:, df1.columns != 'MATCHED'].to_csv(df1_path[:-4]+'_matched.csv', index=False)
return 0 

非常感谢对上述解决方案或其他不同的完全不同的解决方案的任何修改。

PS。如果这有助于更容易解决,您可以忽略基于 A 列的匹配。

标签: pythonpython-3.xpandascosine-similarity

解决方案


这可能不是您问题的完整解决方案,但它的工作topk = 1速度比原来的快 10 倍。

for a_val in df1['A'].unique():
    np1 = df1[df1['A'] == a_val][match_based_on].values
    np2 = df2[df2['A'] == a_val][match_based_on].values
    similarities = cdist(np2, np1, metric='cosine')
    df1_idx = np.argpartition(similarities, kth=topk, axis=1)[:,:topk]
    
    ind = df1[df1['A'] == a_val].iloc[df1_idx.flatten()].index
    df1.loc[ind, 'MATCHED_ID'] = df2[df2['A'] == a_val].ID.values
    df1.loc[df1['MATCHED_ID'] != "", 'MATCHED'] = 1
df1[df1['MATCHED_ID'] != ""].head()

两个主要思想:

  • 我们不应该遍历 DataFrame 中的行。所有方法通常都支持 n 维数组,例如cdist, argsort
  • 如果您需要 N 个值中的前 k 个最大/最小的值,请使用argpartition而不是。对 O(Nlog(N)) 的所有值进行排序,对于 O(N),仅选择最大 k 和最小 N - k 之间的拆分。如果您需要 k 个项目的排序列表,请按 O(klog(k)) 的值对 argpartitioned (idx, value) 对进行排序。argsortargsortargpartition

推荐阅读