pandas DataFrame 中行的高效成对比较

Question

我目前正在处理一个较小的数据集（大约 900 万行）。不幸的是，大多数条目都是字符串，即使强制类别，框架在内存中也只有几 GB。

我想做的是将每一行与其他行进行比较，并对内容进行直接比较。例如，给定

   A   B     C      D
0 cat blue  old Saturday
1 dog red   old Saturday

我想计算

      d_A   d_B   d_C   d_D
0, 0  True  True  True  True
0, 1  False False True  True
1, 0  False False True  True
1, 1  True  True  True  True

显然，组合爆炸将排除每个记录与其他记录的比较。因此我们可以通过应用 groupby 来使用阻塞，比如在 A 列上。

我的问题是，有没有一种方法可以在 pandas 或 dask 中执行此操作，比以下序列更快：

1. 按索引分组
2. 将每个组外连接到自身以生成对
3. dataframe.apply 比较函数在每行对上

作为参考，假设我可以使用大量的核心（数百个）和大约 200G 的内存。

Answer 1

解决方案是使用 numpy 代替步骤 3)。虽然我们无法为每一行创建外连接，但我们可以按 A 列中的值进行分组，并为外连接创建更小的组。

技巧是使用numpy.equal.outer(df1, df2).ravel()当数据帧以这种方式作为输入传递给 numpy 函数时，结果是更快（至少 30 倍）的矢量化结果。例如：

>>> df = pd.DataFrame
   A   B     C      D
0 cat blue  old Saturday
1 dog red   old Saturday

>>> result = pd.DataFrame(columns=["A", "B", "C", "D"], 
                            index=pd.MultiIndex.from_product([df.index, df.index]))
>>> result["A"] = np.equal.outer(df["A"], df["A"]).ravel()
>>> result
        A     B     C     D
0, 0  True   NaN   NaN   NaN  
0, 1  False  NaN   NaN   NaN  
1, 0  False  NaN   NaN   NaN  
1, 1  True   NaN   NaN   NaN  

您可以对每一列重复，或者只是通过按列应用结果来自动化该过程。