Pyspark:将数据帧作为数组类型列连接到另一个数据帧

Question

我正在尝试在 pyspark 中连接两个数据帧，但将一个表作为数组列连接在另一个表上。

例如，对于这些表：

from pyspark.sql import Row

df1 = spark.createDataFrame([
    Row(a = 1, b = 'C', c = 26, d = 'abc'),
    Row(a = 1, b = 'C', c = 27, d = 'def'),
    Row(a = 1, b = 'D', c = 51, d = 'ghi'),
    Row(a = 2, b = 'C', c = 40, d = 'abc'),
    Row(a = 2, b = 'D', c = 45, d = 'abc'),
    Row(a = 2, b = 'D', c = 38, d = 'def')
])

df2 = spark.createDataFrame([
    Row(a = 1, b = 'C', e = 2, f = 'cba'),
    Row(a = 1, b = 'D', e = 3, f = 'ihg'),
    Row(a = 2, b = 'C', e = 7, f = 'cba'),
    Row(a = 2, b = 'D', e = 9, f = 'cba')
])

我想在列a和b上将df1连接到df2，但df1. c和df1. d应该是单个数组类型列。此外，应保留所有名称。新数据框的输出应该能够转换为此json结构（例如前两行）：

{
    "a": 1,
    "b": "C",
    "e": 2,
    "f": "cba",
    "df1": [
            {
                "c": 26,
                "d": "abc"
            },
            {
                "c": 27,
                "d": "def"
            } 
           ]
}

任何关于如何实现这一目标的想法都将不胜感激！

谢谢，

卡罗莱纳州

Answer 1

根据您输入的样本数据:

from pyspark.sql import functions as F


df1 = df1.groupBy("a", "b").agg(
    F.collect_list(F.struct(F.col("c"), F.col("d"))).alias("df1")
)

df1.show()
+---+---+--------------------+
|  a|  b|                 df1|
+---+---+--------------------+
|  1|  C|[[26, abc], [27, ...|
|  1|  D|         [[51, ghi]]|
|  2|  D|[[45, abc], [38, ...|
|  2|  C|         [[40, abc]]|
+---+---+--------------------+

df3 = df1.join(df2, on=["a", "b"])

df3.show()
+---+---+--------------------+---+---+
|  a|  b|                 df1|  e|  f|
+---+---+--------------------+---+---+
|  1|  C|[[26, abc], [27, ...|  2|cba|
|  1|  D|         [[51, ghi]]|  3|ihg|
|  2|  D|[[45, abc], [38, ...|  9|cba|
|  2|  C|         [[40, abc]]|  7|cba|
+---+---+--------------------+---+---+