scala - 如何在Spark SQL中将额外的参数传递给UDF？

Question

我想解析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并将函数应用于每一行是否会非常慢。

有谁知道正确的解决方案？谢谢!

Answer 1

只需使用一些小技巧:

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。