27

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"?

Jacek Laskowski
  • 72,696
  • 27
  • 242
  • 420
TechCrap
  • 930
  • 3
  • 14
  • 28

3 Answers3

53

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   |
+--------+--------+---------+-----+------+
Ramesh Maharjan
  • 41,071
  • 6
  • 69
  • 97
  • 1
    That's a great solution, thanks! Another question is how many times is the UDF called? I added counter to count the number of calls and in the code above the UDF is called 3 times. Any way to fix this? – Kal-ko Jan 29 '19 at 21:38
  • @RameshMaharjan I saw your other answer on processing all columns in `df`, and combined with this, they offer a great solution.However, I am stuck at using the return value from the `UDF` to modify multiple columns using `withColumn` which only takes one column name at a time. Would you know of a workaround? – smaug Feb 08 '19 at 17:04
9

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")
skotlov
  • 186
  • 1
  • 4
3

I miss an explanation about how to assign the multiples values in the case class to several columns in the dataframe.

So, in summary, a complete example in Scala

import org.apache.spark.sql.functions.udf

val df = Seq((1L, 3.0, "a"), (2L, -1.0, "b"), (3L, 0.0, "c")).toDF("x", "y", "z")

case class Foobar(foo: Double, bar: Double)

val foobarUdf = udf((x: Long, y: Double, z: String) => 
  Foobar(x * y, z.head.toInt * y))

val df1 = df.withColumn("foo", foobarUdf($"x", $"y", $"z").getField("foo")).withColumn("bar", foobarUdf($"x", $"y", $"z").getField("bar"))

If you check the schema of the df1 dataframe, you'll get

scala> df1.printSchema
root
 |-- x: long (nullable = false)
 |-- y: double (nullable = false)
 |-- z: string (nullable = true)
 |-- foo: double (nullable = true)
 |-- bar: double (nullable = true)
evinhas
  • 179
  • 1
  • 4