Azure Databricks中PySpark大DataFrame多列主键组合的高效筛选方法

在Azure Databricks中处理包含170列的大型PySpark DataFrame时，需筛选出2-3列的最优组合作为主键，满足以下要求：

• 唯一性：组合需能唯一标识每一行；
• 数据完整性：组合需避免NULL值与重复项；
• 性能：方法需高效且可扩展至大数据集；
• 业务相关性：所选列需符合业务逻辑约束。

曾尝试用itertools.combinations()遍历列组合，通过distinct().count()验证唯一性，但大数据集上多次执行该操作速度极慢，代码如下：

from itertools import combinations

selected_columns = [
    "Message_MessageBody_Header_", 
    "Message_FileSequenceNo", 
    "Message_MessageBody_ArticleInfo_BNo",
    "EnvDate",
    "EnvTime"
]  # 替换为实际列名

total_count = df.count()
print(f"Total records in DataFrame: {total_count}")

missing_columns = [col for col in selected_columns if col not in df.columns]
if missing_columns:
    print(f"Error: The following columns are missing in the DataFrame: {missing_columns}")
else:
    print(f"Selected columns exist in the DataFrame: {selected_columns}")

found_primary_key = False
for r in range(2, len(selected_columns) + 1):
    print(f"\nChecking {r}-column combinations...")

    for combo in combinations(selected_columns, r):
        print(f"\n🔍 Testing combination: {combo}")
        unique_count = df.select(*combo).distinct().count()
        print(f"Unique count for {combo}: {unique_count}")

        df.select(*combo).distinct().show(5, truncate=False)

        if unique_count == total_count:
            print(f"\n✅ Found Primary Key: {combo}")
            found_primary_key = True
            break  # 找到有效主键后停止

    if found_primary_key:
        break  

if not found_primary_key:
    print("\n❌ No unique primary key combination found in the selected columns.")

核心问题：多次执行distinct().count()导致计算成本极高，效率低下。需要找到大型PySpark DataFrame中确定最优主键组合的最高效方式，包括优化的PySpark技术、统计方法或启发式算法，以及性能优化策略或替代方案。

1. 先过滤含NULL值的列，缩小候选范围 #

主键要求无NULL值，先一次性检查所有候选列的NULL值情况，直接排除含NULL的列，减少后续组合数量：

from pyspark.sql.functions import col, count, when

# 计算每列的NULL值数量
null_counts = df.select([count(when(col(c).isNull(), c)).alias(c) for c in selected_columns])
null_counts_dict = null_counts.first().asDict()

# 过滤掉存在NULL值的列
valid_candidates = [c for c in selected_columns if null_counts_dict[c] == 0]
print(f"过滤后无NULL值的候选列: {valid_candidates}")

2. 预计算单列基数，优先测试高基数组合 #

基数（列的distinct值数量）越高的列，越容易和其他列组合出唯一键。先一次性计算所有候选列的基数，按基数降序排序后优先测试高基数列的组合，减少无效遍历：

from pyspark.sql.functions import countDistinct

# 预计算所有候选列的基数
cardinality_df = df.agg(*[countDistinct(col(c)).alias(c) for c in valid_candidates])
cardinality_dict = cardinality_df.first().asDict()

# 按基数降序排列候选列
sorted_candidates = sorted(valid_candidates, key=lambda x: -cardinality_dict[x])
print(f"按基数降序排列的候选列: {sorted_candidates}")

3. 批量验证组合，减少Spark作业提交次数 #

原代码每个组合单独提交Spark作业，开销极大。改为批量处理：将多个组合的哈希值一次性计算，同时验证唯一性，大幅减少作业提交次数：

from pyspark.sql.functions import concat_ws, hash, count as spark_count
from itertools import combinations

def batch_check_combinations(df, candidates, target_count, r):
    combinations_list = list(combinations(candidates, r))
    if not combinations_list:
        return None
    
    # 为每个组合生成唯一哈希列
    hash_cols = []
    for idx, combo in enumerate(combinations_list):
        hash_col_name = f"hash_{idx}"
        # 用concat_ws拼接列后计算哈希（可替换为md5降低碰撞概率）
        hash_expr = hash(concat_ws("||", *[col(c) for c in combo])).alias(hash_col_name)
        hash_cols.append(hash_expr)
    
    # 一次性计算所有组合的哈希值
    hash_df = df.select(*hash_cols)
    
    # 统计每个哈希列中出现次数>1的记录数（即重复数）
    duplicate_counts = hash_df.agg(*[
        spark_count(when(col(hc).isin(hash_df.groupBy(hc).count().filter(col("count")>1).select(hc)), 1)).alias(hc)
        for hc in [f"hash_{idx}" for idx in range(len(combinations_list))]
    ])
    
    duplicate_counts_dict = duplicate_counts.first().asDict()
    
    # 找到重复数为0的组合（即唯一）
    for idx, combo in enumerate(combinations_list):
        if duplicate_counts_dict[f"hash_{idx}"] == 0:
            return combo
    return None

使用示例：

total_count = df.count()
found_key = None

# 仅检查2-3列组合
for r in range(2, 4):
    print(f"\n批量检查{r}列组合...")
    found_key = batch_check_combinations(df, sorted_candidates, total_count, r)
    if found_key:
        print(f"\n✅ 找到有效主键组合: {found_key}")
        break

if not found_key:
    print("\n❌ 未找到符合要求的主键组合")

4. 抽样验证+全量确认 #

先对数据集抽样（如10%）快速验证组合的唯一性，通过抽样的组合再进行全量验证，大幅减少初期计算量：

from itertools import combinations

# 抽取10%的样本（可根据数据规模调整抽样比例）
sample_df = df.sample(withReplacement=False, fraction=0.1, seed=42)
sample_count = sample_df.count()

# 先在样本上验证组合
found_sample_key = None
for r in range(2, 4):
    for combo in combinations(sorted_candidates, r):
        sample_unique = sample_df.select(*combo).distinct().count()
        if sample_unique == sample_count:
            found_sample_key = combo
            break
    if found_sample_key:
        break

# 对样本中通过的组合进行全量验证
if found_sample_key:
    full_unique = df.select(*found_sample_key).distinct().count()
    if full_unique == total_count:
        print(f"\n✅ 全量验证通过，主键组合: {found_sample_key}")
    else:
        print(f"\n⚠️ 样本验证通过但全量存在重复，继续检查其他组合")

5. 利用Spark元数据统计信息加速 #

如果表已执行过ANALYZE TABLE收集统计信息，可以直接从元数据获取列基数，无需重复计算，适合静态数据集：

from pyspark.sql import SparkSession

spark = SparkSession.getActiveSession()
# 需提前执行：ANALYZE TABLE your_table_name COMPUTE STATISTICS FOR ALL COLUMNS
stats_df = spark.sql(f"DESCRIBE EXTENDED your_table_name")
# 解析各列的基数信息（格式因Spark版本略有差异，需提取`numDistinctValues`字段）

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