Spark Structured Streaming如何设置CosmosDB微批处理条目数

问题背景 #

使用Spark Structured Streaming从Cosmos DB读取传感器数据，共24个传感器，采集频率为1次/秒，数据转换处理后需要调用MLFlow分类模型执行推理，要求每个微批恰好包含24条输入数据（或24的整数倍条数据）。
此前尝试在读流配置中使用ItemCountPerTriggerHint、limit、maxItemCount相关参数，也通过trigger(processingTime='x seconds')调整触发间隔控制处理速率，代码可正常无报错运行，但上述配置均未对批DataFrame大小产生实际效果，运行时numInputRows始终在3到100之间随机波动。

当前实现代码 #

"spark.cosmos.accountEndpoint" : cosmosEndpoint,
"spark.cosmos.accountKey" : cosmosMasterKey,
"spark.cosmos.database" : cosmosDatabaseName,
"spark.cosmos.container" : cosmosContainerName,
"spark.cosmos.upsert" : "true"
}

# Configure Catalog Api to be used
spark.conf.set("spark.sql.catalog.cosmosCatalog", "com.azure.cosmos.spark.CosmosCatalog")
spark.conf.set("spark.sql.catalog.cosmosCatalog.spark.cosmos.accountEndpoint", cosmosEndpoint)
spark.conf.set("spark.sql.catalog.cosmosCatalog.spark.cosmos.accountKey", cosmosMasterKey)

# Initiate Cosmos Connection Config Object
changeFeedCfg = {
  "spark.cosmos.accountEndpoint": cosmosEndpoint,
  "spark.cosmos.accountKey": cosmosMasterKey,
  "spark.cosmos.database": cosmosDatabaseName,
  "spark.cosmos.container": cosmosContainerName,
  "spark.cosmos.read.partitioning.strategy": "Default",
  "spark.cosmos.read.inferSchema.enabled" : "false",
  "spark.cosmos.changeFeed.startFrom" : "Now",
  "spark.cosmos.changeFeed.mode" : "Incremental",
  "spark.cosmos.changeFeed.ItemCountPerTriggerHint" : 24,
}

# Load model as a PysparkUDF
loaded_model = mlflow.pyfunc.spark_udf(spark, model_uri='runs:/*********/model', result_type='double')
literal_eval_udf = udf(ast.literal_eval, MapType(StringType(), StringType()))

fixedStream = spark.readStream.format("cosmos.oltp.changeFeed").options(**changeFeedCfg).load()

fixedStream = fixedStream.select('_rawBody').withColumn('temp', regexp_replace('_rawBody', ',"_rid".*', '}')).drop('_rawBody')
fixedStream = fixedStream.withColumn("temp", map_values(literal_eval_udf(col("temp"))))
keys = ['datetime', 'machine', 'id', 'factor', 'value', 'Sensor']
for k in range(len(keys)):
    fixedStream = fixedStream.withColumn(keys[k], fixedStream.temp[k])
fixedStream = fixedStream.select('factor','machine','Sensor','value')

def foreach_batch_function(df, epoch_id):
    df = df.groupBy('factor','machine').pivot("Sensor").agg(first("value"))
    columns = list(df)
    df = df.withColumn('predictions', loaded_model(*columns)).collect()
    df.write.option("mergeSchema","true").format("delta").option("header", "true").mode("append").saveAsTable("poc_industry.test_stream")
    
fixedStream.writeStream.foreachBatch(foreach_batch_function).start()

问题根因 #

此前尝试的参数均无法控制微批大小，核心原因如下：

• spark.cosmos.changeFeed.ItemCountPerTriggerHint属于提示性参数，无强制约束力：Cosmos DB Change Feed按物理分区并行拉取数据，每个分区返回的条目数不受该参数强制限制，多分区返回条目累加后总数自然随机波动。
• 流计算中的limit()是分区级截断逻辑，不会做全局计数校验，无法保证整个微批的总条数符合要求。
• trigger(processingTime)仅控制查询触发的时间间隔，数据上报速率、网络波动、Cosmos DB端返回效率都会影响单个时间窗口内拉取的总条数，不可能稳定维持在固定值。

可行解决方案 #

方案1：Structured Streaming 状态管理攒批（推荐，无侵入、容错性好） #

不依赖数据源端的条数控制，在Spark计算层通过原生状态管理实现全局固定条数切批，容错由Spark Checkpoint机制保证，不会丢数也不会重复计算。
实现逻辑：给所有流数据分配统一的全局路由key，通过flatMapGroupsWithState维护跨微批的缓存状态，每次新数据进入后先存入缓存，当缓存条数达到24的整数倍时，将整批数据输出到后续处理逻辑，不足24条的残留在状态中等待下一个微批凑数。
参考实现代码：

from pyspark.sql.types import *
from pyspark.sql.functions import *
from pyspark.sql.streaming import GroupStateTimeout

# 配置固定批大小
BATCH_SIZE = 24
# 攒批操作的并发设为1，保证全局计数准确
spark.conf.set("spark.sql.shuffle.partitions", "1")

# 定义状态结构：存储跨微批缓存的待处理记录
buffer_record_schema = StructType([
    StructField("factor", StringType()),
    StructField("machine", StringType()),
    StructField("Sensor", StringType()),
    StructField("value", StringType())
])
state_schema = StructType([
    StructField("buffered_records", ArrayType(buffer_record_schema))
])

def fixed_size_batch_func(key, rows_iter, state):
    # 读取历史缓存状态
    if state.exists():
        buffered = state.get().buffered_records
    else:
        buffered = []
    
    # 将当前微批新数据加入缓存
    for row in rows_iter:
        buffered.append(row.asDict())
    
    output = []
    # 凑够整批就输出
    while len(buffered) >= BATCH_SIZE:
        current_batch = buffered[:BATCH_SIZE]
        buffered = buffered[BATCH_SIZE:]
        output.extend(current_batch)
    
    # 更新状态：保留不足一批的残留数据
    if len(buffered) > 0:
        state.update((buffered,))
    else:
        if state.exists():
            state.remove()
    
    return iter(output)

# 对原始流应用固定批切分逻辑
fixed_size_stream = fixedStream \
    .withColumn("global_key", lit(1)) \
    .groupBy("global_key") \
    .flatMapGroupsWithState(
        outputMode="append",
        timeoutConf=GroupStateTimeout.NoTimeout,
        func=fixed_size_batch_func,
        stateSchema=state_schema
    ) \
    .drop("global_key")

# 后续写入逻辑替换为fixed_size_stream即可
def foreach_batch_function(df, epoch_id):
    df = df.groupBy('factor','machine').pivot("Sensor").agg(first("value"))
    columns = list(df)
    df = df.withColumn('predictions', loaded_model(*columns))
    df.write.option("mergeSchema","true").format("delta").option("header", "true").mode("append").saveAsTable("poc_industry.test_stream")

fixed_size_stream.writeStream.foreachBatch(foreach_batch_function).start()

注意：该方案本质是攒批，会产生最多1秒左右的延迟（和传感器采集频率匹配），完全满足工业场景推理的时延要求。

方案2：调整Cosmos DB分区从源头控量 #

如果不想在计算层做攒批逻辑，可以调整Cosmos DB容器配置：

• 将分区键设置为秒级采集时间戳 + 设备ID，保证每一秒24个传感器的上报数据全部落在同一个逻辑分区
• 读流配置中增加"spark.cosmos.read.maxPartitionsPerTrigger": "1"，限制每个微批最多拉取1个分区的数据
• 保留"spark.cosmos.changeFeed.ItemCountPerTriggerHint": 24配置
  该方案限制较大：已上线的Cosmos DB容器修改分区键需要全量导数据，成本较高，适合新建容器的场景。

方案3：驱动端/外部缓存攒批 #

在foreachBatch逻辑中实现攒批：

• 用驱动端内存变量或者外部持久化缓存（如Redis）维护全局待处理队列
• 每个微批进入后先把数据写入队列，再从队列中取出N*24条数据做推理和写入，不足24条的留在队列等待下一批
  该方案实现简单，但容错性差：如果用驱动端内存存缓存，驱动重启会丢数；如果用外部缓存需要额外实现幂等写入逻辑，避免故障恢复后重复计算。

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