[Spark]根据父子节点构建树

绝地灬酷狼 2022-11-27 08:46 156阅读 0赞

> 数据集中有两列child和parent，需要依据此父子节点来构建树，找到根节点以及调用链。

##  一、递归遍历 ##

import org.apache.spark.sql.DataFrame;
    val sqlContext=new org.apache.spark.sql.SQLContext(sc);
    
    val df = sqlContext.createDataFrame(List(
      ("a", "b", "b", "a-b", "20190201"),
      ("b", "c", "c", "b-c", "20190201"),
      ("c", "d", "d", "c-d", "20190201"),
      ("d", "e", "e", "d-e", "20190201")
    )).toDF("child", "parent", "grandparent", "chain", "dt")
    
    var result = df.selectExpr("child as child1","parent as parent1","grandparent as grandparent1","chain as chain1","dt as dt1")
    
    def findroot(dataframe:DataFrame) : DataFrame = {
        val r = result.where("grandparent1 is not null").take(1)
        if(r.isEmpty) {
            return result
        }
        result = dataframe
        result = df.join(result,df("parent") === result("child1"),"left_outer").selectExpr("child as child1","NVL(parent1, parent) as parent1","grandparent1","concat_ws('-',child,chain1) as chain1","dt as dt1")
        findroot(result)
    }
    
    val output = findroot(result)
    
    output.show

运行结果如下：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM3NzcxNDc1_size_16_color_FFFFFF_t_70][]

## 二、优化 ##

import org.apache.spark.sql.functions._
      import org.apache.spark.sql.DataFrame
      def findRoot(df:DataFrame):DataFrame = {
        import df.sparkSession.implicits._
        df.cache()
        var result = df.withColumn("chain",concat_ws("->",df("child"),df("parent"))).withColumn("root",df("parent")).withColumn("tmp",df("parent"))
        var sample = result.where("tmp is not null").take(1)
        while (!sample.isEmpty) {
          result = result.as("result").join(df.as("source"),($"source.child"===$"result.tmp").and($"result.tmp".isNotNull),"left_outer")
            .select(col("result.child"),col("result.parent")
              ,concat_ws ("->",col("result.chain"),col("source.parent")).as("chain")
              ,coalesce(col("source.parent"),column("result.root")).as("root"),col("source.parent").as("tmp")  )
          result.cache()
          sample = result.where("tmp is not null").take(1)
        }
    
        df.unpersist()
        result.drop("tmp")
      } 
    
    val sqlContext=new org.apache.spark.sql.SQLContext(sc);
    
    val df = sqlContext.createDataFrame(List(
      ("a", "b"),
      ("b", "c"),
      ("c", "d"),
      ("d", "e")
    )).toDF("child", "parent")
    
    val result = findRoot(df)
    result.show

运行结果如下：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM3NzcxNDc1_size_16_color_FFFFFF_t_70 1][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM3NzcxNDc1_size_16_color_FFFFFF_t_70]: /images/20221124/7bb3a041e7374b1288f29e7688837003.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM3NzcxNDc1_size_16_color_FFFFFF_t_70 1]: /images/20221124/02340a3d87464497a8bcec4cb9c010ad.png