[Python] Best Practices for Selecting Primary Key Combinations from Multiple Columns

I am working in Azure Databricks with a large PySpark DataFrame that has 170 columns. I need to identify the best possible combination of 2-3 columns to use as the primary key, ensuring:

Uniqueness: The selected combination should uniquely identify each row. Data Integrity: The combination should avoid NULLs and duplicates. Performance: The approach should be efficient and scalable on large datasets. Business Relevance: The chosen columns should align with business logic constraints. What I have tried: I attempted an approach using itertools.combinations() to check uniqueness for different column combinations. However, this method is slow on large datasets due to multiple .distinct().count() operations.

Here’s my code:

from itertools import combinations

selected_columns = [
"Message_MessageBody_Header_",
"Message_FileSequenceNo",
"Message_MessageBody_ArticleInfo_BNo",
"EnvDate",
"EnvTime"
] # Replace with actual column names

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 # Stop once a valid key is found

if found_primary_key:
break

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


Issue: This approach is computationally expensive due to repeated .distinct().count() operations on large data.

Question: What is the most efficient way to determine the optimal primary key combination in a large PySpark DataFrame? Are there any optimized PySpark techniques, statistical methods, or heuristics that can make this process faster?

Any performance optimization strategies or alternative approaches would be greatly appreciated!

Continue reading...