将DataFrame转换为带枚举自定义字段的Dataset失败，求解决方案

问题 #

在将DataFrame转换为带有自定义枚举字段的Dataset对象时遇到报错。现有一个包含country和currency两列的DataFrame，想要转换为使用MyObj样例类的Dataset，其中currency为Scala的Enumeration枚举类型。

原代码如下：

val schema = StructType(Seq(
  StructField("country", StringType),
  StructField("currency", StringType)
))

// Define the sample data
val data = Seq(
  ("France", "EUR"),
  ("USA", "DOLLAR"),
  ("Germany", "EUR")
)

// Create a DataFrame from the sample data
val df = sparkSession.createDataFrame(data).toDF(schema.fieldNames: _*)

class Currency extends Enumeration {
  type Currency = Value
  val EUR = Value("EUR")
  val DOLLAR = Value("DOLLAR")
}

case class MyObj(country: String, currency: Currency)

val dsProduct = df.as[MyObj](Encoders.product[MyObj])

执行程序时出现错误：

Exception in thread "main" org.apache.spark.sql.AnalysisException: Try to map struct<country:string,currency:string> to Tuple1, but failed as the number of fields does not line up.

将currency改为字符串类型可正常运行，但因业务需求必须保留枚举类型，请问如何解决？

解决方案 #

Spark默认的Encoders.product无法直接处理Scala的Enumeration类型，需要手动处理类型转换，具体方案如下：

步骤1：调整枚举定义 #

把class Currency改为object Currency，Scala的Enumeration通常以单例对象形式使用，这样可以直接通过Currency.valueOf根据字符串获取枚举值，避免类实例化带来的映射问题。

步骤2：手动映射DataFrame行到样例类 #

不要直接使用as[MyObj]转换，而是通过map方法遍历每行数据，将字符串类型的currency字段手动转换为枚举值，再封装为MyObj实例。

修改后的完整代码：

import org.apache.spark.sql.{SparkSession, Encoders}
import org.apache.spark.sql.types.{StructType, StructField, StringType}

val schema = StructType(Seq(
  StructField("country", StringType),
  StructField("currency", StringType)
))

// 定义示例数据
val data = Seq(
  ("France", "EUR"),
  ("USA", "DOLLAR"),
  ("Germany", "EUR")
)

// 初始化SparkSession
val sparkSession = SparkSession.builder().master("local").appName("EnumDatasetTest").getOrCreate()
val df = sparkSession.createDataFrame(data).toDF(schema.fieldNames: _*)

// 定义枚举为单例对象
object Currency extends Enumeration {
  type Currency = Value
  val EUR = Value("EUR")
  val DOLLAR = Value("DOLLAR")
}
import Currency._ // 导入枚举类型，简化使用

// 定义目标样例类
case class MyObj(country: String, currency: Currency)

// 手动转换并生成Dataset
val dsProduct = df.map(row => {
  val country = row.getAs[String]("country")
  val currencyEnum = Currency.valueOf(row.getAs[String]("currency"))
  MyObj(country, currencyEnum)
})(Encoders.product[MyObj])

// 验证结果
dsProduct.show()

额外说明 #

• 如果遇到枚举值不存在的情况，可以添加异常处理逻辑，比如使用Currency.values.find(_.toString == currencyStr).getOrElse(...)来避免运行时错误。
• 若需要更通用的枚举类型支持，可以自定义Encoder，但上述方法对于大多数简单场景已经足够高效且易于维护。

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