给定两个简单的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
计算结果的最有效方法是什么?
让我们从建立基准开始。解决此问题的最简单方法是使用临时“键”列:
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。
merge
尽管多对多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()