Berechnung gewichteten Mittelwert in PySpark

Question

Ich versuche, gewichtetes Mittel in pyspark zu berechnen, aber nicht viele Fortschritte

# Example data 
df = sc.parallelize([ 
    ("a", 7, 1), ("a", 5, 2), ("a", 4, 3), 
    ("b", 2, 2), ("b", 5, 4), ("c", 1, -1) 
]).toDF(["k", "v1", "v2"]) 
df.show() 

import numpy as np 
def weighted_mean(workclass, final_weight): 
    return np.average(workclass, weights=final_weight) 

weighted_mean_udaf = pyspark.sql.functions.udf(weighted_mean, 
    pyspark.sql.types.IntegerType()) 

machen, aber wenn ich versuche, diesen Code

df.groupby('k').agg(weighted_mean_udaf(df.v1,df.v2)).show() 

Ich erhalte die auszuführen Fehler

Meine Frage ist, kann ich eine benutzerdefinierte Funktion (mit mehreren Argumenten) als Argument angeben zu agg? Wenn nicht, gibt es eine Alternative, um Operationen wie das gewichtete Mittel nach der Gruppierung mit einem Schlüssel durchzuführen?

Answer 1

Benutzerdefinierte Aggregationsfunktion (UDAF, die auf pyspark.sql.GroupedData funktioniert, aber nicht in PYSPARK unterstützt) ist keine benutzerdefinierte Funktion (UDF, die auf pyspark.sql.DataFrame funktioniert).

Da in pyspark können Sie nicht Ihre eigene UDAF erstellen, und die gelieferten UDAFs kann das Problem nicht lösen, können Sie zu RDD Welt gehen müssen zurück:

from numpy import sum 

def weighted_mean(vals): 
    vals = list(vals) # save the values from the iterator 
    sum_of_weights = sum(tup[1] for tup in vals) 
    return sum(1. * tup[0] * tup[1]/sum_of_weights for tup in vals) 

df.map(
    lambda x: (x[0], tuple(x[1:])) # reshape to (key, val) so grouping could work 
).groupByKey().mapValues(
    weighted_mean 
).collect()