Spark SQL用UDF實現按列特徵重分割槽
歡迎關注，浪尖公眾號，bigdatatip，建議置頂。
這兩天，球友又問了我一個比較有意思的問題:
解決問題之前，要先了解一下Spark 原理，要想進行相同資料歸類到相同分割槽，肯定要有產生shuffle步驟。
比如，F到G這個shuffle過程，那麼如何決定資料到哪個分割槽去的呢？這就有一個分割槽器的概念，預設是hash分割槽器。
假如，我們能在分割槽這個地方著手的話肯定能實現我們的目標。
那麼，在沒有看Spark Dataset的介面之前，浪尖也不知道Spark Dataset有沒有給我門提供這種型別的API，抱著試一試的心態，可以去Dataset類看一下，這個時候會發現有一個
/**
* 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
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實現該功能呢？
推薦閱讀：
Kafka原始碼系列之0.10版本的Producer原始碼解析及效能點講解
歡迎點贊，轉發，星球入口點選原文閱讀。