Find "primary key" for any pyspark dataset

user_9677271 0
Viewed 912

How can I find the combination of columns in a data set (pyspark) that can be considered the primary key.

I tried to generate the combination of all the columns then compare the number of distinct records of each subset with the whole set, but it is very expensive.

from itertools import combinations

l_key = []

for i in range(len(df.columns)+1):
  print(f'Iter:{i+2}..{len(df.columns)+1}')
  for c in list(combinations(df.columns, i+2)):
    if(df.select(*c).distinct().count() == df.count()):
      l_key.append(c)
      print(f'Key:{c}')

Are there any functions or libraries that can generate this type of analysis?

pyspark apache-spark python
1 Answers

You can try it by creating the combinations as columns and then grouping.

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

pdf = pd.DataFrame([[1, 1, 1], [2, 1, 1], [3, 2, 1]], columns=['x', 'y', 'z'])

df = spark.createDataFrame(pdf)

initCols = df.columns

for i in range(len(initCols)+1):
    for c in list(combinations(initCols, i+2)):
        df = df.withColumn(','.join(c), F.concat_ws(',', *c))

finalCols = df.columns

exprs = [F.size(F.collect_set(x)).alias(x) for x in finalCols]

df = df\
.withColumn("aggCol", F.lit("a"))\
.groupBy("aggCol")\
.agg(*exprs)

df.show()

Output:

+------+---+---+---+---+---+---+-----+
|aggCol|  x|  y|  z|x,y|x,z|y,z|x,y,z|
+------+---+---+---+---+---+---+-----+
|     a|  3|  2|  1|  3|  3|  2|    3|
+------+---+---+---+---+---+---+-----+

I believe this should be less expensive. I tested it quickly on a small dataframe (~20k rows, 7 cols) and it didn't take too much time. Let me know how this works out for your dataset.

user_55390270
Related