如何让Spark将CSV中显式浮点值识别为FloatType而非DoubleType

问题场景 #

我有一个包含数百列的CSV文件，数值都显式标记为浮点型（如1.0F），但Spark 3.3自动推断其类型为DoubleType。由于数据量庞大，我希望避免对所有列逐一进行类型转换。我的场景中前5列为字符串类型，其余均为浮点型，但列的具体数量无法预先确定。

测试示例 #

$ cat test.csv
Word    Wt1     Wt2
hello   1.0F    2.0F
hi      2.0F    4.0F

Spark Shell执行结果 #

scala> val x = 2.0F
val x: Float = 2.0

scala> val df = sqlContext.read.format("csv").option("delimiter", "\t").option("header", "true").option("inferSchema", "true").csv("test.csv")
val df: org.apache.spark.sql.DataFrame = [Word: string, Wt1: double ... 1 more field]

scala> df.show()
+-----+---+---+
| Word|Wt1|Wt2|
+-----+---+---+
|hello|1.0|2.0|
|   hi|2.0|4.0|
+-----+---+---+

scala> df.dtypes
val res6: Array[(String, String)] = Array((Word,StringType), (Wt1,DoubleType), (Wt2,DoubleType))

解决方案 #

通过自定义Schema的方式避免全量类型转换：先读取表头获取所有列名，按规则生成对应类型的Schema，再用该Schema读取文件。以下是修正了类型转换问题的代码：

import org.apache.spark.sql.types.{StringType, FloatType, StructType, StructField}
import scala.collection.mutable.ListBuffer

// 读取表头获取所有列名（仅读一行空数据，性能无损耗）
val headerDf = spark.read.format("csv")
  .option("delimiter", "\t")
  .option("header", "true")
  .load("test.csv")
  .limit(0)
val columns = headerDf.columns

// 构建自定义Schema
val schemaBuffer = new ListBuffer[StructField]()
columns.zipWithIndex.foreach { case (colName, idx) =>
  val fieldType = if (idx < 5) StringType else FloatType
  schemaBuffer += StructField(colName, fieldType, nullable = true)
}
val customSchema = StructType(schemaBuffer.toArray) // 关键：将ListBuffer转为Array，解决重载错误

// 使用自定义Schema读取文件
val df = spark.read.format("csv")
  .option("delimiter", "\t")
  .option("header", "true")
  .schema(customSchema)
  .load("test.csv")

// 验证列类型
df.dtypes

执行后，非前5列的类型会被正确识别为FloatType，无需对已有数据做额外转换，避免了大数据量下的性能损耗。

补充说明 #

之前采用@leleogere的方案时出现重载错误，原因是StructType构造方法需要接收Array[StructField]类型参数，但直接传入了ListBuffer，调用toArray()将其转换为Array即可解决该问题。

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