CODEWITHSUNDEEP
scalaapache-sparkapache-spark-sql
Tates90
Tates90

Reputation: 23

Select column based on Map

I have the following dataframe df with the following schema:

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

and a map

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

I want a new column "record" that is either the value of record_sales or record_marketing based on the value of type.

I've tried some variants of this:

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

But nothing has worked. Does anyone have any idea? Thanks!

Upvotes: 2

Views: 44

Answers (1)

blackbishop
blackbishop

Reputation: 32700

You can iterate over the map typemap and use when function to get case/when expressions depending on the value of type column:

val recordCol = typemap.map{case (k,v) => when(col("type") === k, col(v))}.toSeq

val df2 = df.withColumn("record", coalesce(recordCol: _*))

Upvotes: 1

Related Questions