使用数据框中多个其他列的值将新列添加到Dataframe - spark / scala

时间:2017-11-22 23:20:09

标签: scala apache-spark dataframe apache-spark-sql spark-dataframe

我是新手来激发SQL和Dataframes。我有Dataframe我应该根据其他列的值添加新列。我有一个来自excel的Nested IF公式我应该实现(用于向新列添加值),当转换为程序化术语时,它是这样的:

if(k =='yes')
{
  if(!(i==''))
  {
    if(diff(max_date, target_date) < 0)
    {
      if(j == '')
      {
        "pending" //the value of the column
      }
      else {
        "approved" //the value of the column
      }
    }
    else{
      "expired" //the value of the column
    }
  }
  else{
    "" //the value should be empty
  }
}
else{
  "" //the value should be empty
} 

i,j,k are three other columns in the Dataframe.我知道我们可以使用withColumnwhen根据其他列添加新列,但我不确定如何使用该方法实现上述逻辑。

实现上述逻辑以添加新列的简单/有效方法是什么?任何帮助将不胜感激。

谢谢。

1 个答案:

答案 0 :(得分:5)

首先,让我们简化if语句:

if(k == "yes" && i.nonEmpty)
  if(maxDate - targetDate < 0)
    if (j.isEmpty) "pending" 
    else "approved"
  else "expired"
else ""

现在有两种主要方法可以实现这个目标

  1. 使用自定义UDF
  2. 使用spark内置函数:coalescewhenotherwise
  3. 自定义UDF

    现在由于条件的复杂性,编号2会相当棘手。使用自定义UDF应该符合您的需求。

    def getState(i: String, j: String, k: String, maxDate: Long, targetDate: Long): String =  
      if(k == "yes" && i.nonEmpty)
        if(maxDate - targetDate < 0)
          if (j.isEmpty) "pending" 
          else "approved"
        else "expired"
      else ""
    
    val stateUdf = udf(getState _)
    df.withColumn("state", stateUdf($"i",$"j",$"k",lit(0),lit(0)))
    

    只需将点亮(0)和点亮(0)更改为您的日期代码,这应该适合您。

    使用spark内置函数

    如果您发现性能问题,可以切换为使用coalesceotherwisewhen,如下所示:

    val isApproved = df.withColumn("state", when($"k" === "yes" && $"i" =!= "" && (lit(max_date) - lit(target_date) < 0) && $"j" =!= "", "approved").otherwise(null))
    val isPending = isApproved.withColumn("state", coalesce($"state", when($"k" === "yes" && $"i" =!= "" && (lit(max_date) - lit(target_date) < 0) && $"j" === "", "pending").otherwise(null)))
    val isExpired = isPending.withColumn("state", coalesce($"state", when($"k" === "yes" && $"i" =!= "" && (lit(max_date) - lit(target_date) >= 0), "expired").otherwise(null)))
    val finalDf = isExpired.withColumn("state", coalesce($"state", lit("")))
    

    我过去使用自定义udf来处理大型输入源而没有问题,自定义udfs可以使代码更易读,特别是在这种情况下。