使用pandas实现高效的笛卡尔积(CROSS JOIN)

78

这篇文章的内容原本是Pandas合并101的一部分,但由于完整地展现这个主题所需的内容的性质和规模,它被移到自己的问答页面上。

给定两个简单的数据框:

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

使用哪种方法可以最有效地计算出这个结果?


1
你想在 Github 上分享你的意见吗?我认为在 pandas 中添加 cross join 函数非常好,可以匹配 SQL 中的所有 join 函数。https://github.com/pandas-dev/pandas/issues/5401 - BENY
5个回答

99

首先,让我们建立一个基准。解决这个问题最简单的方法是使用临时的“key”列:

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 然后在“键”上执行多对多连接。

虽然这个多对多连接技巧适用于相当大小的DataFrame,但在更大的数据上性能会相对较低。

更快速的实现需要使用NumPy。以下是一些著名的NumPy 1D笛卡尔积实现。我们可以基于这些高效的解决方案来得到我们所需的输出。然而,我的最爱是@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)  

泛化:在唯一或非唯一索引的数据帧上执行CROSS JOIN

免责声明
这些解决方案是针对非混合标量数据类型的数据帧进行优化的。如果处理混合数据类型,请自行承担风险!

此技巧适用于任何类型的数据帧。我们使用上述cartesian_product计算数据帧的数字索引的笛卡尔积,将其用于重新索引数据帧,并且

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
这个解决方案可以推广到多个数据框。例如,

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

更进一步的简化

在处理仅限于两个DataFrames时,可以使用np.broadcast_arrays来实现几乎相同水平的性能,而不需要涉及@senderle的cartesian_product

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

性能比较

在一些带有唯一索引的构想数据帧上进行基准测试,我们得到:

enter image description here

请注意,时间可能因设置、数据和适用的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合并101中的其他主题以继续学习:

* 您在此处


1
为什么列名会变成整数?我试图重命名它们,使用.rename()函数运行后,整数仍然存在。 - OverflowingTheGlass
不行...更加密集了 - 我在整数周围加上引号 - 谢谢 - OverflowingTheGlass
另一个问题。我正在使用cartesian_product_simplified,当我尝试将一个有50K行的df与一个有30K行的df连接时,我(可以预见地)会耗尽内存。有什么建议可以解决内存问题吗? - OverflowingTheGlass
@CameronTaylor,其他的cartesian_product_*函数也会抛出内存错误吗?我猜你可以在这里使用cartesian_product_multi。 - cs95
FYI,我的朋友,https://pandas.pydata.org/docs/reference/api/pandas.merge.html,现在已经有交叉了~ :-) - BENY
显示剩余5条评论

24

在pandas 1.2.0之后,merge现在有了选项cross

left.merge(right, how='cross')
使用 itertoolsproduct 创建数据帧中的值。
import itertools
l=list(itertools.product(left.values.tolist(),right.values.tolist()))
pd.DataFrame(list(map(lambda x : sum(x,[]),l)))
   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

6
这里有一种使用三次 concat 的方法。
m = pd.concat([pd.concat([left]*len(right)).sort_index().reset_index(drop=True),
       pd.concat([right]*len(left)).reset_index(drop=True) ], 1)

    col1  col2 col1  col2
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

enter image description here


0
一种选择是使用pyjanitor中的expand_grid函数:
# pip install pyjanitor
import pandas as pd
import janitor as jn

others = {'left':left, 'right':right}

jn.expand_grid(others = others)

  left      right
  col1 col2  col1 col2
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

-1
我认为最简单的方法是向每个数据框添加一个虚拟列,对其进行内部合并,然后从结果笛卡尔数据帧中删除该虚拟列。
left['dummy'] = 'a'
right['dummy'] = 'a'

cartesian = left.merge(right, how='inner', on='dummy')

del cartesian['dummy']

这个问题已经在被接受的答案中讨论过了。但是现在,left.merge(right, how="cross")已经可以在不需要第二列的情况下完成此操作。 - cs95
不知何故,跨平台对我没起作用。可能是版本问题。 - abhygag

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接