CosmosDB查询极速完成，同款PySpark SQL查询耗时超久求排查优化

核心原因 #

你遇到的问题并非PySpark固有特性，而是自定义查询（customQuery）与Spark Cosmos DB连接器的分区策略不匹配导致的。

为何Cosmos DB端快、Spark端慢？ #

Cosmos DB执行该聚合查询时，直接在服务端完成计算，返回的是分组后的少量结果。但Spark Cosmos DB连接器默认采用全扫描分区策略：

• 连接器会先获取容器所有物理分区，再对每个分区发起查询请求
• 即便你的customQuery已经在Cosmos DB端完成聚合，连接器仍会默认并行扫描所有分区，且在Spark端重复处理数据（实际返回结果量极小）
• df.rdd.isEmpty()属于Spark Action操作，会强制触发整个数据读取流程，大量无效的分区扫描和数据传输直接导致耗时超1小时

为何添加Restrictive分区策略后耗时骤降？ #

配置spark.cosmos.read.partitioning.strategy=Restrictive时，连接器会：

• 先执行查询获取结果的分区键分布，或直接利用Cosmos DB聚合结果的元数据
• 仅针对包含有效数据的分区发起请求，避免空分区/无关分区的扫描
• 跳过不必要的并行扫描逻辑，直接获取Cosmos DB端已聚合完成的结果，因此耗时大幅缩短

优化建议 #

1. 保留自定义查询，配合Restrictive分区策略
   你的自定义SQL已在Cosmos DB端完成过滤和聚合，这是最高效的方式（Cosmos DB擅长处理自身数据的聚合计算）。只需在读取配置中明确指定分区策略：

   df = spark.read.format("cosmos.oltp")\
       .option("spark.synapse.linkedService", "<你的链接服务名>")\
       .option("spark.cosmos.container", "<你的容器名>")\
       .option("spark.cosmos.read.customQuery", """
           SELECT c.Name, count(c.Enabled) as Redeemed
           FROM c
           WHERE NOT IS_NULL(c.Enabled)
           AND c.Name NOT IN ('EXAMPLE1', 'EXAMPLE2')
           GROUP BY c.Name
       """)\
       .option("spark.cosmos.read.partitioning.strategy", "Restrictive")\
       .load()
   

   这种方式既利用了Cosmos DB的服务端计算能力，又通过Restrictive策略避免了Spark端的无效扫描。

2. 不建议改用DataFrame过滤（除非特殊场景）
   若放弃customQuery，改用Spark DataFrame的filter、groupBy操作，需将5000万条数据全部拉取到Spark集群后再计算，会带来巨大的数据传输开销和资源消耗，效率远低于在Cosmos DB端完成聚合。

3. 替换df.rdd.isEmpty()的高效方式
   rdd.isEmpty()会触发全量数据读取，可改用df.head(1)或df.count()判断数据集是否为空，配合Restrictive策略时，能更快触发Cosmos DB查询并返回结果：

   # 更高效的非空判断
   if df.head(1):
       # 业务处理逻辑
   

总结 #

你的自定义查询本身是高效的，问题根源在Spark Cosmos DB连接器的默认分区策略。通过指定Restrictive分区策略，可让连接器正确利用Cosmos DB的服务端查询结果，避免无效分区扫描，大幅降低耗时。无需放弃自定义查询改用DataFrame过滤。

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