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實現該功能呢？