Spark SQL:对数组值使用collect_set

Question

我有一个聚合的数据帧，其中包含使用 collect_set 创建的列。我现在需要再次聚合此数据帧，并再次collect_set应用于该列的值。问题是我需要应用collect_Set验证集合的值 - 并且到目前为止，我看到如何做到这一点的唯一方法是分解聚合的数据帧。有没有更好的方法？

例：

初始数据帧：

country   | continent   | attributes
-------------------------------------
Canada    | America     | A
Belgium   | Europe      | Z
USA       | America     | A
Canada    | America     | B
France    | Europe      | Y
France    | Europe      | X

聚合数据帧(我作为输入接收的数据帧( - country聚合：

country   | continent   | attributes
-------------------------------------
Canada    | America     | A, B
Belgium   | Europe      | Z
USA       | America     | A
France    | Europe      | Y, X

我想要的输出 - continent聚合：

continent   | attributes
-------------------------------------
America     | A, B
Europe      | X, Y, Z

Answer 1

由于此时只能有少数行，因此只需按原样收集属性并展平结果(Spark>= 2.4(

import org.apache.spark.sql.functions.{collect_set, flatten, array_distinct}
val byState = Seq(
  ("Canada", "America", Seq("A", "B")),
  ("Belgium", "Europe", Seq("Z")),
  ("USA", "America", Seq("A")),
  ("France", "Europe", Seq("Y", "X"))
).toDF("country", "continent", "attributes")
byState
  .groupBy("continent")
  .agg(array_distinct(flatten(collect_set($"attributes"))) as "attributes")
  .show

+---------+----------+
|continent|attributes|
+---------+----------+
|   Europe| [Y, X, Z]|
|  America|    [A, B]|
+---------+----------+

在一般情况下，事情更难处理，在许多情况下，如果您期望大型列表，每个组有许多重复项和许多值，则最佳解决方案*是从头开始重新计算结果，即

input.groupBy($"continent").agg(collect_set($"attributes") as "attributes")

一种可能的替代方法是使用Aggregator

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.{Encoder, Encoders}
import scala.collection.mutable.{Set => MSet}

class MergeSets[T, U](f: T => Seq[U])(implicit enc: Encoder[Seq[U]]) extends 
     Aggregator[T, MSet[U], Seq[U]] with Serializable {
  def zero = MSet.empty[U]
  def reduce(acc: MSet[U], x: T) = {
    for { v <- f(x) } acc.add(v)
    acc
  }
  def merge(acc1: MSet[U], acc2: MSet[U]) = {
    acc1 ++= acc2
  }
  def finish(acc: MSet[U]) = acc.toSeq
  def bufferEncoder: Encoder[MSet[U]] = Encoders.kryo[MSet[U]]
  def outputEncoder: Encoder[Seq[U]] = enc
}

并按如下方式应用它

case class CountryAggregate(
  country: String, continent: String, attributes: Seq[String])
byState
  .as[CountryAggregate]
  .groupByKey(_.continent)
  .agg(new MergeSets[CountryAggregate, String](_.attributes).toColumn)
  .toDF("continent", "attributes")
  .show

+---------+----------+
|continent|attributes|
+---------+----------+
|   Europe| [X, Y, Z]|
|  America|    [B, A]|
+---------+----------+

但这显然不是一个对Java友好的选择。

另请参阅如何在 groupBy 之后将值聚合到集合中？(类似，但没有唯一性约束(。

---

* 这是因为explode可能非常昂贵，尤其是在较旧的 Spark 版本中，与访问 SQL 集合的外部表示形式相同。