Spark SQL用UDF實現按列特徵重分區 repatition

轉：https://cloud.tencent.com/developer/article/1371921

解決問題之前，要先了解一下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實現該功能呢？