scala - 如何调用在 Scala 数据框中返回 int 的函数并将其附加

Question

我有一个像这样的数据框

x  y
_  _
1  10
2  30
3  50
4  24
5  36 
6  45

我想附加另一列 z，这将取决于 y 的值。

所以我创建了一个函数

def  GiveNumVal(col: Column) => Integer = {

      if(Column>=0 && Column<15){
        return 1;
      }
      else if(Column>=15 && Column<30){
        return 2;
      }
      else if(Column>=30 && Column<45){
        return 3;
      }

     else if (Column>=45 && Column<=59){
        return 4;
      }
      else{
        return 0;
      }
    }

我打电话

val new_df=df.withColumn("z",GiveNumVal($"y"));

它甚至无法编译。我不确定哪里是错误的部分。感谢任何帮助。

Answer 1

您需要注册要使用的udf或创建一个像这样的udf

import org.apache.spark.sql.functions._
// create dataframe
val df = Seq(
  (1, 10),
  (2, 30),
  (3, 50),
  (4, 24),
  (5, 36),
  (6, 45)
).toDF("x", "y")

//create udf 
def giveNumVal = udf((c : Int) => {
  if(c >=0 && c <15) 1
  else if(c >=15 && c <30) 2
  else if(c >=30 && c <45) 3
  else if (c >=45 && c <=59) 4
  else 0
})

并将其用作

val new_df=df.withColumn("z",giveNumVal($"y"));

如果你有通用功能并且想用udf那么你可以注册为

//general function 
  def giveNumVal = (c : Int) => {
    //implementation here 
  }

//To register 
  val GiveNumVal = spark.sqlContext.udf.register("functionName", giveNumVal)

输出:

+---+---+---+
|x  |y  |z  |
+---+---+---+
|1  |10 |1  |
|2  |30 |3  |
|3  |50 |4  |
|4  |24 |2  |
|5  |36 |3  |
|6  |45 |4  |
+---+---+---+

注意:在 scala 中不需要 return 语句和 ;

希望这有帮助!