web-dev-qa-db-fra.com

Comment puis-je transmettre des paramètres supplémentaires aux FDU dans Spark SQL?

Je veux analyser les colonnes de date dans un DataFrame, et pour chaque colonne de date, la résolution de la date peut changer (ie 2011/01/10 => 2011/01 si la résolution est réglée sur "Mois" ).

J'ai écrit le code suivant:

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

}}

Mais ça ne marche pas. Il semble que je ne peux transmettre que des Column aux UDF. Et je me demande si ce sera très lent si je convertis le DataFrame en RDD et applique la fonction sur chaque ligne.

Quelqu'un connaît-il la bonne solution? Je vous remercie!

scalaapache-sparkapache-spark-sqluser-defined-functions
16
DarkZero

Utilisez juste un peu de curry:

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

et utilisez-le comme suit:

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

En passant, vous devriez jeter un œil à sql.functions.trunc et sql.functions.date_format. Ceux-ci devraient au moins faire partie du travail sans utiliser du tout d'UDF.

Remarque:

Dans Spark 2.2 ou version ultérieure, vous pouvez utiliser la fonction typedLit:

import org.Apache.spark.sql.functions.typedLit

qui prennent en charge une plus large gamme de littéraux comme Seq ou Map.

38
zero323

Vous pouvez créer un littéral Column à passer à un udf à l'aide de la fonction lit(...) définie dans org.Apache.spark.sql.functions

Par exemple:

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