将数据框列值映射到 Scala 字典

Question

我有一个 Scala 字典，其中 stateCodes 映射到州名。
我有一个包含状态代码的列的数据框。 我想用状态名称向数据框中添加一列。

val stateCodes = Map("100" -> "ca",
                 "22" -> "ny",
                 "12" -> "tx,
                 "71" -> "va")

df = 
+------------+--------+
|     stateId|   count|
+------------+--------+
|         100|17243160|
|          71|12398011|
|          12|12302708|
|           3| 9734593|

预期输出

new_df = 
+------------+--------+
|stateName   |   count|
+------------+--------+
|          ca|17243160|
|          va|12398011|
|          tx|12302708|
|     unknown| 9734593|

Answer 1

您可以通过多种方式做到这一点。 一种最简单的方法是将地图转换为数据框并使用连接。

import spark.implicits._
val stateCodesDf = stateCodes.toSeq().toDF("stateId","stateName")
val newDf = df.join(stateCodesDF,Seq("stateId"),"left_outer").
    select("stateName","count")

第二种是写一个简单的UDF

import org.apache.spark.sql.functions.udf
import spark.implicits._
val stateUdf = udf( (stateId:String) => stateCodes.get(stateId) )
val newDf = df.select(stateUdf($"stateId"), $"count")

两种解决方案都在输出中包含表中没有匹配项且将为 NULL 的条目。