PySpark Streaming Kafka手动获取与提交偏移量问题求助

问题背景 #

现有PySpark Streaming代码以earliest模式从Kafka读取流数据，每日有新数据流入，需求是每天从最后提交的偏移量开始拉取数据，需解决两个核心问题：

1. 如何通过PySpark代码手动获取最后提交的偏移量？
2. 如何手动提交偏移量？

用户现有代码：

df = spark.readStream.format("kafka")\
.option("kafka.bootstrap.servers", serveurs)\
.option("subscribe", topic_name) \
.option("startingOffsets", "earliest")\
.option("kafka.group.id" , "consumer-id")\
.option("kafka.security.protocol", "SSL")\
.option("kafka.ssl.truststore.location",  truststore_location ) \
.option("kafka.ssl.truststore.type", "JKS") \
.option("kafka.ssl.truststore.password", truststore_pwd) \
.option("kafka.ssl.keystore.location", keystore_location ) \
.option("kafka.ssl.keystore.type", "JKS") \
.option("kafka.ssl.keystore.password", keystore_pwd) \
.option("enable.auto.commit", "TRUE") \
.load()\
.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")

query = df\
.select("value")\
.writeStream\
.queryName("event")\
.outputMode("append")\
.format("parquet")\
.option("path", HDFS_LOCATION + '/' + str(date.today()))\
.option("checkpointLocation", HDFS_LOCATION + '/' + str(date.today()))\
.trigger(processingTime='30 seconds')\
.option("truncate", False)\
.start()\
.awaitTermination()

一、手动获取最后提交的偏移量 #

PySpark没有直接提供查询Kafka已提交偏移量的API，可通过以下两种方式实现：

方式1：Kafka 消费者API查询（推荐） #

通过Kafka官方客户端查询指定消费者组的已提交偏移量，代码示例：

from kafka import KafkaConsumer
from kafka.structs import TopicPartition
import json

# 复用Kafka连接参数
kafka_params = {
    "bootstrap_servers": serveurs,
    "security_protocol": "SSL",
    "ssl_truststore_location": truststore_location,
    "ssl_truststore_type": "JKS",
    "ssl_truststore_password": truststore_pwd,
    "ssl_keystore_location": keystore_location,
    "ssl_keystore_type": "JKS",
    "ssl_keystore_password": keystore_pwd
}

# 初始化消费者（仅用于查询，不消费数据）
consumer = KafkaConsumer(group_id="consumer-id", **kafka_params)

# 获取topic的所有分区
partitions = [TopicPartition(topic_name, p) for p in consumer.partitions_for_topic(topic_name)]

# 查询已提交的偏移量
committed = consumer.committed(partitions)

# 转换为Spark startingOffsets需要的JSON格式
starting_offsets = {}
for p, offset in zip(partitions, committed):
    # 若未提交过偏移量，默认从0开始
    starting_offsets[f"{topic_name}/{p.partition}"] = offset.offset if offset else 0

starting_offsets_json = json.dumps(starting_offsets)

注意：集群环境需确保所有节点安装kafka-python包；若无法安装，可通过Py4J调用Java的AdminClient实现，逻辑一致。

方式2：从Spark Checkpoint目录读取 #

Structured Streaming会将偏移量存储在checkpoint目录的offsets子目录下，可直接读取该目录的JSON文件：

import json

# 读取checkpoint中的偏移量文件
offset_data = spark.read.text(f"{HDFS_LOCATION}/{str(date.today())}/offsets").collect()
if offset_data:
    offset_json = json.loads(offset_data[0][0])
    kafka_offsets = offset_json.get("sources", [{}])[0].get("startOffset", {})
    starting_offsets_json = json.dumps(kafka_offsets)

注意：该方式依赖Spark checkpoint的内部格式，版本升级可能导致格式变化，稳定性不如Kafka客户端查询。

二、手动提交偏移量 #

要实现手动提交，需先关闭自动提交，再在批次处理完成后提交偏移量，推荐使用foreachBatch算子实现自定义提交逻辑：

from kafka import KafkaConsumer
from kafka.structs import TopicPartition, OffsetAndMetadata

def commit_kafka_offsets(batch_df, batch_id):
    # 获取当前批次的偏移量范围
    offset_ranges = batch_df.select("partition", "offset").distinct().collect()
    
    # 初始化消费者（关闭自动提交）
    kafka_params = {
        "bootstrap_servers": serveurs,
        "security_protocol": "SSL",
        "ssl_truststore_location": truststore_location,
        "ssl_truststore_type": "JKS",
        "ssl_truststore_password": truststore_pwd,
        "ssl_keystore_location": keystore_location,
        "ssl_keystore_type": "JKS",
        "ssl_keystore_password": keystore_pwd
    }
    consumer = KafkaConsumer(group_id="consumer-id", enable_auto_commit=False, **kafka_params)
    
    # 构造待提交的偏移量（偏移量+1，确保下一次从当前批次结束位置开始）
    commit_dict = {}
    for row in offset_ranges:
        tp = TopicPartition(topic_name, row.partition)
        commit_dict[tp] = OffsetAndMetadata(row.offset + 1, None)
    
    # 提交偏移量
    consumer.commit(offsets=commit_dict)
    consumer.close()

# 修改流读取配置：关闭自动提交，设置startingOffsets为查询到的偏移量
df = spark.readStream.format("kafka")\
.option("kafka.bootstrap.servers", serveurs)\
.option("subscribe", topic_name) \
.option("startingOffsets", starting_offsets_json)  # 使用查询到的已提交偏移量
.option("kafka.group.id" , "consumer-id")\
.option("kafka.security.protocol", "SSL")\
.option("kafka.ssl.truststore.location",  truststore_location ) \
.option("kafka.ssl.truststore.type", "JKS") \
.option("kafka.ssl.truststore.password", truststore_pwd) \
.option("kafka.ssl.keystore.location", keystore_location ) \
.option("kafka.ssl.keystore.type", "JKS") \
.option("kafka.ssl.keystore.password", keystore_pwd) \
.option("enable.auto.commit", "FALSE")  # 关闭自动提交
.load()\
.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "offset", "partition")  # 保留偏移量和分区字段

# 修改流写入逻辑，添加foreachBatch提交偏移量
query = df\
.writeStream\
.queryName("event")\
.outputMode("append")\
.format("parquet")\
.option("path", f"{HDFS_LOCATION}/{str(date.today())}")\
.option("checkpointLocation", f"{HDFS_LOCATION}/{str(date.today())}")\
.trigger(processingTime='30 seconds')\
.option("truncate", False)\
.foreachBatch(commit_kafka_offsets)  # 自定义偏移量提交逻辑
.start()\
.awaitTermination()

关键说明：

• 必须关闭enable.auto.commit，避免自动提交与手动提交冲突
• 提交的偏移量需为当前批次最大偏移量+1，确保下一次读取从正确位置开始
• 需在读取流时保留offset和partition字段，用于获取批次偏移范围

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