UDF with Dynamic Data Type

user_2167322 0
Viewed 174

I am trying to write the udf that can remove few keys from Map. But the Map's key and value type is not fixed, it can be String or Array or something else. How should I define such udf. I am using Spark version 2.4.4.

Below is my udf for Map[String, string]:

val mapKeys = //Seq[String]
val mapFilterUdf = udf[Map[String, String], Map[String, String]] {
    map => map.filter{case (key, _) => mapKeys.contains(key)}
}
mapFilterUdf(dataFrame.col("column_name")).as(column.name)
apache-spark-sql apache-spark user-defined-functions scala
2 Answers

you can do a generic factory method for the udf:

import scala.reflect.runtime.universe._

def filterUdfFactory[T](mapKeys:Seq[T])(implicit tag:TypeTag[T]) = udf((map:Map[T,T]) => map.filter{case (k,v) => mapKeys.contains(k)})

then use as e.g. for Strings:

val mapKeys = Seq("k1")

val tt = typeTag[String]
val filterUdf = filterUdfFactory[String](mapKeys)

 val df = Seq(
    Map("k1" -> "v1","k2" -> "v2")
 ).toDF("map")

 df.select(filterUdf($"map"))
.show()

gives:

+----------+
|  UDF(map)|
+----------+
|[k1 -> v1]|
+----------+
user_11385230

You can use Any in your UDF only if you provide the runtime-schema of your column as a second argument to udf:

val mapKeys : Seq[Any] = Seq("k1")

val df = Seq(
    Map("k1" -> "v1","k2" -> "v2")
).toDF("map")

val colSchema = df.select($"map").schema.head.dataType

val filterUdf = udf((map:Map[Any,Any]) => map.filter{case (k:Any,v:Any) => mapKeys.contains(k)},colSchema)

df
.select(filterUdf($"map"))
.show()

gives

+----------+
|  UDF(map)|
+----------+
|[k1 -> v1]|
+----------+

This work also for Row, see : https://stackoverflow.com/a/49714640/1138523

user_11385230
Related