udf.py

def personality_udf(personality_mapping: dict):
    @udf(returnType=IntegerType())
    def inner(cat):
        return personality_mapping.get(cat, -1)

    return inner

def compute_gvector_udf(personality_mapping: dict):
    @udf(returnType=ArrayType(DoubleType()))
    def inner(ngvector, cgvector, suma):
        result = [0.] * len(personality_mapping)
        for i in range(len(ngvector)):
            result[ngvector[i]] = cgvector[i] / suma
        return result

    return inner