Scala Spark DataFrame SQL withColumn - 如何使用函数(x:String)进行转换

时间:2018-04-03 04:30:26

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

我的目标是将列添加到现有DataFrame中,并使用DF中现有列的转换填充列。

我找到的所有示例都使用withColumn来添加列,而when()。otherwise()用于转换。

我希望使用匹配大小写的定义函数(x:String),这允许我使用字符串函数并应用更复杂的转换。

示例DataFrame

val etldf = Seq(   
            ("Total, 20 to 24 years            "),
            ("Men, 20 to 24 years              "),
            ("Women, 20 to 24 years            ")).toDF("A")

使用when()。否则()应用简单转换。我可以将一堆这些放在一起,但很快就会变得混乱。

val newcol = when($"A".contains("Men"), "Male").
  otherwise(when($"A".contains("Women"), "Female").
  otherwise("Both"))
val newdf = etldf.withColumn("NewCol", newcol)      
newdf.select("A","NewCol").show(100, false)

输出如下:

+---------------------------------+------+
|A                                |NewCol|
+---------------------------------+------+
|Total, 20 to 24 years            |Both  |
|Men, 20 to 24 years              |Male  |
|Women, 20 to 24 years            |Female|
+---------------------------------+------+

但是我想说我想要一个稍微复杂的转型:

val newcol = when($"A".contains("Total") && $"A".contains("years"), $"A".indexOf("to").toString())

它不喜欢这个,因为indexOf是一个String函数而不是ColumnName的成员。

我真正想要做的是定义一个可以实现非常复杂的转换并将其传递给withColumn()的函数:

 def AtoNewCol( A : String): String = A match {
   case a if a.contains("Men") => "Male"
   case a if a.contains("Women") => "Female"
   case a if a.contains("Total") && a.contains("years") => a.indexOf("to").toString()
   case other => "Both"
 }
 AtoNewCol("Total, 20 to 24 years            ")  

输出结果为10(“到”的位置)

但我面临同样的类型不匹配:withColumn()想要一个ColumnName对象:

scala> val newdf = etldf.withColumn("NewCol", AtoNewCol($"A"))
<console>:33: error: type mismatch;
found   : org.apache.spark.sql.ColumnName
required: String
val newdf = etldf.withColumn("NewCol", AtoNewCol($"A"))
                                                    ^

如果我更改AtoNewCol的签名(A:org.apache.spark.sql.ColumnName),我在实现中遇到同样的问题:

scala>  def AtoNewCol( A : org.apache.spark.sql.ColumnName): String = A 
match {
 |     case a if a.contains("Men") => "Male"
 |     case a if a.contains("Women") => "Female"
 |     case a if a.contains("Total") && a.contains("years") => a.indexOf("to").toString()
 |     case other => "Both"
 |   }
<console>:30: error: type mismatch;
found   : org.apache.spark.sql.Column
required: Boolean
       case a if a.contains("Men") => "Male"
                           ^
.
.
.
etc.  

我希望有一种语法允许将列的值绑定到函数。

或者可能有一个除了withColum()之外的函数,它可以为转换定义更复杂的函数。

对所有建议开放。

2 个答案:

答案 0 :(得分:3)

您需要的只是udf功能

import org.apache.spark.sql.functions._
def AtoNewCol = udf(( A : String) => A match {
  case a if a.contains("Men") => "Male"
  case a if a.contains("Women") => "Female"
  case a if a.contains("Total") && a.contains("years") => a.indexOf("to").toString()
  case other => "Both"
})

etldf.withColumn("NewCol", AtoNewCol($"A")).show(false)

你应该得到

+---------------------------------+------+
|A                                |NewCol|
+---------------------------------+------+
|Total, 20 to 24 years            |10    |
|Men, 20 to 24 years              |Male  |
|Women, 20 to 24 years            |Female|
+---------------------------------+------+

udf函数逐行工作,数据操作发生在原始数据类型上不像其他内置函数那样按列进行

答案 1 :(得分:2)

您需要为此创建UDF,您可以尝试以下方法。我按原样使用你定义的函数。

def AtoNewCol = udf((A: String) => {
  A match {
    case a if a.contains("Men") => "Male"
    case a if a.contains("Women") => "Female"
    case a if a.contains("Total") && a.contains("years") => a.indexOf("to").toString
    case other => "Both"
  }
})

etldf.withColumn("NewCol", AtoNewCol($"A")).show(false)

//    output
//    +---------------------------------+------+
//    |A                                |NewCol|
//    +---------------------------------+------+
//    |Total, 20 to 24 years            |10    | 
//    |Men, 20 to 24 years              |Male  |
//    |Women, 20 to 24 years            |Female|
//    +---------------------------------+------+