Spark查询CosmosDB小数据集时控制并行度优化RU成本

你的判断完全符合实际触发逻辑，该重复请求现象和Cosmos DB本身的分区设计无关，核心是Cosmos DB Spark连接器的默认分片逻辑在小结果集场景下的适配问题：

• 连接器默认会结合Cosmos DB集合物理分区数、当前集群可用executor核数生成输入分片，每个分片对应一个独立的查询任务，调度到不同executor执行
• 当集合数据量极小（仅数百条文档）时，连接器的分片范围估算逻辑不会自动做小结果集合并，会生成和可用executor数量匹配的分片，且每个分片携带的查询条件完全一致——连接器默认假设每个分片只拉取对应范围的数据，但小数据集下范围切分规则失效，最终导致所有executor拿到的都是全量查询任务，重复向Cosmos DB发起完全相同的跨分区请求
• 该逻辑在Spark 2.4、Spark 3版本的Cosmos DB连接器，以及HDInsight、Synapse托管环境的默认配置下均会触发，和你观测到的executor数量与请求数正相关、多节点同一时间发起相同请求的现象完全吻合。调整集合分区键属于服务端优化，无法解决连接器侧的任务拆分问题。

不需要将整个Spark作业降级为单节点运行，仅针对该特定小查询做单独配置即可，不影响作业其他操作的多节点算力使用，以下方案按落地便捷度排序：

方案1：针对该查询单独设置输入分区数为1（最推荐） #

仅在执行这条小查询的读取配置中，强制连接器只生成1个查询分片，该查询就只会调度到1个executor上执行，从根源上避免重复请求。
不同版本连接器对应参数如下：

• Spark 2.4版本Cosmos DB连接器：配置spark.cosmosdb.inputpartitioncount参数值为1
• Spark 3版本Cosmos DB连接器：配置spark.cosmos.read.partitioning.targetedPartitionCount参数值为1
  注意该配置仅在读取该小集合时传入即可，不要设置为全局Spark参数，避免影响其他大查询的并行读取效率。

方案2：将小查询结果封装为广播变量 #

该查询结果集仅数百条，体量极小，可在读取数据后直接通过广播变量封装，Spark会在Driver端拉取全量数据后分发到所有executor内存，后续作业复用该数据时不会再重复触发Cosmos DB请求。
参考代码：

// 读取Cosmos DB小集合数据
val smallSet = spark.read.format("cosmos.oltp")
  // 配置连接信息、自定义查询语句
  .load()
  .where("Mode != 'anc'")
// 封装为广播变量
val broadcastSmallSet = spark.sparkContext.broadcast(smallSet.collect())
// 后续逻辑直接调用broadcastSmallSet.value即可使用数据

该方案无需调整连接器配置，仅对目标小数据集做单独处理，对现有作业逻辑侵入极低。

方案3：单分区读取后缓存结果 #

如果不想调整连接器参数或使用广播变量，可在读取该集合后通过coalesce(1)强制合并为单分区，同时加缓存确保数据只拉取一次：

val smallSet = spark.read.format("cosmos.oltp")
  // 配置连接信息、查询语句
  .load()
  .where("Mode != 'anc'")
  .coalesce(1)
  .cache() // 必须加缓存，避免后续复用重复触发查询

使用该方案需要注意，必须在读取链路中直接调用分区合并和缓存方法，不要拆分算子链，避免Spark优化时重新拆分任务导致重复请求。

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