Come posso passare parametri extra alle UDF in SparkSql?

Question

Voglio analizzare le colonne della data in un DataFrame e, per ciascuna colonna della data, la risoluzione per la data potrebbe cambiare (ad esempio 2011/01/10 => 2011/01 se la risoluzione è impostata su "Mese").

ho scritto il seguente codice:

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:_*) 

}} 

Tuttavia non funziona. Sembra che io possa solo passare Column s alle UDF. E mi chiedo se sarà molto lento se converto il DataFrame in RDD e applichiamo la funzione su ogni riga.

Qualcuno conosce la soluzione corretta? Grazie!

Answer 1

basta usare un po 'di accattivarsi:

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

e usarlo come segue:

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

Su un lato nota si dovrebbe dare un'occhiata a sql.functions.trunc e sql.functions.date_format. Questi dovrebbero almeno parte del lavoro senza utilizzare le UDF.

Nota:

In Spark 2.2 o versioni successive è possibile utilizzare la funzione typedLit:

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

che supportano una gamma più ampia di letterali come Seq o Map.

Answer 2

È possibile creare un letterale Column per passare ad un'UDF utilizzando la funzione lit(...) definito org.apache.spark.sql.functions

Ad esempio:

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