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

Question

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

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"))

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

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|
//+----+----+----+---------------+

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

How to apply a customized function with multiple parameters to each group of a dataframe and union the resulting dataframes in Scala Spark?

scala

group-by

user-defined-functions

dataframe

apache-spark