Scala 2.13结合Spark遍历case class作为类型参数调用泛型方法问题

问题背景 #

在Spark 2.13开发中定义了如下通用方法读取Parquet数据并转换为对应case class类型的Dataset：

def getDataset[T <: Product: TypeTag](name:String): Dataset[T] = {
    import spark.implicits._

    val ds = spark.read.parquet(BASE_PATH + "/" + name).as[T]
    ds.createOrReplaceTempView(name)
    ds
}

需求是遍历一组case class，批量调用该方法加载数据，初始尝试的写法如下：

case class CLASS1(...)
case class CLASS2(...)
case class CLASS3(...)

Seq(CLASS1, CLASS2, CLASS3, ....).foreach {
  c => getDataset[c??](name=c???)
}

这里遇到的核心问题是：Seq中放入的CLASS1、CLASS2是case class的伴生对象实例，类型为() => Product，既无法直接作为类型参数传入getDataset，也无法直接拿到case class的类名字符串作为入参。

原方案优化（无额外依赖） #

你后续基于TypeTag获取类名的方案已经能解决重复编写类名字符串的问题，调整后的方法如下：

def getDataset[T <: Product: TypeTag]: Dataset[T] = {
    import spark.implicits._

    val name = typeTag[T].tpe.typeSymbol.name.toString
    val ds = spark.read.parquet(BASE_PATH + "/" + name).as[T]
    ds.createOrReplaceTempView(name)
    ds
}

调用方式为：

Seq(getDataset[CLASS1], getDataset[CLASS2], getDataset[CLASS3])

该方案无需引入第三方依赖，代码可读性高，适合case class数量不多的场景。

遍历式写法实现 #

如果希望实现最初预期的遍历调用效果，解决普通Seq丢失元素具体类型的问题，可以通过封装类型元信息的方式实现：

实现步骤 #

1. 定义元信息类，保留每个case class的TypeTag和类名

import scala.reflect.runtime.universe.TypeTag
import org.apache.spark.sql.Dataset

case class DatasetLoader[T <: Product: TypeTag]() {
  val className: String = implicitly[TypeTag[T]].tpe.typeSymbol.name.toString
  def load(): Dataset[T] = {
    import spark.implicits._
    val ds = spark.read.parquet(s"$BASE_PATH/$className").as[T]
    ds.createOrReplaceTempView(className)
    ds
  }
}

2. 构造加载器列表，保留所有类型信息

val loaders = Seq(
  DatasetLoader[CLASS1](),
  DatasetLoader[CLASS2](),
  DatasetLoader[CLASS3]()
)

3. 批量遍历加载

val datasets: Seq[Dataset[_ <: Product]] = loaders.map(_.load())

原理说明 #

普通Seq为同构集合，放入不同类型的伴生对象时，编译期会将元素类型向上统一为公共父类，丢失具体类型信息。而上述方案中每个DatasetLoader实例都绑定了对应case class的TypeTag，编译期能保留所有类型信息，运行时也可以通过TypeTag获取类名和进行类型转换。

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