scala - 在Spark Dataframe中使用函数创建新列

我正在尝试找出Spark中的新数据框API。似乎是向前迈出的一大步，但在做一些本应非常简单的事情时遇到了麻烦。我有一个包含2列的数据框，“ID”和“Amount”。作为一个通用示例，假设我想返回一个名为“code”的新列，该列将基于“Amt”的值返回一个代码。我可以写一个类似这样的函数:

def coder(myAmt:Integer):String {
  if (myAmt > 100) "Little"
  else "Big"
}

当我尝试像这样使用它时:

val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")

myDF.withColumn("Code", coder(myDF("Amt")))

我收到类型不匹配错误

found   : org.apache.spark.sql.Column
required: Integer

我尝试将函数的输入类型更改为org.apache.spark.sql.Column，但随后我开始使用函数编译得到错误，因为它需要if语句中的 bool(boolean) 值。

我做错了吗？有没有比使用withColumn更好/另一种方法？

谢谢你的帮助。

最佳答案

假设您在架构中有“Amt”列:

import org.apache.spark.sql.functions._
val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")
val coder: (Int => String) = (arg: Int) => {if (arg < 100) "little" else "big"}
val sqlfunc = udf(coder)
myDF.withColumn("Code", sqlfunc(col("Amt")))

我认为withColumn是添加列的正确方法