Options
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
08-27-2024 11:14 AM
Hi @Databricks143 ,
This code works in Scala. One thing. It has a number of iterations hard-coded to 10. If there are more levels, you need to adjust
import org.apache.spark.sql.functions._
val data = Seq(
(1, "Alice", None: Option[Int]), // Root employee
(2, "Bob", Some(1)), // Reports to Alice
(3, "Charlie", Some(2)), // Reports to Bob
(4, "David", Some(1)), // Reports to Alice
(5, "Eve", Some(3)) // Reports to Charlie
)
val df = data.toDF("userId", "userName", "managerId")
// Initialize root employee, set EmpLevel to 0
var df_levels = df.filter($"managerId".isNull).withColumn("EmpLevel", lit(0))
// Create a DataFrame to hold the results as they are built up
var final_df = df_levels
// Iteratively process subordinates
val max_iterations = 10 // Adjust based on expected maximum depth
var hasNewLevels = true
for (_ <- 1 to max_iterations if hasNewLevels) {
// Identify subordinates of employees whose levels have already been calculated
val df_new_levels = df.as("emp")
.join(df_levels.as("mgr"), $"emp.managerId" === $"mgr.userId", "inner")
.select(
$"emp.userId",
$"emp.userName",
$"emp.managerId",
($"mgr.EmpLevel" + 1).as("EmpLevel")
)
// If no new levels are calculated, set the flag to false to break the loop
if (df_new_levels.isEmpty) {
hasNewLevels = false
} else {
// Append the new levels to the final result set
final_df = final_df.union(df_new_levels)
// Update df_levels to include only the newly calculated levels for the next iteration
df_levels = df_new_levels
}
}
// Final result
final_df.orderBy("EmpLevel", "userId").show()The output: