Spark DataFrame添加枚举类型列及转含枚举的样例类列表方法咨询

注意：Test样例类的属性名type是Scala保留关键字，定义时需要用反引号包裹，否则会触发语法错误，样例类参考定义如下：

object SaleType extends Enumeration {
  val SALE, PURCHASE = Value
}
case class Test(id: Int, name: String, `type`: SaleType.Value)

方案1：收集后手动映射（实现最简单，适配小数据量场景） #

不需要提前在DataFrame中新增枚举类型列，直接收集原始DataFrame的Row数据，遍历过程中根据业务规则赋值枚举即可，完全规避Spark DataFrame不支持自定义枚举作为列类型的问题：

import org.apache.spark.sql.functions._

val df = Seq(
  (0, "p"),
  (1, "s")
).toDF("id", "name")

val result: List[Test] = df.collect().map { row =>
  val id = row.getAs[Int]("id")
  val name = row.getAs[String]("name")
  val saleType = if (name == "p") SaleType.PURCHASE else SaleType.SALE
  Test(id, name, saleType)
}.toList

如果你的业务逻辑已经提前生成了存储枚举名称的字符串列，也可以用withName方法直接映射枚举：

val dfWithStrType = df.withColumn("type_str", when(col("name") === "p", "PURCHASE").otherwise("SALE"))
val result: List[Test] = dfWithStrType.collect().map { row =>
  Test(
    row.getAs[Int]("id"),
    row.getAs[String]("name"),
    SaleType.withName(row.getAs[String]("type_str"))
  )
}.toList

方案2：自定义Encoder实现Dataset直接转换（适配大数据量场景） #

如果后续还需要对Test类型的数据集做算子操作、不想提前collect到驱动端，可以显式为Test类定义Kryo编码器，在Dataset层面完成类型转换：

import org.apache.spark.sql.Encoders

implicit val testEncoder = Encoders.kryo[Test]

val testDs = df.map { row =>
  val id = row.getAs[Int]("id")
  val name = row.getAs[String]("name")
  val saleType = if (name == "p") SaleType.PURCHASE else SaleType.SALE
  Test(id, name, saleType)
}
// 需要时再collect转List
val result: List[Test] = testDs.collect().toList

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