在pyspark中groupby时保留列名,并使用数组\u zip收集\u列表

8hhllhi2  于 2021-07-13  发布在  Spark
关注(0)|答案(2)|浏览(738)

我有这样一个Dataframe:

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

上述结果如下:

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

我需要分组讨论 foo 然后把其他的都列成单子。结果需要与派生它们的列保持相同的名称。
到目前为止,我得到了这个:

  1. df.groupBy('foo').agg(
  2. F.arrays_zip(
  3. F.collect_list(F.col('bar')),
  4. F.collect_list(F.col('baz')),
  5. ).alias('events')
  6. )

输出返回一个Dataframe,其中 events 列包含行,每个行都有递增的键 0 , 1 , 2 等等。

  1. [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) . 如何做到这一点?

wvt8vs2t

wvt8vs2t1#

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

  1. df1 = df.groupBy('foo').agg(
  2. F.collect_list(
  3. F.struct(
  4. F.col('bar'),
  5. F.col('baz')
  6. )
  7. ).alias("events")
  8. )
  9. print(df1.collect())
  10. # [Row(foo='a', events=[Row(bar=1, baz=4), Row(bar=2, baz=5)]), Row(foo='b', events=[Row(bar=3, baz=6)])]
d5vmydt9

d5vmydt92#

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

  1. df.groupBy('foo').agg(
  2. F.collect_list(F.col('bar')).alias('bar'),
  3. F.collect_list(F.col('baz')).alias('baz'),
  4. ).withColumn(
  5. 'events',
  6. F.arrays_zip(
  7. F.col('bar'),
  8. F.col('baz'),
  9. )
  10. ).drop('bar', 'baz').collect()

结果是:

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

相关问题