Spark 3仅写入型TableProvider如何推断Schema？

一、处理首次写入的Schema推断问题 #

当首次写入时，目标路径无数据可供推断Schema，可通过以下步骤解决：

1. 路径存在性判断 + 实现Schema覆盖接口
   在inferSchema方法中检查目标路径是否存在：

   • 路径存在（追加写入）：从已有数据中读取并返回对应Schema（若需要支持追加逻辑）。
   • 路径不存在（首次写入）：返回空StructType，同时让MySpecialDataSource实现SupportsSchemaOverride接口，允许Spark用DataFrame的Schema覆盖推断的空Schema。

   修改后的代码示例：

   import org.apache.hadoop.fs.Path
   import org.apache.spark.sql.SparkSession
   import org.apache.spark.sql.catalyst.util.CaseInsensitiveStringMap
   import org.apache.spark.sql.types.StructType
   import org.apache.spark.sql.connector.catalog.SupportsSchemaOverride
   
   class MySpecialDataSource extends TableProvider with SupportsSchemaOverride {
     override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
       val path = options.get("path")
       val fs = new Path(path).getFileSystem(SparkSession.active.sparkContext.hadoopConfiguration)
       if (fs.exists(new Path(path))) {
         // 从已有自定义格式数据中读取Schema（按需实现）
         readSchemaFromExistingData(path)
       } else {
         // 首次写入返回空Schema，后续由DataFrame Schema覆盖
         StructType(Seq.empty)
       }
     }
   
     override def getTable(schema: StructType, transforms: Array[Transform], options: util.Map[String, String]): Table = {
       // 此处的schema参数即为DataFrame的Schema，直接用于创建表
       new MySpecialTable(schema, options.get("path"))
     }
   
     // 辅助方法：从已有数据读取Schema（支持追加时实现）
     private def readSchemaFromExistingData(path: String): StructType = {
       ???
     }
   }
   

2. 强制显式指定Schema（可选）
   若你的数据源要求必须显式定义Schema，可在inferSchema中抛出IllegalArgumentException，提示用户通过.schema(...)方法指定Schema，但这种方式对首次写入不够友好，仅适用于严格场景。

二、替代的抽象类实现（仅写入功能） #

如果仅需支持写入、无需读取功能，无需实现TableProvider，可直接实现DataSourceV2 + WriteSupport接口，彻底规避Schema推断逻辑：

import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, PhysicalWriteInfo, WriterCommitMessage, Write, WriteMode}
import org.apache.spark.sql.sources.DataSourceRegister
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.connector.write.DataSourceOptions

class MySpecialDataSourceV2 extends DataSourceV2 with WriteSupport with DataSourceRegister {
  override def createWriter(jobId: String, schema: StructType, mode: WriteMode, options: DataSourceOptions): Write = {
    // 直接接收DataFrame的Schema和写入路径
    new MySpecialWrite(schema, options.get("path"), mode)
  }

  override def shortName(): String = "my-special-format"
}

class MySpecialWrite(schema: StructType, path: String, mode: WriteMode) extends BatchWrite {
  override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = {
    new MySpecialWriterFactory(schema, path)
  }

  override def commit(messages: Array[WriterCommitMessage]): Unit = {
    // 全局提交逻辑：合并分区写入结果
    ???
  }

  override def abort(messages: Array[WriterCommitMessage]): Unit = {
    // 中止逻辑：清理临时数据
    ???
  }
}

class MySpecialWriterFactory(schema: StructType, path: String) extends DataWriterFactory {
  override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
    new MySpecialDataWriter(partitionId, path)
  }
}

class MySpecialDataWriter(partitionId: Int, path: String) extends DataWriter[InternalRow] {
  override def write(record: InternalRow): Unit = {
    // 单条数据写入逻辑
    ???
  }

  override def commit(): WriterCommitMessage = {
    // 分区提交信息
    new MySpecialCommitMessage(partitionId)
  }

  override def abort(): Unit = {
    // 分区写入失败清理逻辑
    ???
  }
}

class MySpecialCommitMessage(partitionId: Int) extends WriterCommitMessage

调用方式保持不变：

df.write
  .format("my.package.MySpecialDataSourceV2")
  .mode("append")
  .save("/tmp/path/to/write/to")

这种方式的优势在于：

• 无需处理Schema推断逻辑，Spark直接将DataFrame的Schema传递到写入流程。
• 代码完全聚焦于写入逻辑，避免Table相关的冗余接口实现。

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