按条件应用 UDF 的优雅方式答案

【问题标题】：Elegant way to apply UDF by condition按条件应用 UDF 的优雅方式
【发布时间】：2019-12-09 14:19:00
【问题描述】：

我有一些输入文件，所有文件共享相同的架构。它们都有一个名为channel_id 的字段，但对于file1、channel_id = 1、对于file2、channel_id = 2。

我需要对这些文件做一些 ETL。但是，对于不同的文件，逻辑是不同的。比如有一个UDF来计算channel_name

val getChannelNameUdf : UserDefinedFunction = udf((channelId: Integer) => {
    if (channelId == 1) {
      "English"
    } else if (channelId == 2) {
      "French"
    } else {
      ""
    }
  })

由于我们有多个频道，使用if-else 似乎不太优雅。是否有更优雅的方式或合适的设计模式来编写代码？非常感谢。

【问题讨论】：

标签： apache-spark apache-spark-sql user-defined-functions

【解决方案1】：

嗨，布鲁克林，欢迎来到 StackOverflow，

您可以在 UDF 中使用模式匹配，但我建议您使用 when 内置函数，而不是定义自己的 UDF。

要回答您的请求，这里是您可能需要的代码：

val getChannelNameUdf = udf[String, Int] { _ match {
  case 1 => "English"
  case 2 => "French"
  case _ => ""
}}

甚至更好，只是匿名函数：

val getChannelNameUdf = udf[String, Int] {
  case 1 => "English"
  case 2 => "French"
  case _ => ""
}

这是一个使用 when 内置函数的示例：

val getChannelName = {col: Column =>
  when(col === 1, "English").when(col === 2, "French").otherwise("")
}
df.withColumn("channelName", getChannelName($"channelId"))

编辑：对于更通用的方法，您可以使用以下定义：

val rules = Map(1 -> "English", 2 -> "French")
val getChannelName = {col: Column =>
  rules.foldLeft(lit("")){case (c, (i,label)) =>
    when(col === i, label).otherwise(c)
  }
}

然后

df.withColumn("channelName", getChannelName($"channelId"))

【讨论】：

【解决方案2】：

是否有更优雅的方式或合适的设计模式来编写代码？

是的！一个简单有效的方法是使用join。

您可以拥有一个包含所有通道引用的文件，假设它具有以下结构：channel_id, channel_name，然后加入 2 个 DataFrame。像这样的：

val df_channels = spark.read.csv("/path/to/channel_file.csv")

val result = df.join(df_channels, Seq("channel_id"),"left")

【讨论】：