如何在 Pyspark Dataframe 中创建多列的所有成对组合？

Posted 2023-04-17

【中文标题】如何在 Pyspark Dataframe 中创建多列的所有成对组合？

【英文标题】：How can I create all pairwise combinations of multiple columns in a Pyspark Dataframe?

【发布时间】：2021-06-03 15:54:25

【问题描述】：

考虑以下 Pyspark 数据框

Col1	Col2	Col3
A	D	G
B	E	H
C	F	I

如何创建以下包含所有列的所有成对组合的数据框？

Col1	Col2	Col3	Col1_Col2_cross	Col1_Col3_cross	Col2_Col3_cross
A	D	G	A,D	A,G	D,G
B	E	H	B,E	B,H	E,H
C	F	I	C,F	C,I	F,I

【问题讨论】：

它是否需要适用于任意数量的列，还是只适用于 3 个？

任意数量的列。我实际上有大约 30 列，但这里只显示了 3 列以简化问题。

【参考方案1】：

您可以使用itertools生成列组合：

import pyspark.sql.functions as F
import itertools

df2 = df.select(
    '*', 
    *[F.concat_ws(',', x[0], x[1]).alias(x[0] + '_' + x[1] + '_cross') 
      for x in itertools.combinations(df.columns, 2)]
)

df2.show()
+----+----+----+---------------+---------------+---------------+
|Col1|Col2|Col3|Col1_Col2_cross|Col1_Col3_cross|Col2_Col3_cross|
+----+----+----+---------------+---------------+---------------+
|   A|   D|   G|            A,D|            A,G|            D,G|
|   B|   E|   H|            B,E|            B,H|            E,H|
|   C|   F|   I|            C,F|            C,I|            F,I|
+----+----+----+---------------+---------------+---------------+