Derive multiple columns from a single column in a Spark DataFrame

by

Generally speaking what you want is not directly possible. UDF can return only a single column at the time. There are two different ways you can overcome this limitation:

  1. Return a column of complex type. The most general solution is a StructType but you can consider ArrayType or MapType as well.

    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("foobar", foobarUdf($"x", $"y", $"z"))
df1.show
// +---+----+---+------------+
// |  x|   y|  z|      foobar|
// +---+----+---+------------+
// |  1| 3.0|  a| [3.0,291.0]|
// |  2|-1.0|  b|[-2.0,-98.0]|
// |  3| 0.0|  c|   [0.0,0.0]|
// +---+----+---+------------+

df1.printSchema
// root
//  |-- x: long (nullable = false)
//  |-- y: double (nullable = false)
//  |-- z: string (nullable = true)
//  |-- foobar: struct (nullable = true)
//  |    |-- foo: double (nullable = false)
//  |    |-- bar: double (nullable = false)

    This can be easily flattened later but usually there is no need for that.

  2. Switch to RDD, reshape and rebuild DF:

    import org.apache.spark.sql.types._
import org.apache.spark.sql.Row

def foobarFunc(x: Long, y: Double, z: String): Seq[Any] = 
  Seq(x * y, z.head.toInt * y)

val schema = StructType(df.schema.fields ++
  Array(StructField("foo", DoubleType), StructField("bar", DoubleType)))

val rows = df.rdd.map(r => Row.fromSeq(
  r.toSeq ++
  foobarFunc(r.getAs[Long]("x"), r.getAs[Double]("y"), r.getAs[String]("z"))))

val df2 = sqlContext.createDataFrame(rows, schema)

df2.show
// +---+----+---+----+-----+
// |  x|   y|  z| foo|  bar|
// +---+----+---+----+-----+
// |  1| 3.0|  a| 3.0|291.0|
// |  2|-1.0|  b|-2.0|-98.0|
// |  3| 0.0|  c| 0.0|  0.0|
// +---+----+---+----+-----+

More Related Contents:

  1. Defining a UDF that accepts an Array of objects in a Spark DataFrame?
  2. How to query JSON data column using Spark DataFrames?
  3. How to define and use a User-Defined Aggregate Function in Spark SQL?
  4. How to define partitioning of DataFrame?
  5. How do I detect if a Spark DataFrame has a column
  6. Caused by: java.lang.NullPointerException at org.apache.spark.sql.Dataset
  7. How can I pass extra parameters to UDFs in Spark SQL?
  8. SparkSQL: How to deal with null values in user defined function?
  9. How to split a dataframe into dataframes with same column values?
  10. Provide schema while reading csv file as a dataframe
  11. Dropping a nested column from Spark DataFrame
  12. Spark Dataframe :How to add a index Column : Aka Distributed Data Index
  13. How to zip two (or more) DataFrame in Spark
  14. About how to add a new column to an existing DataFrame with random values in Scala
  15. Joining Spark dataframes on the key
  16. Spark UDAF with ArrayType as bufferSchema performance issues
  17. Apache Spark how to append new column from list/array to Spark dataframe
  18. Create new Dataframe with empty/null field values
  19. Provide schema while reading csv file as a dataframe in Scala Spark
  20. Renaming column names of a DataFrame in Spark Scala
  21. Replace missing values with mean – Spark Dataframe
  22. What is going wrong with `unionAll` of Spark `DataFrame`?
  23. Apache Spark, add an “CASE WHEN … ELSE …” calculated column to an existing DataFrame
  24. How to get keys and values from MapType column in SparkSQL DataFrame
  25. Spark DataFrames when udf functions do not accept large enough input variables
  26. Filter spark DataFrame on string contains
  27. How to flatmap a nested Dataframe in Spark
  28. How to write unit tests in Spark 2.0+?
  29. How to calculate the size of dataframe in bytes in Spark?
  30. Spark column string replace when present in other column (row)

Leave a Comment