CODEWITHSUNDEEP
dataframescalaapache-sparkapache-spark-sql
Spark_user
Spark_user

Reputation: 41

How to dynamically build column name in spark sql withColumn

I have sample data as below.

input

I would like populate another column exp based on country field value.

Something like below

df.withColumn("exp",col(s"exp_$country"))

So that respective country number can be placed there.

But above code errors out saying:

cannot resolve country

Output I need is

2

Any help appreciated.

Upvotes: 0

Views: 1332

Answers (1)

blackbishop
blackbishop

Reputation: 32660

You can chain multiple when expressions from the list of countries:

val countries = Seq("us", "uk", "ind")

val expCol = countries.foldLeft(lit(null)) { case (acc, country) =>
  when(col("country")===country, col(s"exp_$country")).otherwise(acc)
}

val df1 = df.withColumn("exp", expCol)

Or if you prefer creating a map expression country -> exp from the columns exp_* than use the map to create exp column:

val mapCountries = map(
  df.columns
    .filter(_.startsWith("exp_"))
    .flatMap(c => Seq(lit(c.split("_")(1)), col(c))): _*
)

val df1 = df.withColumn("exp", mapCountries(col("country")))

Upvotes: 2

Related Questions