如何在 Dataframe Spark Scala 中对列表进行分组和连接答案

【问题标题】：How do I groupby and concat a list in a Dataframe Spark Scala如何在 Dataframe Spark Scala 中对列表进行分组和连接
【发布时间】：2018-10-18 20:26:56
【问题描述】：

我有一个包含两列数据的数据框，数据如下

+----+-----------------+
|acct|           device|
+----+-----------------+
|   B|       List(3, 4)|
|   C|       List(3, 5)|
|   A|       List(2, 6)|
|   B|List(3, 11, 4, 9)|
|   C|       List(5, 6)|
|   A|List(2, 10, 7, 6)|
+----+-----------------+

我需要如下结果

+----+-----------------+
|acct|           device|
+----+-----------------+
|   B|List(3, 4, 11, 9)|
|   C|    List(3, 5, 6)|
|   A|List(2, 6, 7, 10)|
+----+-----------------+

我尝试如下，但它似乎不起作用

df.groupBy("acct").agg(concat("device"))

df.groupBy("acct").agg(collect_set("device"))

请告诉我如何使用 Scala 实现这一目标？

【问题讨论】：

@sb0709 没有GroupConcat这样的功能

标签： scala apache-spark dataframe apache-spark-sql

【解决方案1】：

您可以从展开 device 列开始，然后照常继续操作 - 但请注意，它可能不会保留列表的顺序（无论如何，这在任何 group by 中都不能保证）：

val result = df.withColumn("device", explode($"device"))
  .groupBy("acct")
  .agg(collect_set("device"))

result.show(truncate = false)
// +----+-------------------+
// |acct|collect_set(device)|
// +----+-------------------+
// |B   |[9, 3, 4, 11]      |
// |C   |[5, 6, 3]          |
// |A   |[2, 6, 10, 7]      |
// +----+-------------------+

【讨论】：

【解决方案2】：

您可以尝试使用collect_set 和Window。在你的情况下：

df.withColumn("device", collect_set("device").over(Window.partitionBy("acct")))

【讨论】：

我猜这会导致嵌套数组

【解决方案3】：

另一个可能比explode选项性能更好的选项：创建您自己的UserDefinedAggregationFunction，将列表合并到不同的集合中。

您必须将UserDefinedAggregateFunction 扩展如下：

class MergeListsUDAF extends UserDefinedAggregateFunction {

  override def inputSchema: StructType = StructType(Seq(StructField("a", ArrayType(IntegerType))))

  override def bufferSchema: StructType = inputSchema

  override def dataType: DataType = ArrayType(IntegerType)

  override def deterministic: Boolean = true

  override def initialize(buffer: MutableAggregationBuffer): Unit = buffer.update(0, mutable.Seq[Int]())

  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    val existing = buffer.getAs[mutable.Seq[Int]](0)
    val newList = input.getAs[mutable.Seq[Int]](0)
    val result = (existing ++ newList).distinct
    buffer.update(0, result)
  }

  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = update(buffer1, buffer2)

  override def evaluate(buffer: Row): Any = buffer.getAs[mutable.Seq[Int]](0)
}

然后像这样使用它：

val mergeUDAF = new MergeListsUDAF()

df.groupBy("acct").agg(mergeUDAF($"device"))

【讨论】：