Combining Rows that link together in a Spark Dataframe

Question

I have a dataframe who's rows connect to each other through various merges. So far I have got the DF into the format below where I did a groupBy "Merge_To" and collected these into an array then joining this back to my original DF. Which looks like this:

df1
+---+--------+---------+
|Ref|Merge_To|   Merges|
+---+--------+---------+
|  1|      \N|[3, 2, 3]|
|  2|       1|[5, 4, 6]|
|  5|       2|   [8, 7]|
| 10|      \N|   [9, 9]|
| 12|      \N|     [13]|
| 14|      \N|     [15]|
| 16|      18|     [17]|
| 17|      16|     [19]|
| 18|      \N|     [16]|
| 19|      17|     [20]|
+---+--------+---------+

For Refs 1,2,5 and 18,16,17,19,20 these merge together through a chain. Which isn't captured through the groupBy I have done previously. Eventually I want my DF to look like this, which accounts for the merge chains:

+---+--------+------------------------+
|Ref|Merge_To|                  Merges|
+---+--------+------------------------+
|  1|      \N|[3, 2, 3, 5, 4, 6, 8, 7]|
| 10|      \N|                  [9, 9]|
| 12|      \N|                    [13]|
| 14|      \N|                    [15]|
| 18|      \N|        [16, 17, 19, 20]|
+---+--------+------------------------+

I have tried joining df1 to itself filtered when "Merge_To" is not \N

val arrayCombineUDF = udf((a:Seq[String], b:Seq[String]) => a ++ b )

val df1Filter = df1.filter($"Merge_To" !== "\N").
select("Merge_To", "Merges").withColumnRenamed("Merge_To", "Chain_Ref").
withColumnRenamed("Merges", "Chain_Merges")

val df2 = df1.join(df1Filter, $"Ref" === $"Chain_Ref", "left").
withColumn("Merges", when($"Chain_Merges".isNotNull, arrayCombineUDF($"Merges", $"Chain_Merges")).
otherwise($"Merges")).
select("Ref", "Merge_To", "Merges")

df2
+---+--------+----------+------------------+
|Ref|Merge_To|Merge_From|            Merges|
+---+--------+----------+------------------+
|  1|      \N|         3|[3, 2, 3, 5, 4, 6]|
|  2|       1|        \N|   [5, 4, 6, 8, 7]|
|  5|       2|        \N|            [8, 7]|
| 10|      \N|         9|            [9, 9]|
| 12|      \N|        13|              [13]|
| 14|      \N|        \N|              [15]|
| 16|      18|        \N|          [17, 19]|
| 17|      16|        \N|          [19, 20]|
| 18|      \N|        \N|          [16, 17]|
| 19|      17|        \N|              [20]|
+---+--------+----------+------------------+

This kind of gives the result that I am looking for but only really accounts for one level of the merge chain.

I have also tried putting that same join process as above in a while loop trying to get it to repeatedly join. I have also tried to use UDFs with If statements hoping that I could sort out each row into the type of merge it was and combine the ones that are a chain using this.

NB: I know that the arrays aren't distinct but I don't mind and can sort this at the end.

EDIT Here's the original DF

+---+--------+----------+
|Ref|Merge_To|Merge_From|
+---+--------+----------+
|  1|      \N|         3|
|  2|       1|        \N|
|  3|       1|        \N|
|  4|       2|        \N|
|  5|       2|        \N|
|  6|       2|        \N|
|  7|       5|        \N|
|  8|       5|        \N|
|  9|      10|        \N|
| 10|      \N|         9|
| 11|      \N|        \N|
| 12|      \N|        13|
| 13|      \N|        \N|
| 14|      \N|        \N|
| 15|      14|        \N|
| 16|      18|        \N|
| 17|      16|        \N|
| 18|      \N|        \N|
| 19|      17|        \N|
| 20|      19|        \N|
+---+--------+----------+

user_6933993 · Answer

The 11 entry seems to have gone AWOL. Anyway.

Looking at your base data, this is a hierarchical query that can be solved in traditional RDBMSes with nice functionality e.g. CONNECT BY Clause or Recursive WITH View in most RDBMSes.

This attempt you have made stops at level 1 and this is the crux of the problem. In addition this problem is not well-solvable with Spark's partitioning approach for parallelization. What to partition on? Any partition could have data related to the sets you are looking for.

You would be better advised doing processing there and sqooping into your Hive tables, or reading via Spark using JDBC.
You can simulate as per this less than well documented approach https://sqlandhadoop.com/how-to-implement-recursive-queries-in-spark/ here, which is what I sometimes use for BI behind the scene processing.
If you must do in the Spark realm, then how about using graphFrames approach, but runs reasonably slow, as follows with a subset of data and altering your approach slightly - see what you think:

import org.apache.spark.rdd.RDD

 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.graphframes._  
 sc.setCheckpointDir("/checkpoints")

// Subset of your data
val rdd = sc.parallelize( Array(("A", 1, None), ("B", 2, Some(1)), ("C", 3, Some(1)), ("D", 4, Some(2)), ("E", 5, Some(2)), ("F", 6, Some(2)),
                                ("G", 7, Some(5)), ("H", 8, Some(5)), ("X", 9, Some(10)), ("Y", 10, None), ("X2", 12, Some(13)), ("Y3", 13, None)  
                               ))
val df = rdd.toDF("v", "c", "p")

val dfV = df.select($"c".as("id"))
val dfE = df
  .withColumnRenamed("c", "src")
  .withColumnRenamed("p", "dst")

val nGraph = GraphFrame(dfV, dfE)
dfE.cache()
dfV.cache()
val res = nGraph.connectedComponents.run()
val res2 = res.join(df, res("id") === df("c"), "inner")
val res3 = res2.filter("p is not null").groupBy("component").agg(collect_list("id") as "group")
val res4 = res3.join(res2, res3("component") === res2("component") && res2("p").isNull , "inner")
res4.select($"id", $"group").show(false)

returns:

+---+---------------------+
|id |group                |
+---+---------------------+
|1  |[2, 3, 4, 5, 6, 7, 8]|
|10 |[9]                  |
|13 |[12]                 |
+---+---------------------+

My advice: does in RDBMSs. Using graphFrames takes a long time with this approach.

Combining Rows that link together in a Spark Dataframe

1 Answers