Как передать линейную комбинацию столбцов в функции агрегации в Pyspark

Question

Как перевести эту часть sql в эквивалент Pyspark syntx:

Select sum(a*(1-b)) from MTABLE group by (a,b)

Я попытался создать функцию, которая принимает два аргумента и возвращает x * (1-y), но не может передать это как аргумент суммы

  df.groupby('a','b'

 .agg(sum(q1_sum1_udf('a','b'))).show()



def q1_sum1(x,y): 
     return int(x)*(int(y)-1)

Я знаю, что это возможно в scala, но не могу найти то же самое в pyspark

Answer 1

Вы можете сделать это без UDF, так как расчет прост.

df = spark.sql("select * from yourtable")
result = df.groupBy(df.a,df.b) \ 
           .sum(df.a*(1-df.b))
result.show()

Редактировать: Использование UDF.Определите функцию (лямбда-функция показана здесь, но способ определения любой другой функции такой же) и используйте функцию с ранее определенным именем и передайте необходимые аргументы.

from pyspark.sql import functions as func
#Define and register a function
f_sum = func.udf(lambda a,b:a*(1-b))
result = df.groupBy(df.a,df.b).select(f_sum(df.a,df.b))
result.show()