在pyspark中groupby和collect_list with array_zip时保留列名

Question

我有一个数据框架，看起来像这样:

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)。如何才能做到这一点呢?

Answer 1

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

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)])]

Answer 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)])]