如何像 sql 一样在 Spark SQL DataFrame 中求和（当时的情况）？

Question

我是 SparkSQL 的新手，我想计算我的数据中每种状态的百分比。 这是我的数据，如下所示：

A   B
11  1
11  3
12  1
13  3
12  2
13  1
11  1
12  2

所以，我可以这样在 SQL 中做到这一点：

select (C.oneTotal / C.total)   as onePercentage,
       (C.twoTotal / C.total)   as twotPercentage,
       (C.threeTotal / C.total) as threPercentage
from (select count(*) as total,
             A,
             sum(case when B = '1' then 1 else 0 end) as oneTotal,
             sum(case when B = '2' then 1 else 0 end) as twoTotal,
             sum(case when B = '3' then 1 else 0 end) as threeTotal
      from test
      group by A) as C;

但在 SparkSQL DataFrame 中，我首先计算每个状态的 totalCount，如下所示：

// wrong code
val cc = transDF.select("transData.*").groupBy("A")
      .agg(count("transData.*").alias("total"),
        sum(when(col("B") === "1", 1)).otherwise(0)).alias("oneTotal")
        sum(when(col("B") === "2", 1).otherwise(0)).alias("twoTotal")

问题是总和（何时）的结果为零。

我用错了吗？ 如何像我上面的 SQL 一样在 SparkSQL 中实现它？然后计算每个状态的占比？

感谢您的帮助。最后，我用 sum(when) 解决它。以下是我当前的代码。

val cc = transDF.select("transData.*").groupBy("A")
      .agg(count("transData.*").alias("total"),
        sum(when(col("B") === "1", 1).otherwise(0)).alias("oneTotal"),
        sum(when(col("B") === "2", 1).otherwise(0)).alias("twoTotal"))
      .select(col("total"),
        col("A"),
        col("oneTotal") / col("total").alias("oneRate"),
        col("twoTotal") / col("total").alias("twoRate"))

再次感谢。

Answer 1

您可以使用sum(when(... 或count(when..，第二个选项更短：

val df = Seq(
  (11, 1),
  (11, 3),
  (12, 1),
  (13, 3),
  (12, 2),
  (13, 1),
  (11, 1),
  (12, 2)
).toDF("A", "B")

df
  .groupBy($"A")
  .agg(
    count("*").as("total"),
    count(when($"B"==="1",$"A")).as("oneTotal"),
    count(when($"B"==="2",$"A")).as("twoTotal"),
    count(when($"B"==="3",$"A")).as("threeTotal")
  )
  .select(
    $"A",
    ($"oneTotal"/$"total").as("onePercentage"),
    ($"twoTotal"/$"total").as("twoPercentage"),
    ($"threeTotal"/$"total").as("threePercentage")
  )
  .show()

给予

+---+------------------+------------------+------------------+
|  A|     onePercentage|     twoPercentage|   threePercentage|
+---+------------------+------------------+------------------+
| 12|0.3333333333333333|0.6666666666666666|               0.0|
| 13|               0.5|               0.0|               0.5|
| 11|0.6666666666666666|               0.0|0.3333333333333333|
+---+------------------+------------------+------------------+

或者，您可以生成一个带有窗口函数的“长”表：

df
  .groupBy($"A",$"B").count()
  .withColumn("total",sum($"count").over(Window.partitionBy($"A")))
  .select(
    $"A",
    $"B",
    ($"count"/$"total").as("percentage")
  ).orderBy($"A",$"B")
  .show()

+---+---+------------------+
|  A|  B|        percentage|
+---+---+------------------+
| 11|  1|0.6666666666666666|
| 11|  3|0.3333333333333333|
| 12|  1|0.3333333333333333|
| 12|  2|0.6666666666666666|
| 13|  1|               0.5|
| 13|  3|               0.5|
+---+---+------------------+

Answer 2

据我了解，您希望实现问题中显示的上述 sql 的逻辑。

一种方法如下例所示

package examples

import org.apache.log4j.Level
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

object AggTest extends App {
  val logger = org.apache.log4j.Logger.getLogger("org")
  logger.setLevel(Level.WARN)


  val spark = SparkSession.builder.appName(getClass.getName)
    .master("local[*]").getOrCreate

  import spark.implicits._

  val df = Seq(
    (11, 1),
    (11, 3),
    (12, 1),
    (13, 3),
    (12, 2),
    (13, 1),
    (11, 1),
    (12, 2)
  ).toDF("A", "B")

  df.show(false)
  df.createOrReplaceTempView("test")
  spark.sql(
    """
      |select (C.oneTotal / C.total)   as onePercentage,
      |       (C.twoTotal / C.total)   as twotPercentage,
      |       (C.threeTotal / C.total) as threPercentage
      |from (select count(*) as total,
      |             A,
      |             sum(case when B = '1' then 1 else 0 end) as oneTotal,
      |             sum(case when B = '2' then 1 else 0 end) as twoTotal,
      |             sum(case when B = '3' then 1 else 0 end) as threeTotal
      |      from test
      |      group by A) as C
    """.stripMargin).show


}

结果：

+---+---+
|A  |B  |
+---+---+
|11 |1  |
|11 |3  |
|12 |1  |
|13 |3  |
|12 |2  |
|13 |1  |
|11 |1  |
|12 |2  |
+---+---+

+------------------+------------------+------------------+
|     onePercentage|    twotPercentage|    threPercentage|
+------------------+------------------+------------------+
|0.3333333333333333|0.6666666666666666|               0.0|
|               0.5|               0.0|               0.5|
|0.6666666666666666|               0.0|0.3333333333333333|
+------------------+------------------+------------------+