Spark 2.4.0批处理消费Kafka：如何编程检测新分区并更新起始偏移

问题原因 #

你遇到的错误是因为Spark批处理模式（spark.read）下，当指定具体分区偏移量（而非全局的earliest/latest）时，要求必须覆盖Kafka Topic当前所有的分区。而官方文档中关于Newly discovered partitions during a query will start at earliest的描述仅适用于流处理模式（readStream），批处理模式不会自动处理新增分区，因此需要手动检测并更新startingOffsets参数。

解决方案步骤 #

核心思路：先获取Kafka当前所有分区，对比Checkpoint中记录的已消费分区，找出新增分区后，将新增分区的起始偏移量设为earliest（或latest，按需调整），再合并成完整的startingOffsets字符串。

1. 获取Kafka Topic的当前所有分区 #

使用Kafka AdminClient API获取指定Topic的所有分区信息：

import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
import org.apache.kafka.common.TopicPartition
import scala.collection.JavaConverters._

def getKafkaTopicPartitions(kafkaBootstrapServers: String, topic: String): Set[TopicPartition] = {
  val props = new java.util.Properties()
  props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaBootstrapServers)
  val adminClient = AdminClient.create(props)
  try {
    val topicDescriptions = adminClient.describeTopics(Set(topic).asJava).all().get()
    topicDescriptions.get(topic)
      .partitions().asScala
      .map(p => new TopicPartition(topic, p.partition()))
      .toSet
  } finally {
    adminClient.close()
  }
}

2. 读取Checkpoint中记录的已消费分区和偏移量 #

Spark的Checkpoint会将偏移量存储在checkpointDir/offsets目录下的文件中，解析这些文件获取历史偏移量：

import org.apache.spark.sql.execution.streaming.OffsetSeq
import com.fasterxml.jackson.databind.ObjectMapper
import com.fasterxml.jackson.module.scala.DefaultScalaModule
import org.apache.spark.sql.kafka010.KafkaSourceOffset
import java.io.File

def getCheckpointOffsets(checkpointDir: String, topic: String): Map[TopicPartition, Long] = {
  val offsetsDir = new File(s"$checkpointDir/offsets")
  if (!offsetsDir.exists() || offsetsDir.listFiles().isEmpty) {
    // 首次运行，无历史偏移量
    Map.empty
  } else {
    // 取最新的偏移量文件（按批次号文件名排序）
    val latestOffsetFile = offsetsDir.listFiles().sortBy(_.getName.toLong).last
    val mapper = new ObjectMapper().registerModule(DefaultScalaModule)
    val offsetSeq = mapper.readValue(latestOffsetFile, classOf[OffsetSeq])
    
    offsetSeq.offsets.flatMap {
      case Some(kafkaOffset: KafkaSourceOffset) =>
        kafkaOffset.partitionToOffset.filter(_._1.topic() == topic)
      case _ => None
    }.toMap
  }
}

3. 生成更新后的startingOffsets参数 #

对比当前分区和历史分区，合并偏移量并转换成Spark要求的JSON格式：

def generateStartingOffsets(kafkaBootstrapServers: String, topic: String, checkpointDir: String): String = {
  val currentPartitions = getKafkaTopicPartitions(kafkaBootstrapServers, topic)
  val checkpointOffsets = getCheckpointOffsets(checkpointDir, topic)
  val existingPartitions = checkpointOffsets.keys.toSet
  val newPartitions = currentPartitions -- existingPartitions

  // 新增分区设为earliest，可根据业务需求改为latest
  val newPartitionOffsets = newPartitions.map(tp => tp -> "earliest").toMap
  // 合并历史偏移量（转成字符串）和新增分区偏移量
  val combinedOffsets = checkpointOffsets.map { case (tp, offset) => tp -> offset.toString } ++ newPartitionOffsets

  // 转换成Spark所需的JSON格式：{"topicName":{"0":"100","1":"earliest"}}
  val topicOffsetMap = combinedOffsets.groupBy(_._1.topic()).map { case (t, tpOffsets) =>
    t -> tpOffsets.map { case (tp, off) => tp.partition().toString -> off }.toMap
  }
  
  new ObjectMapper().writeValueAsString(topicOffsetMap)
}

4. 在Spark批处理作业中使用 #

调用上述方法生成startingOffsets，传入Kafka读取配置：

val kafkaBootstrapServers = "your-kafka-broker:9092"
val targetTopic = "your-target-topic"
val checkpointDirectory = "/path/to/your/checkpoint"

val startingOffsets = generateStartingOffsets(kafkaBootstrapServers, targetTopic, checkpointDirectory)

val kafkaDF = spark.read
  .format("kafka")
  .option("kafka.bootstrap.servers", kafkaBootstrapServers)
  .option("subscribe", targetTopic)
  .option("startingOffsets", startingOffsets)
  // 可添加其他配置，如endingOffsets等
  .load()

注意事项 #

• 依赖兼容：Spark 2.4.0默认依赖kafka-clients 2.0.0，确保作业的依赖版本一致，避免API兼容问题。
• 权限配置：Spark作业需要具备读取Checkpoint目录的权限，以及访问Kafka集群获取元数据的权限。
• 偏移量策略：新增分区的起始偏移量可根据业务需求选择earliest（消费所有历史数据）或latest（仅消费新增分区的后续数据）。

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