Spark強大的函數擴展功能

在數據分析領域中，沒有人能預見所有的數據運算，以至於將它們都內置好，一切準備完好，用戶只需要考慮用，萬事大吉。擴展性是一個平台的生存之本，一個封閉的平台如何能夠擁抱變化？在對數據進行分析時，無論是演算法也好，分析邏輯也罷，最好的重用單位自然還是：函數。

故而，對於一個大數據處理平台而言，倘若不能支持函數的擴展，確乎是不可想像的。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.lengthnnsqlContext.udf.register("len", len _)nnval booksWithLongTitle = sqlContext.sql("select title, author from books where len(title) > 10")n

編寫的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 > lengthnnsqlContext.udf.register("longLength", lengthLongerThan _)nnval booksWithLongTitle = sqlContext.sql("select title, author from books where longLength(title, 10)")n

若使用DataFrame的API，則可以以字元串的形式將UDF傳入：

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

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

import org.apache.spark.sql.functions._nnval longLength = udf((bookTitle: String, length: Int) => bookTitle.length > length)nnimport sqlContext.implicits._nval booksWithLongTitle = dataFrame.filter(longLength($"title", $"10"))n

注意，代碼片段中的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)))n

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

該UDAF（User Defined Aggregate Function）粉墨登場的時候了。

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

def inputSchema: StructTypenndef bufferSchema: StructTypenndef dataType: DataTypenndef deterministic: Booleannndef initialize(buffer: MutableAggregationBuffer): Unitnndef update(buffer: MutableAggregationBuffer, input: Row): Unitnndef merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unitnndef evaluate(buffer: Row): Anyn

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

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

代碼創建了擁有兩個StructField的StructType。StructField的名字並沒有特別要求，完全可以認為是兩個內部結構的列名佔位符。至於UDAF具體要操作DataFrame的哪個列，取決於調用者，但前提是數據類型必須符合事先的設置，如這裡的DoubleType與DateType類型。這兩個類型被定義在org.apache.spark.sql.types中。

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

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

dataType標明了UDAF函數的返回值類型，deterministic是一個布爾值，用以標記針對給定的一組輸入，UDAF是否總是生成相同的結果。

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

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

update函數的第一個參數為bufferSchema中兩個Field的索引，默認以0開始，所以第一行就是針對「sumOfCurrent」的求和值進行初始化。

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

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

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

merge函數負責合併兩個聚合運算的buffer，再將其存儲到MutableAggregationBuffer中：

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

最後，由evaluate函數完成對聚合Buffer值的運算，得到最後的結果：

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

假設我們創建了這樣一個簡單的DataFrame：

val conf = new SparkConf().setAppName("TestUDF").setMaster("local[*]")n val sc = new SparkContext(conf)n val sqlContext = new SQLContext(sc)nn import sqlContext.implicits._nn val sales = Seq(n (1, "Widget Co", 1000.00, 0.00, "AZ", "2014-01-01"),n (2, "Acme Widgets", 2000.00, 500.00, "CA", "2014-02-01"),n (3, "Widgetry", 1000.00, 200.00, "CA", "2015-01-11"),n (4, "Widgets R Us", 2000.00, 0.0, "CA", "2015-02-19"),n (5, "Ye Olde Widgete", 3000.00, 0.0, "MA", "2015-02-28")n )nn val salesRows = sc.parallelize(sales, 4)n val salesDF = salesRows.toDF("id", "name", "sales", "discount", "state", "saleDate")n salesDF.registerTempTable("sales")n

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

val current = DateRange(Timestamp.valueOf("2015-01-01 00:00:00"), Timestamp.valueOf("2015-12-31 00:00:00"))n val yearOnYear = new YearOnYearBasis(current)nn sqlContext.udf.register("yearOnYear", yearOnYear)n val dataFrame = sqlContext.sql("select yearOnYear(sales, saleDate) as yearOnYear from sales")n dataFrame.show()n

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

通過Spark提供的UDF與UDAF，你可以慢慢實現屬於自己行業的函數庫，讓Spark SQL變得越來越強大，對於使用者而言，卻能變得越來越簡單。