Create a new column in Spark DataFrame using UDF

Question

I have a UDF as below -

  val myUdf = udf((col_abc: String, col_xyz: String) => {
    array(
      struct(
        lit("x").alias("col1"),
        col(col_abc).alias("col2"),
        col(col_xyz).alias("col3")
      )
    )
  }

Now, I want to use this use this in a function as below -

def myfunc(): Column = {
    val myvariable = myUdf($"col_abc", $"col_xyz")
    myvariable
}

And then use this function to create a new column in my DataFrame

val df = df..withColumn("new_col", myfunc())

In summary, I want my column "new_col" to be an type array with values as [[x, x, x]]

I am getting the below error. What am I doing wrong here?

Caused by: java.lang.UnsupportedOperationException: Schema for type org.apache.spark.sql.Column is not supported

mck mck · Accepted Answer · 2021-03-15T10:11:20

Two ways.

Don't use a UDF because you're using pure Spark functions:

val myUdf = ((col_abc: String, col_xyz: String) => {
    array(
      struct(
        lit("x").alias("col1"),
        col(col_abc).alias("col2"),
        col(col_xyz).alias("col3")
      )
    )
  }
)

def myfunc(): Column = {
    val myvariable = myUdf("col_abc", "col_xyz")
    myvariable
}

df.withColumn("new_col", myfunc()).show
+-------+-------+---------------+
|col_abc|col_xyz|        new_col|
+-------+-------+---------------+
|    abc|    xyz|[[x, abc, xyz]]|
+-------+-------+---------------+

Use a UDF which takes in strings and returns a Seq of case class:

case class cols (col1: String, col2: String, col3: String)

val myUdf = udf((col_abc: String, col_xyz: String) => Seq(cols("x", col_abc, col_xyz)))

def myfunc(): Column = {
    val myvariable = myUdf($"col_abc", $"col_xyz")
    myvariable
}

df.withColumn("new_col", myfunc()).show
+-------+-------+---------------+
|col_abc|col_xyz|        new_col|
+-------+-------+---------------+
|    abc|    xyz|[[x, abc, xyz]]|
+-------+-------+---------------+

If you want to pass in Columns to the function, here's an example:

val myUdf = ((col_abc: Column, col_xyz: Column) => {
    array(
      struct(
        lit("x").alias("col1"),
        col_abc.alias("col2"),
        col_xyz.alias("col3")
      )
    )
  }
)

Create a new column in Spark DataFrame using UDF

1 Answers