基于两个或多个列的Spark DataFrame聚合

时间:2017-09-19 22:01:18

标签: apache-spark apache-spark-sql spark-dataframe

我想为一些基于多列的自定义聚合编写UDAF。一个简单的例子是具有两列c1和c2的数据帧。对于每一行,我取c1和c2的最大值(我们称之为cmax),然后取cmax之和。

当我调用df.agg()时,看起来我不能将两个或更多列传递给任何聚合方法,包括UDAF。第一个问题,是真的吗?

对于这个简单的例子,我可以创建另一个名为cmax的列,并在cmax上进行聚合。但实际上,我需要基于N个列组合进行聚合,结果将是大小为N的集合。我想在我的UDAF中循环更新方法中的组合。因此它需要N个中间列,这对我来说似乎不是一个干净的解决方案。第二个问题,我想知道是否可以创建中间列,或者是否有更好的解决方案。

我注意到在RDD中,问题要容易得多。我可以将整个记录传递给我的聚合函数,并且我可以访问所有数据字段。

1 个答案:

答案 0 :(得分:1)

您可以在UDAF中使用尽可能多的列,因为它的签名apply函数接受多个Columns(来自它的源代码)。

 def apply(exprs: Column*): Column

您只需确保inputSchema返回StructType,反映您要用作UDAF输入的列。

对于列c1c2的情况,您的UDAF必须使用以下架构实现inputSchema

def inputSchema: StructType = StructType(Array(StructField("c1", DoubleType), StructField("c2", DoubleType)))

但是,如果您想要更通用的解决方案,您始终可以使用允许返回正确inputSchema的参数初始化自定义UDAF。请参阅下面的示例,该示例允许在构建时定义任意StructType注意,我们不确认StructTypeDoubleType)。< / p>

class MyMaxUDAF(schema: StructType) extends UserDefinedAggregateFunction {

  def inputSchema: StructType = this.schema

  def bufferSchema: StructType = StructType(Array(StructField("maxSum", DoubleType)))

  def dataType: DataType = DoubleType

  def deterministic: Boolean = true

  def initialize(buffer: MutableAggregationBuffer): Unit = buffer(0) = 0.0

  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
buffer(0) = buffer.getDouble(0) + Array.range(0, input.length).map(input.getDouble).max
  }

  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = buffer2 match {
    case Row(buffer2Sum: Double) => buffer1(0) = buffer1.getDouble(0) + buffer2Sum
  }

  def evaluate(buffer: Row): Double = buffer match {
    case Row(totalSum: Double) => totalSum
  }

}

您的DataFrame包含值和聚合键。

val df = spark.createDataFrame(Seq(
  Entry(0, 1.0, 2.0, 3.0), Entry(0, 3.0, 1.0, 2.0), Entry(1, 6.0, 2.0, 2)
))
df.show


+-------+---+---+---+
|groupMe| c1| c2| c3|
+-------+---+---+---+
|      0|1.0|2.0|3.0|
|      0|3.0|1.0|2.0|
|      1|6.0|2.0|2.0|
+-------+---+---+---+

使用UDAF,我们希望max的总和为6.0和6.0

val fields = Array("c1", "c2", "c3")
val struct = StructType(fields.map(StructField(_, DoubleType)))
val myMaxUDAF: MyMaxUDAF = new MyMaxUDAF(struct)
df.groupBy("groupMe").agg(myMaxUDAF(fields.map(df(_)):_*)).show


+-------+---------------------+
|groupMe|mymaxudaf(c1, c2, c3)|
+-------+---------------------+
|      0|                  6.0|
|      1|                  6.0|
+-------+---------------------+

UDAF有一个很好的教程。不幸的是,他们没有涵盖多个论点。

https://ragrawal.wordpress.com/2015/11/03/spark-custom-udaf-example/