Question

我有一个像这样的SparkSQL DataFrame：

name gender age isActive points
-------------------------------
 Bob      M  12     true    100
 Hal      M  16    false     80
 Pat      F  21     true     70
 Lin      F  17    false     40
 Zac      M  18     true     20
 Mei      F  19     true     10
 Sal      M  13    false     10

我有几个这样的功能：

def isEligible(prog: String) (name: String, gender: String, age: Int, isActive: Boolean, points: Int): Boolean

确定某人是否有资格参加特定计划。对于Instance，以下函数调用将告诉我Bob是否有资格使用Program1：

isEligible("Program1", "Bob", "M", 12, true, 100)

一个人可能有资格参加一个以上的课程。我想编写一个接收此DataFrame的函数，并输出如下所示的摘要DataFrame：

prog1 prog2 prog3 prog4
-----------------------
    7     3     2     5

显示符合每个计划资格的人数。 Spark中最好的方法是什么？我知道我可以使用struct和agg函数，但我不知道如何将isEligible函数合并到SparkSQL查询中。

Answer 1

定义程序列表：

val progs = Seq("prog1", "prog2", "prog3", "prog4")

定义表达式

@transient val exprs = progs.map(p => {
  val f = udf(isEligible(p) _)
  sum(f(
    $"name", $"gender", $"age", $"isActive", $"points"
  ).cast("long")).alias(p)
})

df.select(exprs: _*)

您还可以使用强类型数据集：

import org.apache.spark.sql.Row

case class Record(name: String, gender: String, age: Int, 
                  isActive: Boolean, points: Int)

df.as[Record].flatMap {
   case Record(name, gender, age, isActive, points) => 
     progs.filter(p => isEligible(p)(name, gender, age, isActive, points))
}.groupBy().pivot("value", progs).count()

SparkSQL总和如果在多个条件下

1 个答案: