在 agg 中使用 pyspark groupBy 和自定義 function

Question

我想通過我的 spark df 與自定義 agg function 進行分組：

def gini(list_of_values):
    sth is processing here
    return number output

我想得到這樣的東西：

df.groupby('activity')['mean_event_duration_in_hours].agg(gini)

你能幫我解決這個問題嗎？

Answer 1

您可以像這樣創建一個udf ：

import pyspark.sql.functions as F
from pyspark.sql.types import FloatType

def gini(list_of_values):
    # sth is processing here
    return number_output

udf_gini = F.udf(gini, FloatType())

df.groupby('activity')\
    .agg(F.collect_list("mean_event_duration_in_hours").alias("event_duration_list"))\
    .withColumn("gini", udf_gini(F.col("event_duration_list")))

或者像這樣將 gini 定義為 UDF：

@udf(returnType=FloatType())
def gini(list_of_values):
    # sth is processing here
    return number_output

df.groupby('activity')\
    .agg(F.collect_list("mean_event_duration_in_hours").alias("event_duration_list"))\
    .withColumn("gini", gini(F.col("event_duration_list")))