1. 程式人生 > >Spark之UDF、UDAF詳解

Spark之UDF、UDAF詳解

對於一個大資料處理平臺而言,倘若不能支援函式的擴充套件,確乎是不可想象的。Spark首先是一個開源框架,當我們發現一些函式具有通用的性質,自然可以考慮contribute給社群,直接加入到Spark的原始碼中。我們欣喜地看到隨著Spark版本的演化,確實湧現了越來越多對於資料分析師而言稱得上是一柄柄利器的強大函式,例如部落格文章《Spark 1.5 DataFrame API Highlights: Date/Time/String Handling, Time Intervals, and UDAFs》介紹了在1.5中為DataFrame提供了豐富的處理日期、時間和字串的函式;以及在Spark SQL 1.4中就引入的

Window Function

然而,針對特定領域進行資料分析的函式擴充套件,Spark提供了更好地置放之處,那就是所謂的“UDF(User Defined Function)”。

UDF的引入極大地豐富了Spark SQL的表現力。一方面,它讓我們享受了利用Scala(當然,也包括Java或Python)更為自然地編寫程式碼實現函式的福利,另一方面,又能精簡SQL(或者DataFrame的API),更加寫意自如地完成複雜的資料分析。尤其採用SQL語句去執行資料分析時,UDF幫助我們在SQL函式與Scala函式之間左右逢源,還可以在一定程度上化解不同資料來源具有歧異函式的尷尬。想想不同關係資料庫處理日期或時間的函式名稱吧!

用Scala編寫的UDF與普通的Scala函式沒有任何區別,唯一需要多執行的一個步驟是要讓SQLContext註冊它。例如:

def len(bookTitle: String):Int = bookTitle.length

sqlContext.udf.register("len", len _)

val booksWithLongTitle = sqlContext.sql("select title, author from books where len(title) > 10")

編寫的UDF可以放到SQL語句的fields部分,也可以作為where、groupBy或者having子句的一部分。

既然是UDF,它也得保持足夠的特殊性,否則就完全與Scala函式泯然眾人也。這一特殊性不在於函式的實現,而是思考函式的角度,需要將UDF的引數視為資料表的某個列。例如上面len函式的引數bookTitle,雖然是一個普通的字串,但當其代入到Spark SQL的語句中,實參title實際上是表中的一個列(可以是列的別名)。

當然,我們也可以在使用UDF時,傳入常量而非表的列名。讓我們稍稍修改一下剛才的函式,讓長度10作為函式的引數傳入:

def lengthLongerThan(bookTitle: String, length: Int): Boolean = bookTitle.length > length

sqlContext.udf.register("longLength", lengthLongerThan _)

val booksWithLongTitle = sqlContext.sql("select title, author from books where longLength(title, 10)")

若使用DataFrame的API,則可以以字串的形式將UDF傳入:

val booksWithLongTitle = dataFrame.filter("longLength(title, 10)")

DataFrame的API也可以接收Column物件,可以用$符號來包裹一個字串表示一個Column。$是定義在SQLContext物件implicits中的一個隱式轉換。此時,UDF的定義也不相同,不能直接定義Scala函式,而是要用定義在org.apache.spark.sql.functions中的udf方法來接收一個函式。這種方式無需register:

import org.apache.spark.sql.functions._

val longLength = udf((bookTitle: String, length: Int) => bookTitle.length > length)

import sqlContext.implicits._
val booksWithLongTitle = dataFrame.filter(longLength($"title", $"10"))

注意,程式碼片段中的sqlContext是之前已經例項化的SQLContext物件。

不幸,執行這段程式碼會丟擲異常:

cannot resolve '10' given input columns id, title, author, price, publishedDate;

因為採用$來包裹一個常量,會讓Spark錯以為這是一個Column。這時,需要定義在org.apache.spark.sql.functions中的lit函式來幫助:

val booksWithLongTitle = dataFrame.filter(longLength($"title", lit(10)))

普通的UDF卻也存在一個缺陷,就是無法在函式內部支援對錶資料的聚合運算。例如,當我要對銷量執行年度同比計算,就需要對當年和上一年的銷量分別求和,然後再利用同比公式進行計算。此時,UDF就無能為力了。

該UDAF(User Defined Aggregate Function)粉墨登場的時候了。

Spark為所有的UDAF定義了一個父類UserDefinedAggregateFunction。要繼承這個類,需要實現父類的幾個抽象方法:

def inputSchema: StructType

def bufferSchema: StructType

def dataType: DataType

def deterministic: Boolean

def initialize(buffer: MutableAggregationBuffer): Unit

def update(buffer: MutableAggregationBuffer, input: Row): Unit

def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit

def evaluate(buffer: Row): Any

可以將inputSchema理解為UDAF與DataFrame列有關的輸入樣式。例如年同比函式需要對某個可以運算的指標與時間維度進行處理,就需要在inputSchema中定義它們。

  def inputSchema: StructType = {
    StructType(StructField("metric", DoubleType) :: StructField("timeCategory", DateType) :: Nil)
  }

程式碼建立了擁有兩個StructFieldStructTypeStructField的名字並沒有特別要求,完全可以認為是兩個內部結構的列名佔位符。至於UDAF具體要操作DataFrame的哪個列,取決於呼叫者,但前提是資料型別必須符合事先的設定,如這裡的DoubleTypeDateType型別。這兩個型別被定義在org.apache.spark.sql.types中。

bufferSchema用於定義儲存聚合運算時產生的中間資料結果的Schema,例如我們需要儲存當年與上一年的銷量總和,就需要定義兩個StructField

  def bufferSchema: StructType = {
    StructType(StructField("sumOfCurrent", DoubleType) :: StructField("sumOfPrevious", DoubleType) :: Nil)
  }

dataType標明瞭UDAF函式的返回值型別,deterministic是一個布林值,用以標記針對給定的一組輸入,UDAF是否總是生成相同的結果。

顧名思義,initialize就是對聚合運算中間結果的初始化,在我們這個例子中,兩個求和的中間值都被初始化為0d:

  def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer.update(0, 0.0)
    buffer.update(1, 0.0)
  }

update函式的第一個引數為bufferSchema中兩個Field的索引,預設以0開始,所以第一行就是針對“sumOfCurrent”的求和值進行初始化。

UDAF的核心計算都發生在update函式中。在我們這個例子中,需要使用者設定計算同比的時間週期。這個時間週期值屬於外部輸入,但卻並非inputSchema的一部分,所以應該從UDAF對應類的建構函式中傳入。我為時間週期定義了一個樣例類,且對於同比函式,我們只要求輸入當年的時間週期,上一年的時間週期可以通過對年份減1來完成:

case class DateRange(startDate: Timestamp, endDate: Timestamp) {
  def in(targetDate: Date): Boolean = {
    targetDate.before(endDate) && targetDate.after(startDate)
  }
}

class YearOnYearBasis(current: DateRange) extends UserDefinedAggregateFunction {
  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    if (current.in(input.getAs[Date](1))) {
      buffer(0) = buffer.getAs[Double](0) + input.getAs[Double](0)
    }
    val previous = DateRange(subtractOneYear(current.startDate), subtractOneYear(current.endDate))
    if (previous.in(input.getAs[Date](1))) {
      buffer(1) = buffer.getAs[Double](0) + input.getAs[Double](0)
    }
  }
}  

update函式的第二個引數input: Row對應的並非DataFrame的行,而是被inputSchema投影了的行。以本例而言,每一個input就應該只有兩個Field的值。倘若我們在呼叫這個UDAF函式時,分別傳入了銷量銷售日期兩個列的話,則input(0)代表的就是銷量,input(1)代表的就是銷售日期。

merge函式負責合併兩個聚合運算的buffer,再將其儲存到MutableAggregationBuffer中:

  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getAs[Double](0) + buffer2.getAs[Double](0)
    buffer1(1) = buffer1.getAs[Double](1) + buffer2.getAs[Double](1)
  }

最後,由evaluate函式完成對聚合Buffer值的運算,得到最後的結果:

  def evaluate(buffer: Row): Any = {
    if (buffer.getDouble(1) == 0.0)
      0.0
    else
      (buffer.getDouble(0) - buffer.getDouble(1)) / buffer.getDouble(1) * 100
  }

假設我們建立了這樣一個簡單的DataFrame:

    val conf = new SparkConf().setAppName("TestUDF").setMaster("local[*]")
    val sc = new SparkContext(conf)
    val sqlContext = new SQLContext(sc)
    
    import sqlContext.implicits._

    val sales = Seq(
      (1, "Widget Co", 1000.00, 0.00, "AZ", "2014-01-01"),
      (2, "Acme Widgets", 2000.00, 500.00, "CA", "2014-02-01"),
      (3, "Widgetry", 1000.00, 200.00, "CA", "2015-01-11"),
      (4, "Widgets R Us", 2000.00, 0.0, "CA", "2015-02-19"),
      (5, "Ye Olde Widgete", 3000.00, 0.0, "MA", "2015-02-28")
    )

    val salesRows = sc.parallelize(sales, 4)
    val salesDF = salesRows.toDF("id", "name", "sales", "discount", "state", "saleDate")
    salesDF.registerTempTable("sales")

那麼,要使用之前定義的UDAF,則需要例項化該UDAF類,然後再通過udf進行註冊:

    val current = DateRange(Timestamp.valueOf("2015-01-01 00:00:00"), Timestamp.valueOf("2015-12-31 00:00:00"))
    val yearOnYear = new YearOnYearBasis(current)

    sqlContext.udf.register("yearOnYear", yearOnYear)
    val dataFrame = sqlContext.sql("select yearOnYear(sales, saleDate) as yearOnYear from sales")
    dataFrame.show()

在使用上,除了需要對UDAF進行例項化之外,與普通的UDF使用沒有任何區別。但顯然,UDAF更加地強大和靈活。如果Spark自身沒有提供符合你需求的函式,且需要進行較為複雜的聚合運算,UDAF是一個不錯的選擇。


轉自:https://www.jianshu.com/p/833b72adb2b6