Spark UDF as function parameter, UDF is not in function scope

Question

I have a few UDFs that I'd like to pass along as a function argument along with data frames.

One way to do this might be to create the UDF within the function, but that would create and destroy several instances of the UDF without reusing it which might not be the best way to approach this problem.

Here's a sample piece of code -

val lkpUDF = udf{(i: Int) => if (i > 0) 1 else 0}

val df =   inputDF1
    .withColumn("new_col", lkpUDF(col("c1")))
val df2 =   inputDF2.
  .withColumn("new_col", lkpUDF(col("c1")))

Instead of doing the above, I'd ideally want to do something like this -

val lkpUDF = udf{(i: Int) => if (i > 0) 1 else 0}

def appendCols(df: DataFrame, lkpUDF: ?): DataFrame = {

    df
      .withColumn("new_col", lkpUDF(col("c1")))

  }
val df = appendCols(inputDF, lkpUDF)

The above UDF is pretty simple, but in my case it can return a primitive type or a user defined case class type. Any thoughts/ pointers would be much appreciated. Thanks.

Answer 1

Your function with the appropriate signature needs to be this:

import org.apache.spark.sql.UserDefinedFunction

def appendCols(df: DataFrame, func: UserDefinedFunction): DataFrame = {
    df.withColumn("new_col", func(col("col1")))
}

The scala REPL is quite helpful in returning the type of the values initialized.

scala> val lkpUDF = udf{(i: Int) => if (i > 0) 1 else 0}
lkpUDF: org.apache.spark.sql.UserDefinedFunction = UserDefinedFunction(<function1>,IntegerType,List(IntegerType))

Also, if the signature of the function that you pass into the udf wrapper consists of an Any return type (which will be the case if the function can return either a primitive or a user defined case class), the UDF will fail to compile with an exception like so:

java.lang.UnsupportedOperationException: Schema for type Any is not supported