Jak mogę przekazać dodatkowe parametry do UDF w SparkSql?

Question

Chcę przeanalizować kolumny z datami w DataFrame, a dla każdej kolumny daty może się zmienić rozdzielczość daty (tj. 2011/01/10 => 2011/01, jeśli rozdzielczość jest ustawiona na "Miesiąc").

Napisałem następujący kod:

def convertDataFrame(dataframe: DataFrame, schema : Array[FieldDataType], resolution: Array[DateResolutionType]) : DataFrame = 
{ 
    import org.apache.spark.sql.functions._ 
    val convertDateFunc = udf{(x:String, resolution: DateResolutionType) => SparkDateTimeConverter.convertDate(x, resolution)} 
    val convertDateTimeFunc = udf{(x:String, resolution: DateResolutionType) => SparkDateTimeConverter.convertDateTime(x, resolution)} 

    val allColNames = dataframe.columns 
    val allCols = allColNames.map(name => dataframe.col(name)) 

    val mappedCols = 
    { 
    for(i <- allCols.indices) yield 
    { 
     schema(i) match 
     { 
     case FieldDataType.Date => convertDateFunc(allCols(i), resolution(i))) 
     case FieldDataType.DateTime => convertDateTimeFunc(allCols(i), resolution(i)) 
     case _ => allCols(i) 
     } 
    } 
    } 

    dataframe.select(mappedCols:_*) 

}} 

Jednak to nie działa. Wygląda na to, że mogę przekazać tylko Column s do UDF. I zastanawiam się, czy to będzie bardzo powolne, jeśli przekonwertuję DataFrame na RDD i zastosuję funkcję w każdym rzędzie.

Czy ktoś zna prawidłowe rozwiązanie? Dziękuję Ci!

Answer 1

Wystarczy użyć trochę zmiękczania:

def convertDateFunc(resolution: DateResolutionType) = udf((x:String) => 
    SparkDateTimeConverter.convertDate(x, resolution)) 

i używać go w sposób następujący:

case FieldDataType.Date => convertDateFunc(resolution(i))(allCols(i)) 

Na marginesie należy przyjrzeć sql.functions.trunc i sql.functions.date_format. Powinny one co najmniej część pracy bez korzystania UDF w ogóle.

Uwaga:

W Spark 2.2 lub nowszej można korzystać typedLit funkcja:

import org.apache.spark.sql.functions.typedLit 

które obsługują szerszy zakres literałów jak Seq lub Map.

Answer 2

Można tworzyć dosłownym Column przejść do UDF przy użyciu funkcji lit(...) zdefiniowany w org.apache.spark.sql.functions

Na przykład:

val takeRight = udf((s: String, i: Int) => s.takeRight(i)) 
df.select(takeRight($"stringCol", lit(1)))