如何让Spark UDAF自动适配DataFrame各列的原生类型？

问题背景 #

我尝试通过UDAF对DataFrame列进行处理，统计列中可能存在的数据类型分布，已实现基础功能，但目前需要为每个UDAF手动指定对应列的类型（比如my_udaf[String](col(c))），想知道有没有方法实现自动适配。

现有代码及运行结果 #

case class EstimateTypes[A]() 
    extends Aggregator[A, Map[String, Int], Map[String, Int]] with Serializable {      
    
    def zero: Map[String, Int] = Map() // Init the buffer
  
    def reduce(y: Map[String, Int], x: A) = FunEstimateType(y, x)

    def merge(a: Map[String, Int], b: Map[String, Int]) = a ++ b.map{ case (k,v) => k -> (v + a.getOrElse(k,0)) }

    def finish(r: Map[String, Int]) = r
  
    def bufferEncoder: Encoder[Map[String, Int]] = implicitly(ExpressionEncoder[Map[String, Int]])

    def outputEncoder: Encoder[Map[String, Int]] = implicitly(ExpressionEncoder[Map[String, Int]])
}

def my_udaf[A: TypeTag](col: org.apache.spark.sql.Column) = {
  val meanUdaf = udaf(new EstimateTypes[A]())
  meanUdaf(col)
}

val my_cols = Seq("id", "value")
val df = Seq(
  ("Hi", "2023-02-13"), ("Hi", "asdasd"),
  ("Hi", "2023-02-13"), ("Hi", "2023-02-13"),
).toDF("id", "value")

df.select(my_cols.map(c => my_udaf[String](col(c)).alias(c)): _*).show()

运行结果：

+-------------+--------------------+
|           id|               value|
+-------------+--------------------+
|{String -> 4}|{date -> 3, Strin...}|
+-------------+--------------------+

解决方案 #

可以通过两种方式实现自动适配，无需手动指定列类型：

方式一：利用Scala隐式推导保留类型安全 #

调整my_udaf方法的签名，让Scala自动根据输入列的类型推导泛型参数A：

import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.{Encoder, Column}
import org.apache.spark.sql.expressions.Aggregator
import scala.reflect.runtime.universe.TypeTag

case class EstimateTypes[A]() 
    extends Aggregator[A, Map[String, Int], Map[String, Int]] with Serializable {      
    
    def zero: Map[String, Int] = Map() // Init the buffer
  
    def reduce(y: Map[String, Int], x: A) = FunEstimateType(y, x)

    def merge(a: Map[String, Int], b: Map[String, Int]) = a ++ b.map{ case (k,v) => k -> (v + a.getOrElse(k,0)) }

    def finish(r: Map[String, Int]) = r
  
    def bufferEncoder: Encoder[Map[String, Int]] = implicitly(ExpressionEncoder[Map[String, Int]])

    def outputEncoder: Encoder[Map[String, Int]] = implicitly(ExpressionEncoder[Map[String, Int]])
}

// 调整方法，通过隐式参数自动推导A的类型
def my_udaf[A: TypeTag : Encoder](col: Column): Column = {
  val udafInstance = udaf(new EstimateTypes[A]())
  udafInstance(col)
}

// 使用时无需手动指定类型，直接传入列即可
val my_cols = Seq("id", "value")
val df = Seq(
  ("Hi", "2023-02-13"), ("Hi", "asdasd"),
  ("Hi", "2023-02-13"), ("Hi", "2023-02-13"),
).toDF("id", "value")

df.select(my_cols.map(c => my_udaf(col(c)).alias(c)): _*).show()

方式二：统一处理任意类型输入 #

如果你的场景中列类型不确定，可以将Aggregator的输入类型设为Any，彻底省去泛型参数：

import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.{Encoder, Column}
import org.apache.spark.sql.expressions.Aggregator

case class EstimateTypes() 
    extends Aggregator[Any, Map[String, Int], Map[String, Int]] with Serializable {      
    
    def zero: Map[String, Int] = Map() // Init the buffer
  
    def reduce(y: Map[String, Int], x: Any) = FunEstimateType(y, x)

    def merge(a: Map[String, Int], b: Map[String, Int]) = a ++ b.map{ case (k,v) => k -> (v + a.getOrElse(k,0)) }

    def finish(r: Map[String, Int]) = r
  
    def bufferEncoder: Encoder[Map[String, Int]] = implicitly(ExpressionEncoder[Map[String, Int]])

    def outputEncoder: Encoder[Map[String, Int]] = implicitly(ExpressionEncoder[Map[String, Int]])
}

def my_udaf(col: Column): Column = {
  val udafInstance = udaf(new EstimateTypes())
  udafInstance(col)
}

// 使用时直接调用，完全不需要指定类型
val my_cols = Seq("id", "value")
val df = Seq(
  ("Hi", "2023-02-13"), ("Hi", "asdasd"),
  ("Hi", "2023-02-13"), ("Hi", "2023-02-13"),
).toDF("id", "value")

df.select(my_cols.map(c => my_udaf(col(c)).alias(c)): _*).show()

说明 #

• 方式一保留了编译时类型安全，适合列类型明确的场景，Scala会自动根据列的类型推导泛型参数。
• 方式二更灵活，能处理任意类型的列，但会丢失编译时类型检查，适合列类型不确定的场景。
• 确保FunEstimateType函数可以正确处理对应类型的输入，准确识别出子类型（如字符串格式的日期）。

内容的提问来源于stack exchange，提问作者Let's try