Question

我有一些输入文件，并且所有文件共享相同的架构。它们都具有一个名为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似乎不太好。是否有更优雅的方式或合适的设计模式来编写代码？非常感谢。

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

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

根据条件应用UDF的优雅方法

2 个答案: