如何将具有多个参数的自定义函数应用于每组数据帧并在 Scala Spark 中合并生成的数据帧？

Question

我有一个看起来像这样的自定义函数，它返回一个不同的数据帧作为输出

def customizedfun(data : DataFrame, param1 : Boolean, param2 : string) : DataFrame = {...}

我想将此功能应用于每一组

df.groupBy("type")

然后将每个type 的输出数据帧附加到一个数据帧中。

这与其他关于将自定义函数应用于分组数据帧的问题略有不同，因为除了有问题的数据帧df.groupBy("type") 之外，此函数还接受其他输入。

最好的方法是什么？

Answer 1

您可以将原始df 过滤到不同的组，为每个组调用customizedfun，然后合并结果。

我假设customizedfun 是一个简单地将两个参数添加为新列的函数，但它可以是任何函数：

def customizedfun(data : DataFrame, param1 : Boolean, param2 : String) : DataFrame =
  data.withColumn("newCol", lit(s"$param2 $param1"))

我需要两个帮助函数来计算 param1 和 param2 的值，这取决于 type 的值。在现实世界的应用程序中，这些功能可以是例如查找字典。

def calcParam1(typ: Integer): Boolean = typ % 2 == 0
def calcParam2(typ: Integer): String = s"type is $typ"

现在将原来的df过滤到不同的组中，调用customizedfun并将结果合并：

//create some test data
val df = Seq((1, "A", "a"), (1, "B", "b"), (1, "C", "c"), (2, "D", "d"), (2, "E", "e"), (3, "F", "f"))
  .toDF("type", "val1", "val2")
//+----+----+----+
//|type|val1|val2|
//+----+----+----+
//|   1|   A|   a|
//|   1|   B|   b|
//|   1|   C|   c|
//|   2|   D|   d|
//|   2|   E|   e|
//|   3|   F|   f|
//+----+----+----+

//get the distinct values of column type
val distinctTypes = df.select("type").distinct().as[Integer].collect()

//call customizedfun for each group
val resultPerGroup= for( typ <- distinctTypes)
  yield customizedfun( df.filter(s"type = $typ"), calcParam1(typ), calcParam2(typ))

//the final union
val result = resultPerGroup.tail.foldLeft(resultPerGroup.head)(_ union _)

//+----+----+----+---------------+
//|type|val1|val2|         newCol|
//+----+----+----+---------------+
//|   1|   A|   a|type is 1 false|
//|   1|   B|   b|type is 1 false|
//|   1|   C|   c|type is 1 false|
//|   3|   F|   f|type is 3 false|
//|   2|   D|   d| type is 2 true|
//|   2|   E|   e| type is 2 true|
//+----+----+----+---------------+