在 pyspark 中使用 array_zip 和 collect_list 时保留列名
Preserve column names when groupby and collect_list with array_zip in pyspark
我有一个如下所示的数据框:
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')
)
输出 returns 一个数据框,其中 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)
。如何实现?
解决方案是将其分为 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)])]
您可以简单地收集结构列表而不是创建 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)])]
我有一个如下所示的数据框:
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')
)
输出 returns 一个数据框,其中 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)
。如何实现?
解决方案是将其分为 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)])]
您可以简单地收集结构列表而不是创建 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)])]