如何编写Spark Encoder自动将继承Base的Case Class转为Struct？

问题背景 #

希望在Spark DataFrame中存储不同关联类型，同时通过DataSet使用强类型Case Class。定义了如下Trait和Case Class：

trait Base {
  def name: String
}

case class A(name: String, number: Int) extends Base

case class B(name: String, text: String) extends Base

尝试创建val lb = List[Base](A("Alice", 20), B("Bob", "Foo"))并调用lb.toDS()生成DataSet时，会因Trait及其子类无对应Encoder而失败。

当前手动解决方案是定义包含所有字段的Struct类，并实现Base与Struct的互转：

case class Struct(typ: String, name: String, number: Option[Int] = None, text: Option[String] = None)

trait Base {
  def name: String

  def asStruct: Struct = {
    this match {
      case A(name, number) => Struct("A", name, number = Some(number))
      case B(name, text) => Struct("B", name, text = Some(text))
    }
  }
}

case class Struct(typ: String, name: String, number: Option[Int] = None, text: Option[String] = None) {
  def asBase: Base = {
    this match {
      case Struct("A", name, Some(number), None) => A(name, number)
      case Struct("B", name, None, Some(text)) => B(name, text)
      case _ => throw new Exception(s"Invalid Base structure $this")
    }
  }
}

之后通过转换为Struct列表创建DataSet：

val a = A("Alice", 32)
val b = B("Bob", "foo")

val ls = List[Struct](a.asStruct, b.asStruct)

val sparkSession = spark
import sparkSession.implicits._

val df = ls.toDS()

df.show()

输出结果：

+---+-----+------+----+
|typ| name|number|text|
+---+-----+------+----+
|  A|Alice|    32|NULL|
|  B|  Bob|  NULL| foo|
+---+-----+------+----+

现在需要知道：能否编写自定义Encoder，自动通过上述asStruct/asBase逻辑，直接将Base类视为Struct处理，无需手动转换？

解决方案：自定义Base类型的Encoder #

可以复用Struct类型的Encoder，结合Base与Struct的转换逻辑，实现自定义Encoder[Base]，具体步骤如下：

1. 保留原有类型定义 #

维持Base、A、B和Struct的定义不变，确保asStruct和asBase转换方法可用。

2. 实现自定义Encoder #

利用Spark隐式提供的Struct类型Encoder，通过映射转换得到Base的Encoder：

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder

implicit val baseEncoder: Encoder[Base] = {
  // 获取Struct类型的Encoder
  val structEncoder: Encoder[Struct] = Encoders.product[Struct]
  // 映射Struct和Base之间的转换逻辑
  structEncoder.map(
    struct => struct.asBase,
    base => base.asStruct
  )
}

3. 直接创建Base类型的DataSet #

现在可直接将List[Base]转换为DataSet，无需手动转换为Struct：

val sparkSession = spark
import sparkSession.implicits._

val lb = List[Base](A("Alice", 20), B("Bob", "Foo"))
val ds = lb.toDS()

ds.show()

输出结果与手动转换一致：

+---+-----+------+----+
|typ| name|number|text|
+---+-----+------+----+
|  A|Alice|    20|NULL|
|  B|  Bob|  NULL| Foo|
+---+-----+------+----+

4. 验证反向转换 #

从DataSet读取数据时，会自动转换回对应的Base子类实例：

ds.collect().foreach {
  case A(name, num) => println(s"Found A: $name, $num")
  case B(name, txt) => println(s"Found B: $name, $txt")
}

输出：

Found A: Alice, 20
Found B: Bob, Foo

原理说明 #

• Spark的ExpressionEncoder.map方法允许基于现有Encoder，通过提供正向（T => U）和反向（U => T）转换函数，生成新类型的Encoder。
• 这里复用Struct的Encoder处理序列化和反序列化逻辑，只需关注Base与Struct之间的类型转换，无需手动处理Schema和Row的映射。

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