Azure Databricks Scala : How to replace rows following a respective hirarchy - scala

Question

Having in mind the following dataset:
I would like to obtain
As you can see, basically the idea is to follow the path indicated by column ACTUAL_ID until it is null (if it wasn't already)
I tried to use a udf where I was passing the full initial Dataframe and the recursively would find what I want but it seems it is not possible to pass Dataframes to UDFs. I also looked into replacing a value of a row, but it seems that is not possible.
My latest attempt:
def calculateLatestImdate(df: DataFrame, lookupId: String) : String = {
var foundId = df.filter($"ID" === lookupId).select($"ACTUAL_ID").first.getAs[String]("ID");
if (foundId == "" || foundId == null)
{
lookupId
}
else
{
calculateLatestImdate(df, foundId);
}
}
val calculateLatestImdateUdf = udf((df:DataFrame, s:String) => {
calculateLatestImdate(df,s)
})
val df = sc.parallelize(Seq(("1", "", "A"), ("2", "3", "B"), ("3", "6", "C"), ("4", "5", "D"), ("5", "", "E"), ("6", "", "F"))).toDF("ID","ACTUAL_ID", "DATA")
val finalDf = df.withColumn("FINAL_ID", when(isEmpty($"ACTUAL_ID"), $"ID").otherwise(calculateLatestImdateUdf(df, $"ACTUAL_ID")))

Answer 1

This looked a bit like a graph problem to me so I worked up an answer using Scala and graphframes. It makes use of the connectedComponents algorithm and the outDegrees method of the graphframe. I've made an assumption that the end of each tree is unique as per your sample data but this assumption needs to be checked. I'd be interested to see what the performance is like with more data, but let me know what you think of the solution.
The complete script:
// NB graphframes had to be installed separately with the right Scala version
import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import org.graphframes._
// Create the test data
// Vertices dataframe
val v2 = sqlContext.createDataFrame(List(
( 1, 0, "A" ), ( 2, 3, "B" ), ( 3, 6, "C" ),
( 4, 5, "D" ), ( 5, 0, "E" ), ( 6, 0, "F" )
)).toDF("id", "actual_id", "data")
// Edge dataframe
val e2 = sqlContext.createDataFrame(List(
(2, 3, "is linked to"),
(3, 6, "is linked to"),
(4, 5, "is linked to")
)).toDF("src", "dst", "relationship")
// Create the graph frame
val g2 = GraphFrame(v2, e2)
print(g2)
// The connected components adds a component id to each 'group'
sc.setCheckpointDir("/tmp/graphframes-example-connected-components")
val components = g2.connectedComponents.run() // doesn't work on Spark 1.4
display(components)
// "end" of tree nodes have no outDegree, so add that in to the component df
val endOfTree = components.join(g2.outDegrees, Seq("id"), "left")
.select("component", "data")
.where("outDegree is null")
endOfTree.show()
components.as("c").join(endOfTree.as("t"), $"c.component" === $"t.component")
.select($"c.id", $"c.component", $"t.data")
.orderBy("id")
.show()
My results:
If your data is already in a dataframe, it's easy to generate the edges dataframe from your original with just a select and a where filter, eg
// Create the GraphFrame from the dataframe
val v2 = df
val e2 = df
.select("id", "actual_id")
.withColumn("rel", lit("is linked to"))
.where("actual_id > 0")
.toDF("src", "dst", "rel")
val g2 = GraphFrame(v2, e2)
print(g2)
g2.vertices.show()
g2.edges.show()

Answer 2

Believe I have found an answer to my problem.
def calculateLatestId(df: DataFrame) : DataFrame = {
var joinedDf = df.as("df1").join(df.as("df2"), $"df1.ACTUAL_ID" === $"df2.ID", "outer").withColumn("FINAL_ID", when($"df2.ID".isNull, $"df1.ID").when($"df2.ACTUAL_ID".isNotNull, $"df2.ACTUAL_ID").otherwise($"df2.ID")).select($"df1.*", $"FINAL_ID").filter($"df1.ID".isNotNull)
val differentIds = joinedDf.filter($"df1.ACTUAL_ID" =!= $"FINAL_ID")
joinedDf = joinedDf.withColumn("ACTUAL_ID", $"FINAL_ID").drop($"FINAL_ID")
if(differentIds.count > 0)
{
calculateLatestId(joinedDf)
}
else
{
joinedDf = joinedDf.as("df1").join(joinedDf.as("df2"), $"df1.ACTUAL_ID" === $"df2.ID", "inner").select($"df1.ID", $"df2.*").drop($"df2.ID")
joinedDf
}
}
I believe the performance can be improved somehow, probably by reducing the amount of rows after each iteration and at the end do some sort of join + clean up.