Question

我修补了org.apache.spark.sql.Column类以添加chainUDF方法。它适用于不接受参数的udfs，我需要帮助使其成为带参数的udfs的通用。

这是当前的chainUDF方法定义。

object ColumnExt {

  implicit class ColumnMethods(c: Column) {

    def chainUDF(udfName: String): Column = {
      callUDF(udfName, c)
    }

  }

}

以下是chainUDF方法。

def appendZ(s: String): String = {
  s"${s}Z"
}

spark.udf.register("appendZUdf", appendZ _)

def prependA(s: String): String = {
  s"A${s}"
}

spark.udf.register("prependAUdf", prependA _)

val hobbiesDf = Seq(
  ("dance"),
  ("sing")
).toDF("word")

val actualDf = hobbiesDf.withColumn(
  "fun",
  col("word").chainUDF("appendZUdf").chainUDF("prependAUdf")
)

我想更新chainUDF方法定义，因此需要Column个参数的可选列表。像这样：

def appendWord(s: String, word: String): String = {
  s"${s}${word}"
}

spark.udf.register("appendWordUdf", appendWord _)

val hobbiesDf = Seq(
  ("dance"),
  ("sing")
).toDF("word")

val actualDf = hobbiesDf.withColumn(
  "fun",
  col("word").chainUDF("appendZUdf").chainUDF("appendWordUdf", lit("cool"))
)

我认为我们需要将chainUDF方法定义更新为以下内容：

object ColumnExt {

  implicit class ColumnMethods(c: Column) {

    def chainUDF(udfName: String, cols: Column* = some_default_value): Column = {
      callUDF(udfName, c + cols)
    }

  }

}

我确信有一些Scala魔术技巧可以实现这一目标。

Answer 1

签名是：

def callUDF(udfName: String, cols: Column*): Column

所以你不需要魔术：

def chainUDF(udfName: String, cols: Column* = some_default_value): Column = {
  callUDF(udfName, c +: cols: _*)
}

使用callUDF创建链接UDF调用的方法

1 个答案: