根据Map选择列

Question

我有以下数据框架df，具有以下模式:

|-- type: string (nullable = true)
|-- record_sales: array (nullable = false)
|    |-- element: string (containsNull = false)
|-- record_marketing: array (nullable = false)
|    |-- element: string (containsNull = false)

和地图

typemap = Map("sales" -> "record_sales", "marketing" -> "record_marketing")

我想要一个新列"record"即根据type的值计算record_sales或record_marketing的值

我已经尝试了这个的一些变体:

val typeMapCol = typedLit(typemap)
val df2 = df.withColumn("record", col(typeMapCol(col("type")))) 

但是什么都没起作用。有人知道吗?谢谢!

Answer 1

您可以遍历映射typemap并使用when函数根据type列的值获得case/when表达式:

val recordCol = typemap.map{case (k,v) => when(col("type") === k, col(v))}.toSeq
val df2 = df.withColumn("record", coalesce(recordCol: _*))