如何获取 spark 中任何值的映射值？

Question

我正在使用嵌套的 json 创建 etl。我主要使用 explode 和 withColumn 来展平嵌套的 json。 json 中的键之一具有状态名称的值。所以我想获取该值并将其转换为缩写的州名。例如：Texas 应该是 TX.

.withColumn("NewColName", $"col.someKey.keyStoringValue")

这有助于我从 json 中的所需键获取关联值。因此，如果 ...someKey.keyStoringValue 的值为 Texas，我将从我的代码中获取 Texas。但我希望它是 TX。我怎样才能做到这一点？我试图创建一个类似的函数；

def getStateName(state: String): String = {
  val maps = Map("Texas"->"TX", "Virginia"->"VA")
  maps.apply(state)
}

.....
.....
.withColumn("NewColName", getStateName($"col.someKey.keyStoringValue"))
.....
.....

我收到这个错误；

error: type mismatch;
found: org.apache.spark.sql.ColumnName
required: String

Answer 1

当您使用 DataFrame/Spark SQL api 时，您不能应用普通的 Scala 函数。要使用这种映射函数，您需要创建一个 UDF：

import org.apache.spark.sql.functions.udf

val maps = Map("Texas"->"TX", "Virginia"->"VA")
val fun = udf { in: String => maps(in)}

...

.withColumn("NewColName", fun($"col.someKey.keyStoringValue"))

Answer 2

您正在寻找的方法是 when（参见 documentation），它允许您 return 基于条件的值。你会做类似的事情：

df.withColumn(
  "NewColName", 
  when($"col.someKey.keyStoringValue" === "Texas", "TX")
  .when($"col.someKey.keyStoringValue" === "Virginia", "VA")
  .otherwise("Unknown")
)

如何获取 spark 中任何值的映射值？

How to get mapped value for any value in spark?

scala

apache-spark