如何高效查找PySpark数据集的主键列组合？

当前遍历所有列组合逐一校验的实现存在天然效率瓶颈：列组合的枚举复杂度会随列数增长呈指数级上升，且每一轮校验都触发全量distinct().count()的分布式计算，数据量稍大就会出现任务排队、资源溢出的问题。

可直接复用的现成能力 #

目前不需要从零手写遍历逻辑，已有成熟组件可直接实现列组合唯一性校验、主键自动识别：

• 开源数据质量校验生态中，已内置列组合唯一性校验的封装算子，这类组件普遍做了计算逻辑优化，无需手动管理列组合的遍历流程。
• Spark 3.3及以上版本内置的聚合观测能力，可在单次全表扫描中，同时计算多组列组合的重复值比例，不需要每检测一组候选键就重复拉取全表数据。
• 轻量PySpark专用工具包中，已实现自动化的主键检测函数，传入目标DataFrame后可直接返回满足唯一性要求的最小列组合，无需手动遍历所有可能的列组合。

低开销实现优化思路 #

如果不想引入额外依赖，可在原有逻辑基础上调整实现方式，将资源开销降低90%以上：

• 前置筛除无效候选列：先逐列计算单列的基数（即该列的去重值数量），直接把基数远低于全表总记录数、无法对唯一识别产生正向增益的列从优先组合列表中剔除。比如某列只有固定5个枚举值，其作为主键组成部分的贡献度极低，可以放到搜索链路的最后顺位。
• 用贪心搜索代替全量组合遍历：不需要从2列组合开始枚举所有可能的子集，先选当前基数最高的列作为初始候选键，每轮往候选键中加入一个能最大幅度降低重复值的列，直到候选键的去重记录数等于全表总记录数时停止搜索。这种方式最多只需要执行和列数相等的聚合计算，相比原逻辑的指数级组合遍历，效率提升非常明显。
• 合并重复计算任务：将多组待检测列组合的统计逻辑合并到同一次全表扫描中，一次性计算所有候选组合的重复数，避免每检测一组列就重复触发全表数据读取与计算。

上述优化逻辑的最终检测结果，和原全量枚举遍历的校验结果完全一致，只是砍掉了无意义的重复计算与无效组合枚举。

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