Сумма SparkSQL, если на нескольких условиях

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

, которые определяют, имеет ли кто-то право на участие в конкретной программе.Для экземпляра следующий вызов функции скажет мне, имеет ли Боб право на участие в Программе1:

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()