filipniziol
Esteemed Contributor

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:

filipniziol_0-1724782428698.png