Spark SQL 多列聚合的几种方法

一、引子

项目中遇到这样一张表：user

要求对这个表按照sid 进行聚合，将所有的id聚合成一个json，所有的tag聚合成一个json。
在hive和Spark中，对tag的聚合相对简单，用聚合函数collect_list 或者collect_set（hive sql 聚合函数）或者直接基于Spark算子计算。

二、思路

id和id_Type 如何聚合呢？有以下几种方案：

  1. 将id，idType 使用特殊字符拼接成字符串,看待成一列聚合;
  2. 同上，也是拼接成字符串，但使用Json保留原始数据结构;
  3. 将id,idType 通过nestRow 构造成Row 结构。
  4. 转换成rdd，基于rdd算子聚合和Json化

对于方法1：如果目标json 结构要求 按照
{ "idType":
[
{"11111111"},
{"2222222"}
]
}
这种格式存储当id中存在特殊字符时，聚合后的结果如何切割是个棘手的问题，但同时这种方式处理简单，且节省计算时内存空间。
对于方法2：对比方法1，该方法不存在切割字符串的问题，但在聚合时，由于Json串本身的格式，存在存储空间的浪费
对于方法3：目前发现Hive和Spark SQL 都还不支持嵌套的数据类型。对于arrary[Row] 这种数据结构，目前的版本（spark 2.2） 不支持。
对于方法4：会在Spark SQL和rdd 操作做对比，此次不详述

三、实现

package org.hhl.example
import org.apache.spark.sql.functions.{collect_list, collect_set}
import org.apache.spark.sql.{DataFrame, SparkSession}
import org.json4s.DefaultFormats
import org.json4s.JsonDSL._
import org.json4s.jackson.JsonMethods.{compact, parse, render}
import scala.collection.JavaConverters._
/**
  * Created by huanghl4 on 2017/11/6.
  */
object SparkSQL {
  // 获取SparkSession, spark 操作得入口
  val spark = SparkSession.builder()
    .appName(s"${this.getClass.getSimpleName}")
    .enableHiveSupport().getOrCreate()
  // 通过字符串拼接，实现多列聚合
  def multiColumnAggWithConcatStr = {
    // 拼接
    val data = spark.sql("select sid,id,idType,tag from hive.user").as[UserTag].map(x=> (x.sid,x.id + "|" + x.idType,x.tag)).toDF("sid","vid","tag")
    // 或
    //val data = spark.sql("select sid,concat(id,'|',idType),tag from hive.user").map(x=> (x.getString(0),x.getString(1),x.getString(2))
    // 聚合, 聚合函数必须导入org.apache.spark.sql.functions._
    import org.apache.spark.sql.functions._
    val dataAgg = data
      .groupBy("sid")
      .agg(
        collect_set("vid") as "ids",
        collect_list("tag") as "tags"
      ).select("sid","ids","tags").map(x =>{
      val sid = x.getString(0)
      val ids = x.getList[String](1).asScala.toList
      val tag = x.getList[String](2).asScala.toList
      (sid,strToJson(ids),listToJson(tag))
    }).toDF("sid","ids","tags")
    // 数据传输到ElasticSearch
    saveToES(dataAgg)
  }
  //通过Json实现多列聚合
  def multiColumnAggWithJson = {
    val data = spark.sql("select sid,id,idType,tag from hive.user").as[UserTag].map(x=>
      (x.sid,listToJson(List(x.id,x.idType)),x.tag))
    val dataAgg = data
      .groupBy("sid")
      .agg(
        collect_set("vid") as "ids",
        collect_list("tag") as "tags"
      ).select("sid","ids","tags").map(x =>{
      val sid = x.getString(0)
      val ids = x.getList[String](1).asScala.toList
      val tag = x.getList[String](2).asScala.toList
      (sid,strJsonToJson(ids),listToJson(tag))
    }).toDF("sid","ids","tags")
    // 数据传输到ElasticSearch
    saveToES(dataAgg)
  }

  type strList= List[String]

  def strToJson(ids:strList):String = {
    // 构造ids 的Json 结构
     val id = ids.map(x=>{
       val vid = x.split("\\|")
       (vid(0),vid(1))
     }).groupBy(_._2).map(x=>(x._1,x._2.map(_._1)))
    val json = id.map{x =>(
      x._1-> x._2
    )}
    compact(render(json))
  }
  def strJsonToJson(ids:strList):String = {
    // 构造ids 的Json 结构
    val id = ids.map(x=>{
      val vid = jsonToList(x)
      (vid(0),vid(1))
    }).groupBy(_._2).map(x=>(x._1,x._2.map(_._1)))
    val json = id.map{x =>(
      x._1-> x._2
      )}
    compact(render(json))
  }
  def listToJson(l:strList):String = compact(render(l))
  def jsonToList(str:String):strList = {
    implicit val formats = DefaultFormats
    val json = parse(str)
    json.extract[strList]
  }
  def saveToES(df:DataFrame) = {
  }
  case class UserTag(sid:String,id:String,idType:String,tag:String)
}

四、总结

1. 多列聚合可采用字符串拼接或者Json化后再聚合
2. 字符串拼接难点在于需判断数据中是否可能存在拼接字符；Json 化的难点在于聚合时由于数据量巨大，有可能带来数据倾斜问题，且处理起来较为复杂。
   可参考GitHub上实现：https://github.com/Smallhi/example/blob/master/src/main/scala/org/hhl/example/SparkSQL.scala
   如有问题联系：huanghl0817@gmail.com