如何在Databricks高效识别150张Delta表的主键列

问题背景 #

我需要用Python/PySpark在Databricks的Delta表中识别主键列，部分表数据量达500万行，甚至存在由10列组成的复合主键。

最初采用的方法是统计全表行数与列组合的去重行数，判断是否相等来识别主键。但这种方法在列数较多时组合数爆炸（比如12列中5列为主键，需要计算C(12,1)+C(12,2)+…+C(12,5)种组合），在一张530万行、12列的表上耗时11-13分钟。尝试过采样（10%/1%）但耗时无明显改善，使用approx_count_distinct因偏差无法满足精确性要求。

原实现代码如下：

from itertools import combinations
from pyspark.sql import SparkSession

# Catalog and schema
catalog_name = "cat"
schema_name = "sch"

# Optional: restrict to specific tables; leave empty to process all
restricted_tables =  ["sales"]# e.g., ["customers", "orders"]

# Progressive sample fractions
sample_fractions = [0.2]

# Max number of columns to consider for exact PK detection
max_columns = 6

# Get all tables in the schema (DataFrame, no collect)
tables_df = spark.sql(f"SHOW TABLES IN {catalog_name}.{schema_name}")

# Apply restriction if any tables are listed
if restricted_tables:
    tables_df = tables_df.filter(tables_df.tableName.isin(restricted_tables))

metadata_results = []

for row in tables_df.toLocalIterator():
    catalog = catalog_name
    schema = row['database']
    table_name = row['tableName']
    full_table_name = f"{catalog}.{schema}.{table_name}"
    
    df = spark.table(full_table_name)
    columns = df.columns
    num_columns = len(columns)
    found_keys = False
    pk_columns = None
    pk_num_columns = None
    status = "success"
    
    # Start checking from 1 column to total columns
    for num_cols in range(1, len(columns)+1):
        print(f"Checking {num_cols} column combinations...")
        if num_cols > max_columns:
            status = "Exceeded 6 columns"
            break
        
        for col_combo in combinations(columns, num_cols):
            is_unique = True
            
            for fraction in sample_fractions:
                print(f"Checking {num_cols} {fraction} sample...")
                sample_df = df.sample(fraction=fraction, seed=42)
                total_count = sample_df.count()
                distinct_count = sample_df.select(*col_combo).distinct().count()
                
                if total_count != distinct_count:
                    is_unique = False
                    break
            
            if is_unique:
                pk_columns = list(col_combo)
                pk_num_columns = len(col_combo)
                found_keys = True
                break
        
        if found_keys or status == "Exceeded 6 columns":
            break
    
    metadata_results.append({
        "catalog": catalog,
        "schema": schema,
        "table_name": table_name,
        "pk_columns": pk_columns,
        "num_columns": num_columns,
        "pk_num_columns": pk_num_columns,
        "status": status
    })

# Convert to DataFrame
metadata_df = spark.createDataFrame(metadata_results)

核心优化思路 #

1. 利用Delta表内置统计信息快速筛选 #

Delta表会自动维护表级和列级统计数据（需开启自动统计收集），可以先通过DESCRIBE DETAIL获取总行数，快速排除不可能成为主键的列：

• 单列基数小于总行数的，直接排除单列主键可能；
• 后续组合优先从基数高的列中选取，减少无效检查。

2. 剪枝组合空间 #

• 排除非唯一单列：先验证所有单列的唯一性，后续复合组合仅从可能组成唯一组合的列中选取；
• 提前终止逻辑：找到最小的主键组合后立即停止，无需检查更多组合；
• 采样预筛选：先用小样本快速排除明显不唯一的组合，仅对样本通过的组合做全表验证。

3. 复用计算资源 #

避免每个组合单独触发Spark作业，复用采样数据进行初步检查，减少数据扫描次数。

改进后的代码 #

from itertools import combinations
from pyspark.sql import functions as F

# Catalog and schema
catalog_name = "cat"
schema_name = "sch"

# Optional: restrict to specific tables; leave empty to process all
restricted_tables = ["sales"]

# Sample fraction for pre-screening (fast elimination)
pre_sample_fraction = 0.1

# Max number of columns to consider for PK
max_columns = 10

# Get all tables in the schema
tables_df = spark.sql(f"SHOW TABLES IN {catalog_name}.{schema_name}")
if restricted_tables:
    tables_df = tables_df.filter(tables_df.tableName.isin(restricted_tables))

metadata_results = []

for row in tables_df.toLocalIterator():
    catalog = catalog_name
    schema = row['database']
    table_name = row['tableName']
    full_table_name = f"{catalog}.{schema}.{table_name}"
    
    # Step 1: Get Delta table total row count from metadata
    detail_df = spark.sql(f"DESCRIBE DETAIL {full_table_name}")
    total_rows = detail_df.select("numRows").collect()[0][0]
    if total_rows == 0:
        metadata_results.append({
            "catalog": catalog,
            "schema": schema,
            "table_name": table_name,
            "pk_columns": None,
            "num_columns": 0,
            "pk_num_columns": 0,
            "status": "Empty table"
        })
        continue
    
    df = spark.table(full_table_name)
    columns = df.columns
    num_columns = len(columns)
    found_keys = False
    pk_columns = None
    pk_num_columns = None
    status = "success"
    
    # Step 2: Pre-screen columns with sample to eliminate obvious non-unique candidates
    sample_df = df.sample(fraction=pre_sample_fraction, seed=42)
    sample_row_count = sample_df.count()
    
    # Collect columns that are unique in sample (potential candidates)
    candidate_columns = []
    for col in columns:
        distinct_count = sample_df.select(col).distinct().count()
        if distinct_count == sample_row_count:
            candidate_columns.append(col)
    
    # Fallback to all columns if no candidates found (sample might have false negatives)
    if not candidate_columns:
        candidate_columns = columns
    
    # Step 3: Check combinations starting from smallest column count
    for num_cols in range(1, min(max_columns, num_columns) + 1):
        print(f"Checking {num_cols}-column combinations for table {table_name}...")
        
        # Generate combinations from candidate columns
        for col_combo in combinations(candidate_columns, num_cols):
            # Fast check with sample first
            sample_distinct = sample_df.select(*col_combo).distinct().count()
            if sample_distinct != sample_row_count:
                continue
            
            # Verify with full table only if sample passes
            full_distinct = df.select(*col_combo).distinct().count()
            if full_distinct == total_rows:
                pk_columns = list(col_combo)
                pk_num_columns = num_cols
                found_keys = True
                break
        
        if found_keys:
            break
    
    if not found_keys:
        status = "No valid PK found within max column limit"
    
    metadata_results.append({
        "catalog": catalog,
        "schema": schema,
        "table_name": table_name,
        "pk_columns": pk_columns,
        "num_columns": num_columns,
        "pk_num_columns": pk_num_columns,
        "status": status
    })

# Convert results to DataFrame for visualization
metadata_df = spark.createDataFrame(metadata_results)
display(metadata_df)

额外优化建议 #

• 开启Delta自动统计：执行以下命令开启自动统计收集，减少手动计算开销：

  SET spark.databricks.delta.properties.defaults.autoOptimize.optimizeWrite = true;
  SET spark.databricks.delta.properties.defaults.autoOptimize.autoCompact = true;
  

• 列排序优化：将列按基数从高到低排序，优先检查高基数列的组合，更快定位主键；
• 并行处理多表：将表列表转为RDD并行处理，避免单循环遍历的串行瓶颈。

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