scala - Apache 星火。 UDF 列基于另一列，而不将其名称作为参数传递。

Question

DataSet 中有 firm 列，我正在向该 DataSet 添加另一列 - firm_id 示例:

private val firms: mutable.Map[String, Integer] = ...
private val firmIdFromCode: (String => Integer) = (code: String) => firms(code)
val firm_id_by_code: UserDefinedFunction = udf(firmIdFromCode)
...
val ds = dataset.withColumn("firm_id", firm_id_by_code($"firm"))

有没有一种方法可以消除传递 $"firm" 作为参数(此列始终存在于 DS 中)。我正在为此寻找一些东西:

val ds = dataset.withColumn("firm_id", firm_id_by_code)

Answer 1

您可以在定义 udf 时提供它将使用的列。

val someUdf = udf{ /*udf code*/}.apply($"colName")

// Usage in dataset
val ds = dataset.withColumn("newColName",someUdf)