Spark SQL用UDF实现按列特征重分区

Spark SQL用UDF实现按列特征重分区

浪尖 浪尖聊大数据
欢迎关注，浪尖公众号，bigdatatip，建议置顶。

这两天，球友又问了我一个比较有意思的问题:

解决问题以前，要先了解一下Spark 原理，要想进行相同数据归类到相同分区，确定要有产生shuffle步骤。

好比，F到G这个shuffle过程，那么如何决定数据到哪一个分区去的呢？这就有一个分区器的概念，默认是hash分区器。

假如，咱们能在分区这个地方着手的话确定能实现咱们的目标。

那么，在没有看Spark Dataset的接口以前，浪尖也不知道Spark Dataset有没有给我门提供这种类型的API，抱着试一试的心态，能够去Dataset类看一下，这个时候会发现有一个函数叫作repartition。

/**
  * Returns a new Dataset partitioned by the given partitioning expressions, using
  * `spark.sql.shuffle.partitions` as number of partitions.
  * The resulting Dataset is hash partitioned.
  *
  * This is the same operation as "DISTRIBUTE BY" in SQL (Hive QL).
  *
  * @group typedrel
  * @since 2.0.0
  */
 @scala.annotation.varargs
 def repartition(partitionExprs: Column*): Dataset[T] = {
   repartition(sparkSession.sessionState.conf.numShufflePartitions, partitionExprs: _*)
 }

能够传入列表达式来进行从新分区，产生的新的Dataset的分区数是由参数spark.sql.shuffle.partitions决定，那么是否是能够知足咱们的需求呢？

明显，直接用是不行的，能够间接使用UDF来实现该功能。

方式一-简单重分区

首先，实现一个UDF截取列值共同前缀，固然根据业务需求来写该udf

val substring = udf{(str: String) => {
     str.substring(0,str.length-1)
   }}

注册UDF

spark.udf.register("substring",substring)

建立Dataset

val sales = spark.createDataFrame(Seq(
     ("Warsaw1", 2016, 100),
     ("Warsaw2", 2017, 200),
     ("Warsaw3", 2016, 100),
     ("Warsaw4", 2017, 200),
     ("Beijing1", 2017, 200),
     ("Beijing2", 2017, 200),
     ("Warsaw4", 2017, 200),
     ("Boston1", 2015, 50),
     ("Boston2", 2016, 150)
   )).toDF("city", "year", "amount")

执行充分去操做

val res = sales.repartition(substring(col("city")))

打印分区ID及对应的输出结果

res.foreachPartition(partition=>{
     println("---------------------> Partition start ")
     println("partitionID is "+TaskContext.getPartitionId())
     partition.foreach(println)
     println("=====================> Partition stop ")
   })

浪尖这里spark.sql.shuffle.partitions设置的数值为10.
输出结果截图以下：

方式二-SQL实现

对于Dataset的repartition产生的shuffle是不须要进行聚合就能够产生shuffle使得按照字段值进行归类到某些分区。

SQL的实现要实现重分区要使用group by，而后udf跟上面同样，须要进行聚合操做。

完整代码以下：

val sales = spark.createDataFrame(Seq(
     ("Warsaw1", 2016, 100),
     ("Warsaw2", 2017, 200),
     ("Warsaw3", 2016, 100),
     ("Warsaw4", 2017, 200),
     ("Beijing1", 2017, 200),
     ("Beijing2", 2017, 200),
     ("Warsaw4", 2017, 200),
     ("Boston1", 2015, 50),
     ("Boston2", 2016, 150)
   )).toDF("city", "year", "amount")

   sales.registerTempTable("temp");
   val substring = udf{(str: String) => {
     str.substring(0,str.length-1)
   }}
   spark.udf.register("substring",substring)

   val res = spark.sql("select sum(amount) from temp group by substring(city)")
//
   res.foreachPartition(partition=>{
     println("---------------------> Partition start ")
     println("partitionID is "+TaskContext.getPartitionId())
     partition.foreach(println)
     println("=====================> Partition stop ")
   })

输出结果以下：

由上面的结果也能够看到task执行结束时间是无序的。

浪尖在这里主要是讲了Spark SQL 如何实现按照本身的需求对某列重分区。

那么，浪尖在这里就顺带问一下，如何用Spark Core实现该功能呢？