Comment récupérer les clés d'un SparseVector de pyspark ?

Question

J'ai effectué une transformation tf-idf et je veux maintenant obtenir les clés et les valeurs du résultat.

J'utilise le code udf suivant pour obtenir les valeurs :

def extract_values_from_vector(vector):
    return vector.values.tolist()

extract_values_from_vector_udf = udf(lambda vector:extract_values_from_vector(vector), ArrayType(DoubleType()))

extract = rescaledData.withColumn("extracted_keys", extract_keys_from_vector_udf("features"))

Ainsi, si le vecteur épars ressemble à : features=SparseVector(123241, {20672 : 4.4233, 37393 : 0.0, 109847 : 3.7096, 118474 : 5.4042}))

Les clés extraites dans mon extrait ressembleront à : [4.4233, 0.0, 3.7096, 5.4042]

Ma question est la suivante : comment puis-je obtenir les clés du dictionnaire SparseVector ? Par exemple, keys = [20672, 37393, 109847, 118474] ?

J'essaie le code suivant mais il ne fonctionne pas

def extract_keys_from_vector(vector):
    return vector.indices.tolist()
extract_keys_from_vector_udf = spf.udf(lambda vector:extract_keys_from_vector(vector), ArrayType(DoubleType()))

Le résultat qu'il m'a donné est le suivant : [null,null,null,null]

Quelqu'un peut-il nous aider ? Merci d'avance !

Answer 1

Puisque la réponse se trouve dans les commentaires ci-dessus, je me suis dit que j'allais prendre le temps (en attendant d'écrire un parquet bien sûr) d'écrire la réponse.

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

def extract_keys_from_vector(vector):
    return vector.indices.tolist()

feature_extract = F.UserDefinedFunction(lambda vector: extract_keys_from_vector(vector), ArrayType(IntegerType()))

df = df.withColumn("features", feature_extract(F.col("features")))