CODEWITHSUNDEEP
scalaapache-spark
Jivan
Jivan

Reputation: 23098

Dynamic column selection in Spark (based on another column's value)

With the given Spark DataFrame:

> df.show()

+---+-----+---+---+---+---+
| id|delay| p1| p2| p3| p4|
+---+-----+---+---+---+---+
|  1|    3|  a|  b|  c|  d|
|  2|    1|  m|  n|  o|  p|
|  3|    2|  q|  r|  s|  t|
+---+-----+---+---+---+---+

How to select a column dynamically so that the new, col column is the result of the p{delay} existing column?

> df.withColumn("col", /* ??? */).show()

+---+-----+---+---+---+---+----+
| id|delay| p1| p2| p3| p4| col|
+---+-----+---+---+---+---+----+
|  1|    3|  a|  b|  c|  d|   c|   // col = p3
|  2|    1|  m|  n|  o|  p|   m|   // col = p1
|  3|    2|  q|  r|  s|  t|   r|   // col = p2
+---+-----+---+---+---+---+----+

Upvotes: 1

Views: 1670

Answers (2)

zero323
zero323

Reputation: 330433

The simplest solution I can think of is to use array with delay as an index:

import org.apache.spark.sql.functions.array

df.withColumn("col", array($"p1", $"p2", $"p3", $"p4")($"delay" - 1))

Upvotes: 8

akuiper
akuiper

Reputation: 215127

One option is create a map from number to column names, and then use foldLeft to update the col column with corresponding values:

val cols = (1 to 4).map(i => i -> s"p$i")

(cols.foldLeft(df.withColumn("col", lit(null))){ 
   case (df, (k, v)) => df.withColumn("col", when(df("delay") === k, df(v)).otherwise(df("col"))) 
}).show
+---+-----+---+---+---+---+---+    
| id|delay| p1| p2| p3| p4|col|
+---+-----+---+---+---+---+---+
|  1|    3|  a|  b|  c|  d|  c|
|  2|    1|  m|  n|  o|  p|  m|
|  3|    2|  q|  r|  s|  t|  r|
+---+-----+---+---+---+---+---+

Upvotes: 1

Related Questions