如何修改现有DataFrame Schema并按预定义Schema写入Parquet至S3

问题概述 #

我有一个包含100+字段的CSV文件，需对这些字段转换生成80+新字段，最终仅将这些新字段（含预定义Schema中未填充的额外字段）以Parquet格式写入S3。预定义Parquet Schema包含120个字段，其中80+是已生成的新字段，剩余字段未填充，不想通过select逐个指定字段，请问能否传入预定义Schema让额外字段自动填充null？

示例CSV数据 #

aid, productId, ts, orderId
1000,100,1674128580179,edf9929a-f253-487
1001,100,1674128580179,cc41a026-63df-410
1002,100,1674128580179,9732755b-1207-471
1003,100,1674128580179,51125ddd-4129-48a
1001,200,1674128580179,f4917676-b08d-41e
1004,200,1674128580179,dc80559d-16e6-4fa
1005,200,1674128580179,c9b743eb-457b-455
1006,100,1674128580179,e8611141-3e0e-4d5
1002,200,1674128580179,30be34c7-394c-43a

预定义Parquet Schema #

def getPartitionFieldsSchema() = {
  List(
    Map("name" -> "company", "type" -> "long",
      "nullable" -> true, "metadata" -> Map()),
    Map("name" -> "epoch_day", "type" -> "long",
      "nullable" -> true, "metadata" -> Map()),
    Map("name" -> "account", "type" -> "string",
      "nullable" -> true, "metadata" -> Map()),
  )
}

val schemaMap = Map("type" -> "struct",
  "fields" -> getPartitionFieldsSchema)

原始示例代码 #

val dataDf = spark
  .read
  .format("csv")
  .option("header", "true")
  .option("inferSchema", "true")
  .load("./scripts/input.csv")


dataDf
  .withColumn("company",lit(col("aid")/100))
  .withColumn("epoch_day",lit(col("ts")/86400))
  .write   // how to write only company, epoch_day, account ?
  .mode("append")
  .csv("/tmp/data2")

解决方案 #

可以通过解析预定义Schema，自动为DataFrame补全所有缺失字段（填充null），无需逐个调用select。具体步骤如下：

1. 将SchemaMap转换为Spark StructType #

先把预定义的schemaMap转换成Spark可识别的StructType对象：

import org.apache.spark.sql.types._

// 将schemaMap转换为StructType
val targetSchema = DataType.fromJson(schemaMap.toString()).asInstanceOf[StructType]

2. 自动补全缺失字段 #

遍历targetSchema中的所有字段，对每个字段做如下处理：

• 若DataFrame已存在该字段，保留并匹配Schema定义的类型
• 若不存在，添加该字段并填充null，同时匹配Schema类型

val transformedDf = dataDf
  .withColumn("company", (col("aid") / 100).cast(LongType)) // 修正原代码：去掉多余lit，直接计算并转类型
  .withColumn("epoch_day", (col("ts") / 86400).cast(LongType))

// 自动补全所有Schema字段
val finalDf = targetSchema.foldLeft(transformedDf) { (df, field) =>
  if (df.columns.contains(field.name)) {
    df.withColumn(field.name, col(field.name).cast(field.dataType))
  } else {
    df.withColumn(field.name, lit(null).cast(field.dataType))
  }
}

3. 写入Parquet到S3 #

直接使用finalDf写入Parquet，输出会严格遵循预定义Schema，缺失字段自动填充null：

finalDf
  .write
  .mode("append")
  .parquet("s3://your-bucket/path/to/output")

关键说明 #

• 无需手动逐个指定字段，通过遍历Schema自动处理所有字段，适配字段数量多的场景
• 强制字段类型与预定义Schema匹配，避免类型不一致问题
• 缺失字段自动填充null，完全满足需求

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