getting the values of a column with keys - spark scala

Viewed 32

I have a map[String,String] like this

val map1 = Map( "S" -> 1 , "T" -> 2, "U" -> 3) 

and a Dataframe with a column called mappedcol ( type array[string] ). Here are the first and second rows of the column : [S,U] , [U,U] and I would like to map every row of this column to get the value of the key so I would have [1,3] instead of [S,U] and [3,3] instead of [U,U]. How can I do this effectively?

Thanks

1 Answers

The map can be tranformed into an SQL expression based on transform and when:

var ex = "transform(value, v -> case ";
for ((k,v) <- map1) ex += s"when v = '${k}' then ${v} "
ex += "else 99 end)"

ex now contains the string

transform(value, v -> case when v = 'S' then 1 when v = 'T' then 2 when v = 'U' then 3 else 99 end)

This expression can now be used to calculate a new column:

import org.apache.spark.sql.functions._

df.withColumn("result", expr(ex)).show();

Output:

+---+------+------+
| id| value|result|
+---+------+------+
|  1|[S, U]|[1, 3]|
|  2|[U, U]|[3, 3]|
+---+------+------+
Related