package scala

 import org.apache.log4j.{Level, Logger}

 import org.apache.spark.sql.types.{StringType, StructField, StructType}

 import org.apache.spark.sql.{Row, SparkSession}

 /**

   * 计算分组topN

   * Created by Administrator on 2019/11/20.

   */

 object GroupTopN {

   Logger.getLogger("org").setLevel(Level.WARN) // 设置日志级别

   def main(args: Array[String]) {

     //创建测试数据

     val test_data = Array("CJ20191120,201911", "CJ20191120,201910", "CJ20191105,201910", "CJ20191105,201909", "CJ20191111,201910")

     val spark = SparkSession.builder().appName("GroupTopN").master("local[2]").getOrCreate()

     val sc = spark.sparkContext

     val test_data_rdd = sc.parallelize(test_data).map(row => {

       val Array(scene, cycle) = row.split(",")

       Row(scene, cycle)

     })

     // 设置数据模式

     val structType = StructType(Array(

       StructField("scene", StringType, true),

       StructField("cycle", StringType, true)

     ))

     // 转换为df

     val test_data_df = spark.createDataFrame(test_data_rdd, structType)

     test_data_df.createOrReplaceTempView("test_data_df")

     // 拼接周期

     val scene_ws = spark.sql("select scene,concat_ws(',',collect_set(cycle)) as cycles from test_data_df group by scene")

     scene_ws.count()

     scene_ws.show()

     scene_ws.createOrReplaceTempView("scene_ws")

     /**

       * 定义参数确定N的大小，暂定为1

       */

     val sum = 1

     // 创建广播变量，把N的大小广播出去

     val broadcast = sc.broadcast(sum)

     /**

       * 定义Udf实现获取组内的前N个数据

       */

     spark.udf.register("getTopN", (cycles : String) => {

       val sum = broadcast.value

       var mid = ""

       if(cycles.contains(",")){ // 多值

         val cycle = cycles.split(",").sorted.reverse // 降序排序

         val min = Math.min(cycle.length, sum)

         for(i <- 0 until min){

           if(mid.equals("")){

             mid = cycle(i)

           }else{

             mid += "," + cycle(i)

           }

         }

       }else{ // 单值

         mid = cycles

       }

       mid

     })

     val result = spark.sql("select scene,getTopN(cycles) cycles from scene_ws")

     result.show()

     spark.stop()

   }

 }