使用 Pandas 进行高性能笛卡尔积 (CROSS JOIN)
- 2024-12-06 08:40:00
- admin 原创
- 148
问题描述:
这篇文章的内容原本是打算成为Pandas Merging 101的一部分
,但由于需要充分阐述这一主题所需的内容的性质和大小,它已被移至其自己的问答中。
给定两个简单的 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
计算这个结果的最有效方法是什么?
解决方案 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
然后对“键”执行多对多 JOIN。
虽然多对多 JOIN 技巧适用于合理大小的 DataFrames,但您会看到较大数据的性能相对较低。
更快的实现需要 NumPy。以下是一些著名的1D 笛卡尔积 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)
概括:在唯一或非唯一索引 DataFrames上使用 CROSS JOIN
免责声明
这些解决方案针对具有非混合标量数据类型的 DataFrames 进行了优化。如果处理混合数据类型,请自行承担风险!
此技巧适用于任何类型的 DataFrame。我们使用上述方法计算 DataFrames 数字索引的笛卡尔积cartesian_product
,然后使用它重新索引 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
性能比较
在一些具有唯一索引的 DataFrames 上对这些解决方案进行基准测试,我们得到
请注意,时间可能会根据您的设置、数据和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 中的其他主题继续学习:
合并基础 - 连接的基本类型
基于索引的连接
推广到多个 DataFrames
交叉连接 *
你在这里
解决方案 2:
从 pandas 1.2.0 开始,merge
现在有选项cross
left.merge(right, how='cross')
使用itertools
product
并重新创建数据框中的值
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
解决方案 3:
以下是三重方法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
解决方案 4:
一个选项是使用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
解决方案 5:
对于大型数据集来说,此方法是最好且最有效的方法。
使用:pd.factorize + np.repeat + np.tile + np.hstack
import pandas as pd
import numpy as np
df1 = pd.DataFrame({'col1': ['A', 'B', 'C'], 'col2': [1, 2, 3]})
df2 = pd.DataFrame({'col1': ['X', 'Y', 'Z'], 'col2': [20, 30, 50]})
print(df1)
print(df2)
df1['col1_codes'],col1_unique_df1 = pd.factorize(df1['col1'])
df2['col1_codes'],col1_unique_df2 = pd.factorize(df2['col1'])
df1_repeated = np.repeat(df1[['col1_codes','col2']].values,df2.shape[0],axis=0)
print(df1_repeated)
df2_tiled = np.tile(df2[['col1_codes','col2']].values, (df1.shape[0],1))
print(df2_tiled)
# Combine results
nphstack = np.hstack([df1_repeated,df2_tiled])
print(nphstack)
res = pd.DataFrame(nphstack,columns = ['col1_codes_df1','col2_df1','col1_codes_df2','col2_df2'])
print(res)
"""
col1_codes_df1 col2_df1 col1_codes_df2 col2_df2
0 0 1 0 20
1 0 1 1 30
2 0 1 2 50
3 1 2 0 20
4 1 2 1 30
5 1 2 2 50
6 2 3 0 20
7 2 3 1 30
8 2 3 2 50
"""
# Map the codes back to the original strings
res['col1_df1'] = res['col1_codes_df1'].map(dict(enumerate(col1_unique_df1)))
res['col1_df2'] = res['col1_codes_df2'].map(dict(enumerate(col1_unique_df2)))
print(res)
"""
col1_codes_df1 col2_df1 col1_codes_df2 col2_df2 col1_df1 col1_df2
0 0 1 0 20 A X
1 0 1 1 30 A Y
2 0 1 2 50 A Z
3 1 2 0 20 B X
4 1 2 1 30 B Y
5 1 2 2 50 B Z
6 2 3 0 20 C X
7 2 3 1 30 C Y
8 2 3 2 50 C Z
"""
# Drop the code columns if not needed
res = res.drop(columns=['col1_codes_df1', 'col1_codes_df2'])
print(res)
"""
col2_df1 col2_df2 col1_df1 col1_df2
0 1 20 A X
1 1 30 A Y
2 1 50 A Z
3 2 20 B X
4 2 30 B Y
5 2 50 B Z
6 3 20 C X
7 3 30 C Y
8 3 50 C Z
"""
解决方案 6:
我认为最简单的方法是向每个数据框添加一个虚拟列,对其进行内部合并,然后从生成的笛卡尔数据框中删除该虚拟列:
left['dummy'] = 'a'
right['dummy'] = 'a'
cartesian = left.merge(right, how='inner', on='dummy')
del cartesian['dummy']