Как передать дополнительные параметры в пользовательские функции в Spark SQL?

Я хочу проанализировать столбцы даты в DataFrame и для каждого столбца даты, разрешение для даты может измениться (т.е. 2011/01/10 = > 2011/01, если для разрешения установлено значение "Месяц" ).

Я написал следующий код:

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

}}

Однако это не сработает. Кажется, что я могу передать только Column в UDF. И мне интересно, будет ли это очень медленно, если я преобразую DataFrame в RDD и применил функцию в каждой строке.

Кто-нибудь знает правильное решение? Спасибо!

Ответ 1

Просто используйте немного currying:

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

и используйте его следующим образом:

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

На боковой ноте вы должны взглянуть на sql.functions.trunc и sql.functions.date_format. Они должны по крайней мере частично выполнять работу без использования UDF вообще.

Примечание

В Spark 2.2 или более поздней версии вы можете использовать функцию typedLit:

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

которые поддерживают более широкий диапазон литералов, таких как Seq или Map.

Ответ 2

Вы можете создать литерал Column для перехода к udf с помощью функции lit(...), определенной в org.apache.spark.sql.functions

Например:

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