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

发布于 2021-02-02 23:17:49

给定两个简单的DataFrames;

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

计算结果的最有效方法是什么?

关注者
0
被浏览
107
1 个回答
  • 面试哥
    面试哥 2021-02-02
    为面试而生,有面试问题,就找面试哥。

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

    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)
    
      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然后对“键”执行多对多JOIN。

    尽管多对多JOIN技巧适用于大小合理的DataFrame,但你会在较大数据上看到相对较低的性能。

    更快的实现将需要NumPy。这是一些著名的一维笛卡尔积的NumPy实现。我们可以基于其中一些性能解决方案来获得所需的输出。但是,我最喜欢的是@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

    免责声明

    这些解决方案针对具有非混合标量dtype的DataFrames进行了优化。如果处理混合dtype,请自担风险!

    此技巧适用于任何类型的DataFrame。我们使用上述方法计算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
    

    进一步简化
    cartesian_product当只处理两个 DataFrame 时,可能会出现一个不涉及@senderle的简单解决方案。使用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()
    


知识点
面圈网VIP题库

面圈网VIP题库全新上线,海量真题题库资源。 90大类考试,超10万份考试真题开放下载啦

去下载看看