在pyspark中使用array_zip进行groupby和collect_list时保留列名答案

【问题标题】：Preserve column names when groupby and collect_list with array_zip in pyspark在pyspark中使用array_zip进行groupby和collect_list时保留列名
【发布时间】：2021-05-10 07:38:06
【问题描述】：

我有一个如下所示的数据框：

df = spark.createDataFrame([
    Row(foo='a', bar=1, baz=4),
    Row(foo='a', bar=2, baz=5),
    Row(foo='b', bar=3, baz=6),
])

上面的结果如下：

[Row(bar=1, baz=4, foo='a'), Row(bar=2, baz=5, foo='a'), Row(bar=3, baz=6, foo='b')]

我需要在foo 上进行分组，然后将其他所有内容收集为列表。结果需要保留与派生它们的列相同的名称。

到目前为止，我得到了这个：

df.groupBy('foo').agg(
    F.arrays_zip(
        F.collect_list(F.col('bar')),
        F.collect_list(F.col('baz')),
    ).alias('events')
)

输出返回一个数据帧，其中events 列包含行，每个行都有递增0、1、2 等的键。

[Row(foo='a', events=[Row(0=1, 1=4), Row(0=2, 1=5)]), Row(foo='b', events=[Row(0=3, 1=6)])]

我想保留原始列名作为键。这样像Row(0=1, 1=4) 这样的行将改为Row(bar=1, baz=4)。如何实现？

【问题讨论】：

标签： apache-spark pyspark apache-spark-sql pyspark-dataframes

【解决方案1】：

您可以简单地收集结构列表，而不是创建 2 个数组列并压缩它们：

df1 = df.groupBy('foo').agg(
    F.collect_list(
        F.struct(
            F.col('bar'), 
            F.col('baz')
        )
    ).alias("events")
)

print(df1.collect()) 

# [Row(foo='a', events=[Row(bar=1, baz=4), Row(bar=2, baz=5)]), Row(foo='b', events=[Row(bar=3, baz=6)])]

【讨论】：

【解决方案2】：

解决方案是将其分为两个步骤：

df.groupBy('foo').agg(
    F.collect_list(F.col('bar')).alias('bar'),
    F.collect_list(F.col('baz')).alias('baz'),
).withColumn(
    'events',
    F.arrays_zip(
        F.col('bar'),
        F.col('baz'),
    )
).drop('bar', 'baz').collect()

导致：

[Row(foo='a', events=[Row(bar=2, baz=5), Row(bar=1, baz=4)]), Row(foo='b', events=[Row(bar=3, baz=6)])]

【讨论】：