How to use UDF to return multiple columns?

22
votes

Is it possible to create a UDF which would return the set of columns?

I.e. having a data frame as follows:

| Feature1 | Feature2 | Feature 3 |
| 1.3      | 3.4      | 4.5       |

Now I would like to extract a new feature, which can be described as a vector of let's say two elements (e.g. as seen in a linear regression - slope and offset). Desired dataset shall look as follows:

| Feature1 | Feature2 | Feature 3 | Slope | Offset |
| 1.3      | 3.4      | 4.5       | 0.5   | 3      |

Is it possible to create multiple columns with single UDF or do I need to follow the rule: "single column per single UDF"?

2
apache-sparkapache-spark-sql

2 Answers

46
votes

Struct method

You can define the udf function as 

def myFunc: (String => (String, String)) = { s => (s.toLowerCase, s.toUpperCase)}

import org.apache.spark.sql.functions.udf
val myUDF = udf(myFunc)

and use .* as 

val newDF = df.withColumn("newCol", myUDF(df("Feature2"))).select("Feature1", "Feature2", "Feature 3", "newCol.*")

I have returned Tuple2 for testing purpose (higher order tuples can be used according to how many multiple columns are required) from udf function and it would be treated as struct column. Then you can use .* to select all the elements in separate columns and finally rename them.

You should have output as 

+--------+--------+---------+---+---+
|Feature1|Feature2|Feature 3|_1 |_2 |
+--------+--------+---------+---+---+
|1.3     |3.4     |4.5      |3.4|3.4|
+--------+--------+---------+---+---+

You can rename _1 and _2

Array method

udf function should return an array 

def myFunc: (String => Array[String]) = { s => Array("s".toLowerCase, s.toUpperCase)}

import org.apache.spark.sql.functions.udf
val myUDF = udf(myFunc)

And the you can select elements of the array and use alias to rename them

val newDF = df.withColumn("newCol", myUDF(df("Feature2"))).select($"Feature1", $"Feature2", $"Feature 3", $"newCol"(0).as("Slope"), $"newCol"(1).as("Offset"))

You should have 

+--------+--------+---------+-----+------+
|Feature1|Feature2|Feature 3|Slope|Offset|
+--------+--------+---------+-----+------+
|1.3     |3.4     |4.5      |s    |3.4   |
+--------+--------+---------+-----+------+
6
votes

Also, you can return the case class:

case class NewFeatures(slope: Double, offset: Int)

val getNewFeatures = udf { s: String =>
      NewFeatures(???, ???)
    }

df
  .withColumn("newF", getNewFeatures($"Feature1"))
  .select($"Feature1", $"Feature2", $"Feature3", $"newF.slope", $"newF.offset")