Spark 2.4.4 如何定义支持动态数据类型Map的过滤UDF

核心问题说明 #

原有代码仅支持Map[String, String]类型的核心限制是硬编码了泛型参数，要兼容任意键值类型的Map，需要解决Scala类型擦除和Spark编码器匹配的问题，以下是两种可直接使用的实现方案：

方案1：带泛型的类型安全UDF #

适合需要编译期类型校验的场景，实现如下：
首先导入依赖包：

import scala.reflect.runtime.universe.TypeTag
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.Column

定义通用UDF生成方法：

def getMapFilterUDF[K: TypeTag, V: TypeTag](allowedKeys: Seq[K]): Column => Column = {
  val filterUdf = udf[Map[K, V], Map[K, V]] { map =>
    if (map == null) null else map.filter { case (key, _) => allowedKeys.contains(key) }
  }
  col => filterUdf(col)
}

调用示例：

// 处理Map[String, Int]类型列
val allowedStrKeys = Seq("k1", "k2", "k3")
val strMapFilter = getMapFilterUDF[String, Int](allowedStrKeys)
df.select(strMapFilter(df.col("map_column")).as("filtered_map"))

// 处理Map[Array[String], Double]类型列
val allowedArrayKeys = Seq(Array("a", "b"), Array("c", "d"))
val arrayMapFilter = getMapFilterUDF[Array[String], Double](allowedArrayKeys)
df.select(arrayMapFilter(df.col("complex_map_column")).as("filtered_complex_map"))

方案2：无类型通用UDF #

无需显式指定泛型，兼容所有Map键值类型，实现如下：

import org.apache.spark.sql.functions.udf
import scala.collection.JavaConverters._

val allowedKeys: Seq[Any] = Seq("test_key", 1024, Array("x", "y"))
val universalMapFilterUdf = udf((map: java.util.Map[Any, Any]) => {
  if (map == null) null
  map.asScala.filter { case (key, _) => allowedKeys.contains(key) }.asJava
})

调用示例：

df.select(universalMapFilterUdf(df.col("any_map_column")).as("filtered_map"))

最优方案：使用Spark内置函数 #

Spark 2.4已经内置了map_filter函数，性能远高于自定义UDF，无需额外开发，直接调用即可：

import org.apache.spark.sql.functions.map_filter

val mapKeys = Seq("k1", "k2", "k3")
df.select(
  map_filter(df.col("map_column"), (key, value) => mapKeys.contains(key)).as("filtered_map")
)

内容的提问来源于stack exchange，提问作者user2167322