Пользовательская функция SparkSQL в предложении when

Question

У меня есть DataFrame, подобный этому:

id val1 val2
------------
 1  v11  v12
 2  v21  v22
 3  v31  v32
 4  v41  v42
 5  v51  v52
 6  v61  v62

Каждая строка представляет человека, который может принадлежать к одной или нескольким группам. У меня есть функция, которая принимает значения для каждой строки и определяет, соответствует ли этот человек критериям для определенной группы:

def isInGroup: Boolean = f(group: Int)(id: String, v1: String, v2: String)

и я пытаюсь вывести DataFrame следующим образом:

Group1 Group2 Group3 Group4
---------------------------
     3      0      6      1

Вот мой код, который пока не работает. К сожалению, предложение when принимает только параметр типа Column, и моя функция не работает. Пользовательские функции тоже не работают. Я действительно хотел бы придерживаться select / struct /, как это было, если это возможно.

val summaryDF = dataDF
    .select(struct(
        sum(when(isInGroup(1)($"id", $"val1", $"val2"), value = 1)).as("Group1")),
        sum(when(isInGroup(2)($"id", $"val1", $"val2"), value = 1)).as("Group2")),
        sum(when(isInGroup(3)($"id", $"val1", $"val2"), value = 1)).as("Group3")),
        sum(when(isInGroup(4)($"id", $"val1", $"val2"), value = 1)).as("Group4"))
    ))

Answer 1

Как я показал в моем предыдущем ответе , вам понадобится udf:

import org.apache.spark.sql.functions.udf 

def isInGroupUDF(group: Int) = udf(isInGroup(group) _)

sum(when(
  isInGroupUDF(1)($"id", $"val1", $"val2"), 1
)).as("Group1")

Если вы хотите избежать перечисления столбцов, вы можете попробовать, например, с аргументами по умолчанию:

def isInGroupUDF(group: Int, id: Column = $"id", 
                 v1: Column = $"val1", v2: Column = $"val2") = {
  val f = udf(isInGroup(group) _)
  f(id, v1, v2)
}

sum(when(
  isInGroupUDF(1), 1
)).as("Group1")