我想在Spark DataFrame中创建一个跨列的sumproduct。我有一个DataFrame,看起来像这样:
id val1 val2 val3 val4
123 10 5 7 5
我还有一张地图,看起来像:
val coefficents = Map("val1" -> 1, "val2" -> 2, "val3" -> 3, "val4" -> 4)
我想在DataFrame的每一列中取一个值,将其乘以映射中的相应值,然后在一个新列中返回结果,本质上是:
(10*1) + (5*2) + (7*3) + (5*4) = 61
我试过这个:
val myDF1 = myDF.withColumn("mySum", {var a:Double = 0.0; for ((k,v) <- coefficients) a + (col(k).cast(DoubleType)*coefficients(k));a})
但是得到一个错误,即"+"方法已重载。即使我解决了这个问题,我也不确定这是否可行。有什么想法吗?我总是可以动态地将SQL查询构建为文本字符串,并以这种方式进行,但我希望有更雄辩的东西。
欢迎提出任何想法。
代码的问题是您试图将Column
添加到Double
。cast(DoubleType)
只影响存储值的类型,而不影响列本身的类型。由于Double
没有提供*(x: org.apache.spark.sql.Column): org.apache.spark.sql.Column
方法,所以一切都失败了。
为了让它发挥作用,你可以做这样的事情:
import org.apache.spark.sql.Column
import org.apache.spark.sql.functions.{col, lit}
val df = sc.parallelize(Seq(
(123, 10, 5, 7, 5), (456, 1, 1, 1, 1)
)).toDF("k", "val1", "val2", "val3", "val4")
val coefficients = Map("val1" -> 1, "val2" -> 2, "val3" -> 3, "val4" -> 4)
val dotProduct: Column = coefficients
// To be explicit you can replace
// col(k) * v with col(k) * lit(v)
// but it is not required here
// since we use * f Column.* method not Int.*
.map{ case (k, v) => col(k) * v } // * -> Column.*
.reduce(_ + _) // + -> Column.+
df.withColumn("mySum", dotProduct).show
// +---+----+----+----+----+-----+
// | k|val1|val2|val3|val4|mySum|
// +---+----+----+----+----+-----+
// |123| 10| 5| 7| 5| 61|
// |456| 1| 1| 1| 1| 10|
// +---+----+----+----+----+-----+
问题似乎是您实际上没有使用a
for((k, v) <- coefficients) a + ...
你可能是指a += ...
此外,一些清理withColumn
调用内部代码块的建议:
您不需要调用coefficients(k)
,因为您已经从for((k,v) <- coefficients)
获得了v
中的值
Scala非常擅长制作一行代码,但如果你必须在这一行中放入分号,那就有点作弊了:p我建议把求和部分分解为每个表达式一行。
求和表达式可以重写为fold
,从而避免使用var
(惯用的Scala通常避免var
s),例如
import org.apache.spark.sql.functions.lit
coefficients.foldLeft(lit(0.0)){
case (sumSoFar, (k,v)) => col(k).cast(DoubleType) * v + sumSoFar
}
我不确定这是否可以通过DataFrame API实现,因为您只能使用列,而不能使用任何预定义的闭包(例如参数映射)。
我在下面概述了一种使用DataFrame的底层RDD:的方法
import org.apache.spark.sql.types._
import org.apache.spark.sql.Row
// Initializing your input example.
val df1 = sc.parallelize(Seq((123, 10, 5, 7, 5))).toDF("id", "val1", "val2", "val3", "val4")
// Return column names as an array
val names = df1.columns
// Grab underlying RDD and zip elements with column names
val rdd1 = df1.rdd.map(row => (0 until row.length).map(row.getInt(_)).zip(names))
// Tack on accumulated total to the existing row
val rdd2 = rdd0.map { seq => Row.fromSeq(seq.map(_._1) :+ seq.map { case (value: Int, name: String) => value * coefficents.getOrElse(name, 0) }.sum) }
// Create output schema (with total)
val totalSchema = StructType(df1.schema.fields :+ StructField("total", IntegerType))
// Apply schema to create output dataframe
val df2 = sqlContext.createDataFrame(rdd1, totalSchema)
// Show output:
df2.show()
...
+---+----+----+----+----+-----+
| id|val1|val2|val3|val4|total|
+---+----+----+----+----+-----+
|123| 10| 5| 7| 5| 61|
+---+----+----+----+----+-----+