Spark Structured Streaming 的 Stateful 操作
Structured Streaming 内部使用 StateStore 模块实现增量持续查询,和故障恢复
StateStore 模块提供了分片的、分版本的、可迁移的、高可用的 key-value store
而在应用层面主要是使用 mapGroupsWithState 和 flatMapGroupsWithState 实现状态操作
参考这篇文章的例子 https://blog.csdn.net/wangpei1949/article/details/105028892
object MapGroupsWithStateExample {
def main(args: Array[String]) {
val spark = SparkSession.builder.appName("MapGroupsWithStateExample").getOrCreate()
spark.udf.register("timezoneToTimestamp", timezoneToTimestamp _)
val jsonSchema =
"""{
"type":"struct",
"fields":[
{
"name":"eventTime",
"type":"string",
"nullable":true
},
{
"name":"eventType",
"type":"string",
"nullable":true
},
{
"name":"userID",
"type":"string",
"nullable":true
}
]
}"""
val inputTable = spark
.readStream
.format("kafka")
.option("kafka.bootstrap.servers", "kafka01:9092,kafka02:9092,kafka03:9092")
.option("subscribe", "test_1")
.load()
val resultTable = inputTable
.select(from_json(col("value").cast("string"), DataType.fromJson(jsonSchema)).as("value"))
.select($"value.*")
.withColumn("timestamp",
functions.callUDF("timezoneToTimestamp",
functions.col("eventTime"),
lit("yyyy-MM-dd HH:mm:ss"),
lit("GMT+8")))
.filter($"timestamp".isNotNull && $"eventType".isNotNull && $"userID".isNotNull)
.withWatermark("timestamp", "2 minutes")
.groupByKey((row: Row) => {
// 分钟 + userID 作为每个 group 的 key
val timestamp = row.getAs[Timestamp]("timestamp")
val sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm")
val currentEventTimeMinute = sdf.format(new Date(timestamp.getTime))
currentEventTimeMinute + "," + row.getAs[String]("userID")
})
.mapGroupsWithState[(String, Long), (String, String, Long)](GroupStateTimeout.EventTimeTimeout())(
// [] 里的 (String, Long) 和 (String, String, Long) 分别代表状态类型和返回类型
// () 里的 GroupStateTimeout.EventTimeTimeout() 指定如何判断状态超时
// 这里开始是 group map 的函数
// 接受 3 个参数: 当前 group 的 key,当前 group 的所有数据,处理这个 group 时的状态
// 每个 (分钟 + userID) 会维护一个状态
(groupKey: String, currentBatchRows: Iterator[Row], groupState: GroupState[(String, Long)]) => {
println("当前组对应的 Key: " + groupKey)
println("当前 Watermark: " + groupState.getCurrentWatermarkMs())
println("当前组的状态是否存在: " + groupState.exists)
println("当前组的状态是否过期: " + groupState.hasTimedOut)
var totalValue = 0L
if (groupState.hasTimedOut) {
// 当前组状态已过期,则清除状态
println("清除状态...")
groupState.remove()
} else if (groupState.exists) {
// 当前组状态已存在,则根据需要处理
println("增量聚合....")
// 历史值: 从状态中获取
val historyValue = groupState.get._2
// 当前值: 从当前组的新数据计算得到
val currentValue = currentBatchRows.size
// 总值 = 历史 + 当前
totalValue = historyValue + currentValue
// 更新状态
val newState = (groupKey, totalValue)
groupState.update(newState)
// 事件时间模式下,不需要设置超时时间,会根据 Watermark 机制自动超时
// 处理时间模式下,可设置个超时时间,根据超时时间清理状态,避免状态无限增加
// groupState.setTimeoutDuration(1 * 10 * 1000)
} else {
// 当前组状态不存在,则初始化状态
println("初始化状态...")
totalValue = currentBatchRows.size
val initialState = (groupKey, totalValue * 1L)
groupState.update(initialState)
}
if (totalValue != 0) {
val groupKeyArray = groupKey.split(",")
(groupKeyArray(0), groupKeyArray(1), totalValue)
} else {
null
}
}
)
.filter(_ != null)
.toDF("minute", "userID", "pv")
// Query Start
val query = resultTable
.writeStream
.format("console")
.option("truncate", "false")
.outputMode("update")
.trigger(Trigger.ProcessingTime("2 seconds"))
.start()
query.awaitTermination()
}
def timezoneToTimestamp(dateTime: String, dataTimeFormat: String, dataTimeZone: String): Timestamp = {
var output: Timestamp = null
try {
if (dateTime != null) {
val format = DateTimeFormatter.ofPattern(dataTimeFormat)
val eventTime = LocalDateTime.parse(dateTime, format).atZone(ZoneId.of(dataTimeZone));
output = new Timestamp(eventTime.toInstant.toEpochMilli)
}
} catch {
case ex: Exception => println("error")
}
output
}
}