Question

我们说我有2个数据框。

DF1在各行的A列中可能具有值{3,4,5}。

DF2在各行的A列中可能具有值{4,5,6}。

我可以使用distinct_set（A）将这些元素聚合成一组不同的元素，假设所有这些行都属于同一个分组。

此时我在结果数据框中有一个集合。反正是否将该集合与另一集合聚合？基本上，如果我有第一个聚合产生的2个数据帧，我希望能够聚合它们的结果。

Answer 1

虽然explode和collect_set可以解决这个问题，但编写自定义聚合器以合并集合本身更有意义。它们下面的结构是WrappedArray。

case class SetMergeUDAF() extends UserDefinedAggregateFunction {

  def deterministic: Boolean = false

  def inputSchema: StructType = StructType(StructField("input", ArrayType(LongType)) :: Nil)

  def bufferSchema: StructType = StructType(StructField("buffer", ArrayType(LongType)) :: Nil)

  def dataType: DataType = ArrayType(LongType)

  def initialize(buf: MutableAggregationBuffer): Unit = {
    buf(0) = mutable.WrappedArray.empty[LongType]
  }

  def update(buf: MutableAggregationBuffer, input: Row): Unit = {
    if (!input.isNullAt(0)) {
      val result : mutable.WrappedArray[LongType] = mutable.WrappedArray.empty[LongType]
      val x = result ++ (buf.getAs[mutable.WrappedArray[Long]](0).toSet ++ input.getAs[mutable.WrappedArray[Long]](0).toSet).toArray[Long]
      buf(0) = x
    }
  }

  def merge(buf1: MutableAggregationBuffer, buf2: Row): Unit = {
    val result : mutable.WrappedArray[LongType] = mutable.WrappedArray.empty[LongType]
    val x = result ++ (buf1.getAs[mutable.WrappedArray[Long]](0).toSet ++ buf2.getAs[mutable.WrappedArray[Long]](0).toSet).toArray[Long]
    buf1(0) = x
  }

  def evaluate(buf: Row): Any = buf.getAs[mutable.WrappedArray[LongType]](0)
}

Spark SQL - 聚合集合？

1 个答案: