python - pandas 的高性能笛卡尔积(CROSS JOIN)

Question

The contents of this post were originally meant to be a part of Pandas Merging 101, but due to the nature and size of the content required to fully do justice to this topic, it has been moved to its own QnA.

给定两个简单的 DataFrame；

left = pd.DataFrame({'col1' : ['A', 'B', 'C'], 'col2' : [1, 2, 3]})
right = pd.DataFrame({'col1' : ['X', 'Y', 'Z'], 'col2' : [20, 30, 50]})

left

  col1  col2
0    A     1
1    B     2
2    C     3

right

  col1  col2
0    X    20
1    Y    30
2    Z    50

这些帧的叉积可以计算出来，看起来像这样:

A       1      X      20
A       1      Y      30
A       1      Z      50
B       2      X      20
B       2      Y      30
B       2      Z      50
C       3      X      20
C       3      Y      30
C       3      Z      50

计算此结果的最高效方法是什么？

Answer 1

让我们从建立一个基准开始。解决此问题的最简单方法是使用临时“关键”列:

Pandas <= 1.1.X

def cartesian_product_basic(left, right):
    return (
       left.assign(key=1).merge(right.assign(key=1), on='key').drop('key', 1))

cartesian_product_basic(left, right)

Pandas >= 1.2

left.merge(right, how="cross") # implements the technique above

  col1_x  col2_x col1_y  col2_y
0      A       1      X      20
1      A       1      Y      30
2      A       1      Z      50
3      B       2      X      20
4      B       2      Y      30
5      B       2      Z      50
6      C       3      X      20
7      C       3      Y      30
8      C       3      Z      50

其工作原理是为两个 DataFrame 分配一个具有相同值(例如 1)的临时“键”列。 merge 然后在 "key"上执行多对多 JOIN。

虽然多对多 JOIN 技巧适用于大小合理的 DataFrame，但您会发现在较大数据上的性能相对较低。

更快的实现需要 NumPy。这里有一些著名的NumPy implementations of 1D cartesian product .我们可以在其中一些高性能解决方案的基础上获得我们想要的输出。然而，我最喜欢的是@senderle 的第一个实现。

def cartesian_product(*arrays):
    la = len(arrays)
    dtype = np.result_type(*arrays)
    arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype)
    for i, a in enumerate(np.ix_(*arrays)):
        arr[...,i] = a
    return arr.reshape(-1, la)  

泛化:唯一或非唯一索引数据帧上的交叉连接

Disclaimer
These solutions are optimised for DataFrames with non-mixed scalar dtypes. If dealing with mixed dtypes, use at yourown risk!

这个技巧适用于任何类型的 DataFrame。我们使用上述 cartesian_product 计算 DataFrames 的数字索引的笛卡尔积，使用它来重新索引 DataFrames，并且

def cartesian_product_generalized(left, right):
    la, lb = len(left), len(right)
    idx = cartesian_product(np.ogrid[:la], np.ogrid[:lb])
    return pd.DataFrame(
        np.column_stack([left.values[idx[:,0]], right.values[idx[:,1]]]))

cartesian_product_generalized(left, right)

   0  1  2   3
0  A  1  X  20
1  A  1  Y  30
2  A  1  Z  50
3  B  2  X  20
4  B  2  Y  30
5  B  2  Z  50
6  C  3  X  20
7  C  3  Y  30
8  C  3  Z  50

np.array_equal(cartesian_product_generalized(left, right),
               cartesian_product_basic(left, right))
True

而且，按照类似的思路，

left2 = left.copy()
left2.index = ['s1', 's2', 's1']

right2 = right.copy()
right2.index = ['x', 'y', 'y']
    

left2
   col1  col2
s1    A     1
s2    B     2
s1    C     3

right2
  col1  col2
x    X    20
y    Y    30
y    Z    50

np.array_equal(cartesian_product_generalized(left, right),
               cartesian_product_basic(left2, right2))
True

此解决方案可以推广到多个 DataFrame。例如，

def cartesian_product_multi(*dfs):
    idx = cartesian_product(*[np.ogrid[:len(df)] for df in dfs])
    return pd.DataFrame(
        np.column_stack([df.values[idx[:,i]] for i,df in enumerate(dfs)]))

cartesian_product_multi(*[left, right, left]).head()

   0  1  2   3  4  5
0  A  1  X  20  A  1
1  A  1  X  20  B  2
2  A  1  X  20  C  3
3  A  1  X  20  D  4
4  A  1  Y  30  A  1

进一步简化

在处理两个 DataFrame 时，可以使用不涉及@senderle 的cartesian_product 的更简单的解决方案。使用 np.broadcast_arrays，我们可以达到几乎相同水平的性能。

def cartesian_product_simplified(left, right):
    la, lb = len(left), len(right)
    ia2, ib2 = np.broadcast_arrays(*np.ogrid[:la,:lb])

    return pd.DataFrame(
        np.column_stack([left.values[ia2.ravel()], right.values[ib2.ravel()]]))

np.array_equal(cartesian_product_simplified(left, right),
               cartesian_product_basic(left2, right2))
True

---

性能比较

在一些人为的具有唯一索引的 DataFrame 上对这些解决方案进行基准测试，我们有

请注意，时间可能会根据您的设置、数据和选择的 cartesian_product 辅助函数(如适用)而有所不同。

性能基准代码
这是计时脚本。这里调用的所有函数都在上面定义。

from timeit import timeit
import pandas as pd
import matplotlib.pyplot as plt

res = pd.DataFrame(
       index=['cartesian_product_basic', 'cartesian_product_generalized', 
              'cartesian_product_multi', 'cartesian_product_simplified'],
       columns=[1, 10, 50, 100, 200, 300, 400, 500, 600, 800, 1000, 2000],
       dtype=float
)

for f in res.index: 
    for c in res.columns:
        # print(f,c)
        left2 = pd.concat([left] * c, ignore_index=True)
        right2 = pd.concat([right] * c, ignore_index=True)
        stmt = '{}(left2, right2)'.format(f)
        setp = 'from __main__ import left2, right2, {}'.format(f)
        res.at[f, c] = timeit(stmt, setp, number=5)

ax = res.div(res.min()).T.plot(loglog=True) 
ax.set_xlabel("N"); 
ax.set_ylabel("time (relative)");

plt.show()

---

---

继续阅读

跳到 Pandas Merging 101 中的其他主题以继续学习:

• Merging basics - basic types of joins

• Index-based joins

• Generalizing to multiple DataFrames

• Cross join ^*

_{*你在这里}