计算总和,计算多个顶部K值火花

时间:2017-08-12 21:59:49

标签: apache-spark apache-spark-sql apache-spark-dataset

我有一个格式为

的输入数据框
+---------------------------------+
|name| values |score    |row_number|
+---------------------------------+
|A    |1000   |0        |1        |
|B    |947    |0        |2        |
|C    |923    |1        |3        |
|D    |900    |2        |4        |
|E    |850    |3        |5        |
|F    |800    |1        |6        |
+---------------------------------+

我需要在得分>时获得总和(值) 0和row_number<当得分> 1时,K(i,e)所有值的和。 0表示数据帧中的前k个值。

我可以通过对前100个值运行以下查询来实现此目的

val top_100_data = df.select(
      count(when(col("score") > 0 and col("row_number")<=100, col("values"))).alias("count_100"),
      sum(when(col("score") > 0 and col("row_number")<=100, col("values"))).alias("sum_filtered_100"),
      sum(when(col("row_number") <=100, col(values))).alias("total_sum_100")
    )

但是,我需要获取前100,200,300 ...... 2500的数据。这意味着我需要运行此查询25次,最后联合25个数据帧。

我是新手,并且仍然会想出很多事情。什么是解决这个问题的最佳方法?

谢谢!

1 个答案:

答案 0 :(得分:1)

您可以创建Array限制为

val topFilters = Array(100, 200, 300) // you can add more

然后,您可以遍历topFilters数组并创建所需的dataframe我建议您使用join而不是union,因为join将为您提供单独的columns,而unions将为您提供单独的rows 。您可以执行以下操作

dataframe视为

+----+------+-----+----------+
|name|values|score|row_number|
+----+------+-----+----------+
|A   |1000  |0    |1         |
|B   |947   |0    |2         |
|C   |923   |1    |3         |
|D   |900   |2    |200       |
|E   |850   |3    |150       |
|F   |800   |1    |250       |
+----+------+-----+----------+

您可以使用上面定义的topFilters数组作为

import sqlContext.implicits._
import org.apache.spark.sql.functions._
var finalDF : DataFrame = Seq("1").toDF("rowNum")
for(k <- topFilters) {
  val top_100_data = df.select(lit("1").as("rowNum"), sum(when(col("score") > 0 && col("row_number") < k, col("values"))).alias(s"total_sum_$k"))
  finalDF = finalDF.join(top_100_data, Seq("rowNum"))
}
finalDF.show(false)

哪个应该给你最终dataframe

+------+-------------+-------------+-------------+
|rowNum|total_sum_100|total_sum_200|total_sum_300|
+------+-------------+-------------+-------------+
|1     |923          |1773         |3473         |
+------+-------------+-------------+-------------+

你可以为你的25个限制做同样的事情。

如果您打算使用union,那么这个想法与上述类似。

我希望答案很有帮助

<强>更新

如果您需要union,那么您可以使用上面定义的相同限制数组应用以下逻辑

var finalDF : DataFrame = Seq((0, 0, 0, 0)).toDF("limit", "count", "sum_filtered", "total_sum")
for(k <- topFilters) {
  val top_100_data = df.select(lit(k).as("limit"), count(when(col("score") > 0 and col("row_number")<=k, col("values"))).alias("count"),
    sum(when(col("score") > 0 and col("row_number")<=k, col("values"))).alias("sum_filtered"),
    sum(when(col("row_number") <=k, col("values"))).alias("total_sum"))
  finalDF = finalDF.union(top_100_data)
}
finalDF.filter(col("limit") =!= 0).show(false)

应该给你

+-----+-----+------------+---------+
|limit|count|sum_filtered|total_sum|
+-----+-----+------------+---------+
|100  |1    |923         |2870     |
|200  |3    |2673        |4620     |
|300  |4    |3473        |5420     |
+-----+-----+------------+---------+