Flatten 'trees' with relationships expressed in two columns with Pandas or Pyspark

user_478213 0
Viewed 183

I have collection of family trees that have a format similar to the below follows:

    A
   / \
  B   C
 / \ / \
D  E F  G
  / \
 .. ..

Expressed in the following two columns (which contain multiple trees):

child parent
A
B A
C A
D B
... ...

What is the most efficient way to flatten this so that in a new column I get the upper most parent?

i.e. B = A, D = A...?

child parent highest_parent
A A
B A A
C A A
D B A
... ... ...

Ideally I would like to do this in Spark (given the size of the dataset) but could also try Pandas?

I currently cant wrap my head around doing this efficiently without some very intensive recursive function for each level (even if my trees are at most 3 levels deep).

pyspark pandas python
3 Answers

In pandas you may check with networkx

df=df.dropna()
import networkx as nx
G=nx.from_pandas_edgelist(df, 'parent', 'child',create_using=nx.DiGraph())
def find_root(G,node):
    if len(list(G.predecessors(node)))>0:
        root = find_root(G,list(G.predecessors(node))[0])
    else:
        root = node
    return root

df['child'].apply(lambda x : find_root(G,x))

Out[109]: 
1    A
2    A
3    A
user_79645270

enter image description here

import pandas as pd
import numpy as np
df = pd.DataFrame({
    "child": ['A','B','C','D','E','F','G','H','I','A1','B1','C1','D1','E1','F1','G1','H1','I1'],
    "parent": [np.NaN,'A','A','B','B','C','C','G','G',np.NaN,'A1','A1','B1','B1','C1','C1','G1','G1']
})

upper_parent_list = list(df[df['parent'].isna()]['child'])
['A', 'A1']

df['upper_parent'] = df['parent'].fillna(df['child'])

   child parent upper_parent
0      A    NaN            A
1      B      A            A
2      C      A            A
3      D      B            B
4      E      B            B
5      F      C            C
6      G      C            C
7      H      G            G
8      I      G            G
9     A1    NaN           A1
10    B1     A1           A1
11    C1     A1           A1
12    D1     B1           B1
13    E1     B1           B1
14    F1     C1           C1
15    G1     C1           C1
16    H1     G1           G1
17    I1     G1           G1

while df['upper_parent'].isin(upper_parent_list).sum()!=df.shape[0]:
    for up_par in upper_parent_list:
        child_list = list(df[df['upper_parent'].isin([up_par])]['child'])
        df['upper_parent'] = np.where(df['parent'].isin(child_list), up_par, df['upper_parent'])
print(df)

   child parent upper_parent
0      A    NaN            A
1      B      A            A
2      C      A            A
3      D      B            A
4      E      B            A
5      F      C            A
6      G      C            A
7      H      G            A
8      I      G            A
9     A1    NaN           A1
10    B1     A1           A1
11    C1     A1           A1
12    D1     B1           A1
13    E1     B1           A1
14    F1     C1           A1
15    G1     C1           A1
16    H1     G1           A1
17    I1     G1           A1
user_140665120

here is one way in pyspark, with only 3 levels. Note that the last row has 4 levels in the example and fails, not your case hopefully but to see it

import pandas as pd
import pyspark.sql.functions as F

# create toy data
pdf = pd.DataFrame({'child':list('ABCDEFGHIJKLM'), 
                    'parent':['','A','A','B','B','C','C','E', '', 'I','J','K','L']})
# convert to spark dataframe
df = spark.createDataFrame(pdf)

# coalesce the column parent
df = df.withColumn('parent', F.when(F.col('parent')!='', F.col('parent'))
                              .otherwise(F.col('child')))

# do self join using alias to direct to the right columns
res = (
    df.alias('df1')
      .join(df.alias('df2'), F.col('df1.parent') == F.col('df2.child'))
      .join(df.alias('df3'), F.col('df2.parent') == F.col('df3.child'))
      .select(['df1.child', 'df1.parent',F.col('df3.parent').alias('highest_parent')])
)

and you get

res.orderBy('child').show()
+-----+------+--------------+
|child|parent|highest_parent|
+-----+------+--------------+
|    A|     A|             A|
|    B|     A|             A|
|    C|     A|             A|
|    D|     B|             A|
|    E|     B|             A|
|    F|     C|             A|
|    G|     C|             A|
|    H|     E|             A|
|    I|     I|             I|
|    J|     I|             I|
|    K|     J|             I|
|    L|     K|             I|
|    M|     L|             J| <--this one 4 levels so fail, could add a join if needed
+-----+------+--------------+
user_92747320
Related