如何在Scala Spark中使用“ .withColumn”和许多条件为数据集创建新列

时间:2018-12-19 05:27:46

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

我有以下输入数组

val bins = (("bin1",1.0,2.0),("bin2",3.0,4.0),("bin3",5.0,6.0))

基本上,字符串“ bin1”是指引用列中的值,在该列上对数据帧进行了过滤-根据边界条件,从另一列中创建了一个新列,该剩余列在数组中保留了两个double。

var number_of_dataframes = bins.length
var ctempdf = spark.createDataFrame(sc.emptyRDD[Row],train_data.schema)
ctempdf = ctempdf.withColumn(colName,col(colName))
val t1 = System.nanoTime
for ( x<- 0 to binputs.length-1)

{
      var tempdf = train_data.filter(col(refCol) === bins(x)._1)
      //println(binputs(x)._1)
      tempdf = tempdf.withColumn(colName,
                                 when(col(colName) < bins(x)._2, bins(x)._2)
                                 when(col(colName) > bins(x)._3, bins(x)._3)
                                 otherwise(col(colName)))
      ctempdf = ctempdf.union(tempdf)
val duration = (System.nanoTime - t1) / 1e9d
println(duration)     
}

上面的代码对于每个增加的bin值都以缓慢的方式缓慢工作- 有没有一种方法可以极大地加快此速度-因为此代码再次嵌套在另一个循环中。

我已经使用了检查点/持久性/缓存,但这些都没有帮助

1 个答案:

答案 0 :(得分:2)

这里不需要迭代联合。使用o.a.s.sql.functions.map创建文字map<string, struct<double, double>>(就功能而言,其行为类似于延迟的string => struct<lower: dobule, upper: double>

import org.apache.spark.sql.functions._

val bins: Seq[(String, Double Double)] = Seq(
  ("bin1",1.0,2.0),("bin2",3.0,4.0),("bin3",5.0,6.0))

val binCol = map(bins.map { 
  case (key, lower, upper) => Seq(
    lit(key), 
    struct(lit(lower) as "lower", lit(upper) as "upper")) 
}.flatten: _*)

定义这样的表达式(这些是预定义映射中的简单查找,因此binCol(col(refCol))被延迟struct<lower: dobule, upper: double>,其余apply则采用lowerupper字段):

val lower = binCol(col(refCol))("lower")
val upper =  binCol(col(refCol))("upper")
val c = col(colName)

并使用CASE ... WHEN ...Spark Equivalent of IF Then ELSE

val result = when(c.between(lower, upper), c)
  .when(c < lower, lower)
  .when(c > upper, upper)

选择并删除NULL

df
  .withColumn(colName, result)
  // If value is still NULL it means we didn't find refCol key in binCol keys.
  // To mimic .filter(col(refCol) === ...) we drop the rows
  .na.drop(Seq(colName))

此解决方案假定NULL开头没有colName值,但可以轻松调整以处理不满足此假设的情况。

如果仍然不清楚该过程,建议您使用文字逐步进行跟踪:

spark.range(1).select(binCol as "map").show(false)
+------------------------------------------------------------+
|map                                                         |
+------------------------------------------------------------+
|[bin1 -> [1.0, 2.0], bin2 -> [3.0, 4.0], bin3 -> [5.0, 6.0]]|
+------------------------------------------------------------+
spark.range(1).select(binCol(lit("bin1")) as "value").show(false)
+----------+
|value     |
+----------+
|[1.0, 2.0]|
+----------+
spark.range(1).select(binCol(lit("bin1"))("lower") as "value").show
+-----+
|value|
+-----+
|  1.0|
+-----+

并进一步参考Querying Spark SQL DataFrame with complex types